feat: implement token budget management in agent and orchestrator (#71)

* feat: implement token budget management in agent and orchestrator

* fix: resolve TypeScript type errors in event and trace handlers

* feat: add budget exceeded event handling in agent and orchestrator

---------

Co-authored-by: MrAvalonApple <74775400+ibrahimkazimov@users.noreply.github.com>

src/agent/agent.ts

@@ -150,6 +150,7 @@ export class Agent {
       agentName: this.name,
       agentRole: this.config.systemPrompt?.slice(0, 50) ?? 'assistant',
       loopDetection: this.config.loopDetection,
+      maxTokenBudget: this.config.maxTokenBudget,
     }
 
     this.runner = new AgentRunner(
@@ -328,6 +329,16 @@ export class Agent {
       const result = await runner.run(messages, runOptions)
       this.state.tokenUsage = addUsage(this.state.tokenUsage, result.tokenUsage)
 
+      if (result.budgetExceeded) {
+        let budgetResult = this.toAgentRunResult(result, false)
+        if (this.config.afterRun) {
+          budgetResult = await this.config.afterRun(budgetResult)
+        }
+        this.transitionTo('completed')
+        this.emitAgentTrace(callerOptions, agentStartMs, budgetResult)
+        return budgetResult
+      }
+
       // --- Structured output validation ---
       if (this.config.outputSchema) {
         let validated = await this.validateStructuredOutput(
@@ -461,6 +472,7 @@ export class Agent {
         tokenUsage: mergedTokenUsage,
         toolCalls: mergedToolCalls,
         structured: validated,
+        ...(retryResult.budgetExceeded ? { budgetExceeded: true } : {}),
       }
     } catch {
       // Retry also failed
@@ -472,6 +484,7 @@ export class Agent {
         tokenUsage: mergedTokenUsage,
         toolCalls: mergedToolCalls,
         structured: undefined,
+        ...(retryResult.budgetExceeded ? { budgetExceeded: true } : {}),
       }
     }
   }
@@ -502,7 +515,7 @@ export class Agent {
           const result = event.data as import('./runner.js').RunResult
           this.state.tokenUsage = addUsage(this.state.tokenUsage, result.tokenUsage)
 
-          let agentResult = this.toAgentRunResult(result, true)
+          let agentResult = this.toAgentRunResult(result, !result.budgetExceeded)
           if (this.config.afterRun) {
             agentResult = await this.config.afterRun(agentResult)
           }
@@ -598,6 +611,7 @@ export class Agent {
       toolCalls: result.toolCalls,
       structured,
       ...(result.loopDetected ? { loopDetected: true } : {}),
+      ...(result.budgetExceeded ? { budgetExceeded: true } : {}),
     }
   }
 

src/agent/runner.ts

@@ -29,6 +29,7 @@ import type {
   LoopDetectionConfig,
   LoopDetectionInfo,
 } from '../types.js'
+import { TokenBudgetExceededError } from '../errors.js'
 import { LoopDetector } from './loop-detector.js'
 import { emitTrace } from '../utils/trace.js'
 import type { ToolRegistry } from '../tool/framework.js'
@@ -70,6 +71,8 @@ export interface RunnerOptions {
   readonly agentRole?: string
   /** Loop detection configuration. When set, detects stuck agent loops. */
   readonly loopDetection?: LoopDetectionConfig
+  /** Maximum cumulative tokens (input + output) allowed for this run. */
+  readonly maxTokenBudget?: number
 }
 
 /**
@@ -117,6 +120,8 @@ export interface RunResult {
   readonly turns: number
   /** True when the run was terminated or warned due to loop detection. */
   readonly loopDetected?: boolean
+  /** True when the run was terminated due to token budget limits. */
+  readonly budgetExceeded?: boolean
 }
 
 // ---------------------------------------------------------------------------
@@ -217,6 +222,7 @@ export class AgentRunner {
    *  - `{ type: 'text', data: string }` for each text delta
    *  - `{ type: 'tool_use', data: ToolUseBlock }` when the model requests a tool
    *  - `{ type: 'tool_result', data: ToolResultBlock }` after each execution
+ *  - `{ type: 'budget_exceeded', data: TokenBudgetExceededError }` on budget trip
    *  - `{ type: 'done', data: RunResult }` at the very end
    *  - `{ type: 'error', data: Error }` on unrecoverable failure
    */
@@ -232,6 +238,7 @@ export class AgentRunner {
     const allToolCalls: ToolCallRecord[] = []
     let finalOutput = ''
     let turns = 0
+    let budgetExceeded = false
 
     // Build the stable LLM options once; model / tokens / temp don't change.
     // toToolDefs() returns LLMToolDef[] (inputSchema, camelCase) — matches
@@ -318,6 +325,21 @@ export class AgentRunner {
           yield { type: 'text', data: turnText } satisfies StreamEvent
         }
 
+        const totalTokens = totalUsage.input_tokens + totalUsage.output_tokens
+        if (this.options.maxTokenBudget !== undefined && totalTokens > this.options.maxTokenBudget) {
+          budgetExceeded = true
+          finalOutput = turnText
+          yield {
+            type: 'budget_exceeded',
+            data: new TokenBudgetExceededError(
+              this.options.agentName ?? 'unknown',
+              totalTokens,
+              this.options.maxTokenBudget,
+            ),
+          } satisfies StreamEvent
+          break
+        }
+
         // Extract tool-use blocks for detection and execution.
         const toolUseBlocks = extractToolUseBlocks(response.content)
 
@@ -516,6 +538,7 @@ export class AgentRunner {
       tokenUsage: totalUsage,
       turns,
       ...(loopDetected ? { loopDetected: true } : {}),
+      ...(budgetExceeded ? { budgetExceeded: true } : {}),
     }
 
     yield { type: 'done', data: runResult } satisfies StreamEvent

src/errors.ts

@@ -0,0 +1,19 @@
+/**
+ * @fileoverview Framework-specific error classes.
+ */
+
+/**
+ * Raised when an agent or orchestrator run exceeds its configured token budget.
+ */
+export class TokenBudgetExceededError extends Error {
+  readonly code = 'TOKEN_BUDGET_EXCEEDED'
+
+  constructor(
+    readonly agent: string,
+    readonly tokensUsed: number,
+    readonly budget: number,
+  ) {
+    super(`Agent "${agent}" exceeded token budget: ${tokensUsed} tokens used (budget: ${budget})`)
+    this.name = 'TokenBudgetExceededError'
+  }
+}

src/index.ts

@@ -107,6 +107,7 @@ export {
 
 export { createAdapter } from './llm/adapter.js'
 export type { SupportedProvider } from './llm/adapter.js'
+export { TokenBudgetExceededError } from './errors.js'
 
 // ---------------------------------------------------------------------------
 // Memory

src/orchestrator/orchestrator.ts

@@ -63,6 +63,7 @@ import { Team } from '../team/team.js'
 import { TaskQueue } from '../task/queue.js'
 import { createTask } from '../task/task.js'
 import { Scheduler } from './scheduler.js'
+import { TokenBudgetExceededError } from '../errors.js'
 
 // ---------------------------------------------------------------------------
 // Internal constants
@@ -83,6 +84,12 @@ function addUsage(a: TokenUsage, b: TokenUsage): TokenUsage {
   }
 }
 
+function resolveTokenBudget(primary?: number, fallback?: number): number | undefined {
+  if (primary === undefined) return fallback
+  if (fallback === undefined) return primary
+  return Math.min(primary, fallback)
+}
+
 /**
  * Build a minimal {@link Agent} with its own fresh registry/executor.
  * Registers all built-in tools so coordinator/worker agents can use them.
@@ -266,6 +273,10 @@ interface RunContext {
   readonly runId?: string
   /** AbortSignal for run-level cancellation. Checked between task dispatch rounds. */
   readonly abortSignal?: AbortSignal
+  cumulativeUsage: TokenUsage
+  readonly maxTokenBudget?: number
+  budgetExceededTriggered: boolean
+  budgetExceededReason?: string
 }
 
 /**
@@ -409,6 +420,23 @@ async function executeQueue(
       }
 
       ctx.agentResults.set(`${assignee}:${task.id}`, result)
+      ctx.cumulativeUsage = addUsage(ctx.cumulativeUsage, result.tokenUsage)
+      const totalTokens = ctx.cumulativeUsage.input_tokens + ctx.cumulativeUsage.output_tokens
+      if (
+        !ctx.budgetExceededTriggered
+        && ctx.maxTokenBudget !== undefined
+        && totalTokens > ctx.maxTokenBudget
+      ) {
+        ctx.budgetExceededTriggered = true
+        const err = new TokenBudgetExceededError('orchestrator', totalTokens, ctx.maxTokenBudget)
+        ctx.budgetExceededReason = err.message
+        config.onProgress?.({
+          type: 'budget_exceeded',
+          agent: assignee,
+          task: task.id,
+          data: err,
+        } satisfies OrchestratorEvent)
+      }
 
       if (result.success) {
         // Persist result into shared memory so other agents can read it
@@ -446,6 +474,10 @@ async function executeQueue(
 
     // Wait for the entire parallel batch before checking for newly-unblocked tasks.
     await Promise.all(dispatchPromises)
+    if (ctx.budgetExceededTriggered) {
+      queue.skipRemaining(ctx.budgetExceededReason ?? 'Skipped: token budget exceeded.')
+      break
+    }
 
     // --- Approval gate ---
     // After the batch completes, check if the caller wants to approve
@@ -524,8 +556,8 @@ async function buildTaskPrompt(task: Task, team: Team): Promise<string> {
  */
 export class OpenMultiAgent {
   private readonly config: Required<
-    Omit<OrchestratorConfig, 'onApproval' | 'onProgress' | 'onTrace' | 'defaultBaseURL' | 'defaultApiKey'>
+    Omit<OrchestratorConfig, 'onApproval' | 'onProgress' | 'onTrace' | 'defaultBaseURL' | 'defaultApiKey' | 'maxTokenBudget'>
-  > & Pick<OrchestratorConfig, 'onApproval' | 'onProgress' | 'onTrace' | 'defaultBaseURL' | 'defaultApiKey'>
+  > & Pick<OrchestratorConfig, 'onApproval' | 'onProgress' | 'onTrace' | 'defaultBaseURL' | 'defaultApiKey' | 'maxTokenBudget'>
 
   private readonly teams: Map<string, Team> = new Map()
   private completedTaskCount = 0
@@ -545,6 +577,7 @@ export class OpenMultiAgent {
       defaultProvider: config.defaultProvider ?? 'anthropic',
       defaultBaseURL: config.defaultBaseURL,
       defaultApiKey: config.defaultApiKey,
+      maxTokenBudget: config.maxTokenBudget,
       onApproval: config.onApproval,
       onProgress: config.onProgress,
       onTrace: config.onTrace,
@@ -592,11 +625,13 @@ export class OpenMultiAgent {
    * @param prompt - The user prompt to send.
    */
   async runAgent(config: AgentConfig, prompt: string): Promise<AgentRunResult> {
+    const effectiveBudget = resolveTokenBudget(config.maxTokenBudget, this.config.maxTokenBudget)
     const effective: AgentConfig = {
       ...config,
       provider: config.provider ?? this.config.defaultProvider,
       baseURL: config.baseURL ?? this.config.defaultBaseURL,
       apiKey: config.apiKey ?? this.config.defaultApiKey,
+      maxTokenBudget: effectiveBudget,
     }
     const agent = buildAgent(effective)
     this.config.onProgress?.({
@@ -611,6 +646,18 @@ export class OpenMultiAgent {
 
     const result = await agent.run(prompt, traceOptions)
 
+    if (result.budgetExceeded) {
+      this.config.onProgress?.({
+        type: 'budget_exceeded',
+        agent: config.name,
+        data: new TokenBudgetExceededError(
+          config.name,
+          result.tokenUsage.input_tokens + result.tokenUsage.output_tokens,
+          effectiveBudget ?? 0,
+        ),
+      })
+    }
+
     this.config.onProgress?.({
       type: 'agent_complete',
       agent: config.name,
@@ -681,6 +728,24 @@ export class OpenMultiAgent {
     const decompositionResult = await coordinatorAgent.run(decompositionPrompt, decompTraceOptions)
     const agentResults = new Map<string, AgentRunResult>()
     agentResults.set('coordinator:decompose', decompositionResult)
+    const maxTokenBudget = this.config.maxTokenBudget
+    let cumulativeUsage = addUsage(ZERO_USAGE, decompositionResult.tokenUsage)
+
+    if (
+      maxTokenBudget !== undefined
+      && cumulativeUsage.input_tokens + cumulativeUsage.output_tokens > maxTokenBudget
+    ) {
+      this.config.onProgress?.({
+        type: 'budget_exceeded',
+        agent: 'coordinator',
+        data: new TokenBudgetExceededError(
+          'coordinator',
+          cumulativeUsage.input_tokens + cumulativeUsage.output_tokens,
+          maxTokenBudget,
+        ),
+      })
+      return this.buildTeamRunResult(agentResults)
+    }
 
     // ------------------------------------------------------------------
     // Step 2: Parse tasks from coordinator output
@@ -724,19 +789,45 @@ export class OpenMultiAgent {
       config: this.config,
       runId,
       abortSignal: options?.abortSignal,
+      cumulativeUsage,
+      maxTokenBudget,
+      budgetExceededTriggered: false,
+      budgetExceededReason: undefined,
     }
 
     await executeQueue(queue, ctx)
+    cumulativeUsage = ctx.cumulativeUsage
 
     // ------------------------------------------------------------------
     // Step 5: Coordinator synthesises final result
     // ------------------------------------------------------------------
+    if (
+      maxTokenBudget !== undefined
+      && cumulativeUsage.input_tokens + cumulativeUsage.output_tokens > maxTokenBudget
+    ) {
+      return this.buildTeamRunResult(agentResults)
+    }
     const synthesisPrompt = await this.buildSynthesisPrompt(goal, queue.list(), team)
     const synthTraceOptions: Partial<RunOptions> | undefined = this.config.onTrace
       ? { onTrace: this.config.onTrace, runId: runId ?? '', traceAgent: 'coordinator' }
       : undefined
     const synthesisResult = await coordinatorAgent.run(synthesisPrompt, synthTraceOptions)
     agentResults.set('coordinator', synthesisResult)
+    cumulativeUsage = addUsage(cumulativeUsage, synthesisResult.tokenUsage)
+    if (
+      maxTokenBudget !== undefined
+      && cumulativeUsage.input_tokens + cumulativeUsage.output_tokens > maxTokenBudget
+    ) {
+      this.config.onProgress?.({
+        type: 'budget_exceeded',
+        agent: 'coordinator',
+        data: new TokenBudgetExceededError(
+          'coordinator',
+          cumulativeUsage.input_tokens + cumulativeUsage.output_tokens,
+          maxTokenBudget,
+        ),
+      })
+    }
 
     this.config.onProgress?.({
       type: 'agent_complete',
@@ -808,6 +899,10 @@ export class OpenMultiAgent {
       config: this.config,
       runId: this.config.onTrace ? generateRunId() : undefined,
       abortSignal: options?.abortSignal,
+      cumulativeUsage: ZERO_USAGE,
+      maxTokenBudget: this.config.maxTokenBudget,
+      budgetExceededTriggered: false,
+      budgetExceededReason: undefined,
     }
 
     await executeQueue(queue, ctx)

src/types.ts

@@ -90,11 +90,12 @@ export interface LLMResponse {
  * - `text`        — incremental text delta
  * - `tool_use`    — the model has begun or completed a tool-use block
  * - `tool_result` — a tool result has been appended to the stream
+ * - `budget_exceeded` — token budget threshold reached for this run
  * - `done`        — the stream has ended; `data` is the final {@link LLMResponse}
  * - `error`       — an unrecoverable error occurred; `data` is an `Error`
  */
 export interface StreamEvent {
-  readonly type: 'text' | 'tool_use' | 'tool_result' | 'loop_detected' | 'done' | 'error'
+  readonly type: 'text' | 'tool_use' | 'tool_result' | 'loop_detected' | 'budget_exceeded' | 'done' | 'error'
   readonly data: unknown
 }
 
@@ -208,6 +209,8 @@ export interface AgentConfig {
   readonly tools?: readonly string[]
   readonly maxTurns?: number
   readonly maxTokens?: number
+  /** Maximum cumulative tokens (input + output) allowed for this run. */
+  readonly maxTokenBudget?: number
   readonly temperature?: number
   /**
    * Maximum wall-clock time (in milliseconds) for the entire agent run.
@@ -307,6 +310,8 @@ export interface AgentRunResult {
   readonly structured?: unknown
   /** True when the run was terminated or warned due to loop detection. */
   readonly loopDetected?: boolean
+  /** True when the run stopped because token budget was exceeded. */
+  readonly budgetExceeded?: boolean
 }
 
 // ---------------------------------------------------------------------------
@@ -375,6 +380,7 @@ export interface OrchestratorEvent {
     | 'task_complete'
     | 'task_skipped'
     | 'task_retry'
+    | 'budget_exceeded'
     | 'message'
     | 'error'
   readonly agent?: string
@@ -385,6 +391,8 @@ export interface OrchestratorEvent {
 /** Top-level configuration for the orchestrator. */
 export interface OrchestratorConfig {
   readonly maxConcurrency?: number
+  /** Maximum cumulative tokens (input + output) allowed per orchestrator run. */
+  readonly maxTokenBudget?: number
   readonly defaultModel?: string
   readonly defaultProvider?: 'anthropic' | 'copilot' | 'grok' | 'openai' | 'gemini'
   readonly defaultBaseURL?: string

tests/agent-hooks.test.ts

@@ -4,7 +4,7 @@ import { Agent } from '../src/agent/agent.js'
 import { AgentRunner } from '../src/agent/runner.js'
 import { ToolRegistry } from '../src/tool/framework.js'
 import { ToolExecutor } from '../src/tool/executor.js'
-import type { AgentConfig, AgentRunResult, LLMAdapter, LLMMessage, LLMResponse } from '../src/types.js'
+import type { AgentConfig, AgentRunResult, LLMAdapter, LLMMessage, LLMResponse, StreamEvent } from '../src/types.js'
 
 // ---------------------------------------------------------------------------
 // Mock helpers
@@ -243,7 +243,7 @@ describe('Agent hooks — beforeRun / afterRun', () => {
     }
     const { agent, calls } = buildMockAgent(config, 'streamed')
 
-    const events = []
+    const events: StreamEvent[] = []
     for await (const event of agent.stream('original')) {
       events.push(event)
     }
@@ -263,7 +263,7 @@ describe('Agent hooks — beforeRun / afterRun', () => {
     }
     const { agent } = buildMockAgent(config, 'original')
 
-    const events = []
+    const events: StreamEvent[] = [] 
     for await (const event of agent.stream('hi')) {
       events.push(event)
     }
@@ -280,7 +280,7 @@ describe('Agent hooks — beforeRun / afterRun', () => {
     }
     const { agent } = buildMockAgent(config, 'unreachable')
 
-    const events = []
+    const events: StreamEvent[] = []
     for await (const event of agent.stream('hi')) {
       events.push(event)
     }
@@ -297,7 +297,7 @@ describe('Agent hooks — beforeRun / afterRun', () => {
     }
     const { agent } = buildMockAgent(config, 'streamed output')
 
-    const events = []
+    const events: StreamEvent[] = []
     for await (const event of agent.stream('hi')) {
       events.push(event)
     }

tests/token-budget.test.ts

@@ -0,0 +1,245 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+import { OpenMultiAgent } from '../src/orchestrator/orchestrator.js'
+import { Agent } from '../src/agent/agent.js'
+import { ToolRegistry } from '../src/tool/framework.js'
+import { ToolExecutor } from '../src/tool/executor.js'
+import type { AgentConfig, LLMChatOptions, LLMMessage, LLMResponse, OrchestratorEvent } from '../src/types.js'
+
+let mockAdapterResponses: string[] = []
+let mockAdapterUsage: Array<{ input_tokens: number; output_tokens: number }> = []
+
+vi.mock('../src/llm/adapter.js', () => ({
+  createAdapter: async () => {
+    let callIndex = 0
+    return {
+      name: 'mock',
+      async chat(_msgs: LLMMessage[], options: LLMChatOptions): Promise<LLMResponse> {
+        const text = mockAdapterResponses[callIndex] ?? 'default mock response'
+        const usage = mockAdapterUsage[callIndex] ?? { input_tokens: 10, output_tokens: 20 }
+        callIndex++
+        return {
+          id: `resp-${callIndex}`,
+          content: [{ type: 'text', text }],
+          model: options.model ?? 'mock-model',
+          stop_reason: 'end_turn',
+          usage,
+        }
+      },
+      async *stream() {
+        yield { type: 'done' as const, data: {} }
+      },
+    }
+  },
+}))
+
+function agentConfig(name: string, maxTokenBudget?: number): AgentConfig {
+  return {
+    name,
+    model: 'mock-model',
+    provider: 'openai',
+    systemPrompt: `You are ${name}.`,
+    maxTokenBudget,
+  }
+}
+
+describe('token budget enforcement', () => {
+  beforeEach(() => {
+    mockAdapterResponses = []
+    mockAdapterUsage = []
+  })
+
+  it('enforces agent-level maxTokenBudget in runAgent', async () => {
+    mockAdapterResponses = ['over budget']
+    mockAdapterUsage = [{ input_tokens: 20, output_tokens: 15 }]
+
+    const events: OrchestratorEvent[] = []
+    const oma = new OpenMultiAgent({
+      defaultModel: 'mock-model',
+      onProgress: e => events.push(e),
+    })
+
+    const result = await oma.runAgent(agentConfig('solo', 30), 'test')
+
+    expect(result.success).toBe(false)
+    expect(result.budgetExceeded).toBe(true)
+    expect(result.messages).toHaveLength(1)
+    expect(result.messages[0]?.role).toBe('assistant')
+    expect(result.messages[0]?.content[0]).toMatchObject({ type: 'text', text: 'over budget' })
+    expect(events.some(e => e.type === 'budget_exceeded')).toBe(true)
+  })
+
+  it('emits budget_exceeded stream event without error transition', async () => {
+    mockAdapterResponses = ['over budget']
+    mockAdapterUsage = [{ input_tokens: 20, output_tokens: 15 }]
+
+    const agent = new Agent(
+      agentConfig('streamer', 30),
+      new ToolRegistry(),
+      new ToolExecutor(new ToolRegistry()),
+    )
+
+    const eventTypes: string[] = []
+    for await (const event of agent.stream('test')) {
+      eventTypes.push(event.type)
+    }
+
+    expect(eventTypes).toContain('budget_exceeded')
+    expect(eventTypes).toContain('done')
+    expect(eventTypes).not.toContain('error')
+    expect(agent.getState().status).toBe('completed')
+  })
+
+  it('does not skip in-progress sibling tasks when team budget is exceeded mid-batch', async () => {
+    mockAdapterResponses = ['done-a', 'done-b', 'done-c']
+    mockAdapterUsage = [
+      { input_tokens: 15, output_tokens: 10 }, // A => 25
+      { input_tokens: 15, output_tokens: 10 }, // B => 50 total (exceeds 40)
+      { input_tokens: 15, output_tokens: 10 }, // C should never run
+    ]
+
+    const events: OrchestratorEvent[] = []
+    const oma = new OpenMultiAgent({
+      defaultModel: 'mock-model',
+      maxTokenBudget: 40,
+      onProgress: e => events.push(e),
+    })
+    const team = oma.createTeam('team-siblings', {
+      name: 'team-siblings',
+      agents: [agentConfig('worker-a'), agentConfig('worker-b')],
+      sharedMemory: false,
+    })
+
+    await oma.runTasks(team, [
+      { title: 'Task A', description: 'A', assignee: 'worker-a' },
+      { title: 'Task B', description: 'B', assignee: 'worker-b' },
+      { title: 'Task C', description: 'C', assignee: 'worker-a', dependsOn: ['Task A'] },
+    ])
+
+    const completedTaskIds = new Set(
+      events.filter(e => e.type === 'task_complete').map(e => e.task).filter(Boolean) as string[],
+    )
+    const skippedTaskIds = new Set(
+      events.filter(e => e.type === 'task_skipped').map(e => e.task).filter(Boolean) as string[],
+    )
+
+    const overlap = [...completedTaskIds].filter(id => skippedTaskIds.has(id))
+    expect(overlap).toHaveLength(0)
+  })
+
+  it('does not trigger budget events when budget is not exceeded', async () => {
+    mockAdapterResponses = ['done-a', 'done-b']
+    mockAdapterUsage = [
+      { input_tokens: 10, output_tokens: 10 },
+      { input_tokens: 10, output_tokens: 10 },
+    ]
+    const events: OrchestratorEvent[] = []
+    const oma = new OpenMultiAgent({
+      defaultModel: 'mock-model',
+      maxTokenBudget: 100,
+      onProgress: e => events.push(e),
+    })
+    const team = oma.createTeam('team-a', {
+      name: 'team-a',
+      agents: [agentConfig('worker-a'), agentConfig('worker-b')],
+      sharedMemory: false,
+    })
+
+    const result = await oma.runTasks(team, [
+      { title: 'A', description: 'Do A', assignee: 'worker-a' },
+      { title: 'B', description: 'Do B', assignee: 'worker-b', dependsOn: ['A'] },
+    ])
+
+    expect(result.success).toBe(true)
+    expect(events.some(e => e.type === 'budget_exceeded')).toBe(false)
+  })
+
+  it('enforces team budget in runTasks and skips remaining tasks', async () => {
+    mockAdapterResponses = ['done-a', 'done-b', 'done-c']
+    mockAdapterUsage = [
+      { input_tokens: 20, output_tokens: 15 }, // A => 35
+      { input_tokens: 20, output_tokens: 15 }, // B => 70 total (exceeds 60)
+      { input_tokens: 20, output_tokens: 15 }, // C should not run
+    ]
+
+    const events: OrchestratorEvent[] = []
+    const oma = new OpenMultiAgent({
+      defaultModel: 'mock-model',
+      maxTokenBudget: 60,
+      onProgress: e => events.push(e),
+    })
+    const team = oma.createTeam('team-b', {
+      name: 'team-b',
+      agents: [agentConfig('worker')],
+      sharedMemory: false,
+    })
+
+    const result = await oma.runTasks(team, [
+      { title: 'A', description: 'A', assignee: 'worker' },
+      { title: 'B', description: 'B', assignee: 'worker', dependsOn: ['A'] },
+      { title: 'C', description: 'C', assignee: 'worker', dependsOn: ['B'] },
+    ])
+
+    expect(result.totalTokenUsage.input_tokens + result.totalTokenUsage.output_tokens).toBe(70)
+    expect(events.some(e => e.type === 'budget_exceeded')).toBe(true)
+    expect(events.some(e => e.type === 'task_skipped')).toBe(true)
+  })
+
+  it('counts retry token usage before enforcing team budget', async () => {
+    mockAdapterResponses = ['attempt-1', 'attempt-2', 'should-skip']
+    mockAdapterUsage = [
+      { input_tokens: 20, output_tokens: 15 }, // attempt 1
+      { input_tokens: 20, output_tokens: 15 }, // attempt 2
+      { input_tokens: 20, output_tokens: 15 }, // next task (should skip)
+    ]
+
+    const events: OrchestratorEvent[] = []
+    const oma = new OpenMultiAgent({
+      defaultModel: 'mock-model',
+      maxTokenBudget: 50,
+      onProgress: e => events.push(e),
+    })
+    const team = oma.createTeam('team-c', {
+      name: 'team-c',
+      agents: [agentConfig('retry-worker', 1)],
+      sharedMemory: false,
+    })
+
+    const result = await oma.runTasks(team, [
+      { title: 'Retrying task', description: 'Will exceed internal budget', assignee: 'retry-worker', maxRetries: 1 },
+      { title: 'Later task', description: 'Should be skipped', assignee: 'retry-worker', dependsOn: ['Retrying task'] },
+    ])
+
+    expect(result.totalTokenUsage.input_tokens + result.totalTokenUsage.output_tokens).toBe(70)
+    expect(events.some(e => e.type === 'budget_exceeded')).toBe(true)
+    expect(events.some(e => e.type === 'error')).toBe(true)
+  })
+
+  it('enforces orchestrator budget in runTeam', async () => {
+    mockAdapterResponses = [
+      '```json\n[{"title":"Task A","description":"Do A","assignee":"worker"}]\n```',
+      'worker result',
+      'synthesis should not run when budget exceeded',
+    ]
+    mockAdapterUsage = [
+      { input_tokens: 20, output_tokens: 15 }, // decomposition => 35
+      { input_tokens: 20, output_tokens: 15 }, // task => 70 total (exceeds 60)
+      { input_tokens: 20, output_tokens: 15 }, // synthesis should not execute
+    ]
+
+    const events: OrchestratorEvent[] = []
+    const oma = new OpenMultiAgent({
+      defaultModel: 'mock-model',
+      maxTokenBudget: 60,
+      onProgress: e => events.push(e),
+    })
+    const team = oma.createTeam('team-d', {
+      name: 'team-d',
+      agents: [agentConfig('worker')],
+      sharedMemory: false,
+    })
+
+    const result = await oma.runTeam(team, 'Do work')
+    expect(result.totalTokenUsage.input_tokens + result.totalTokenUsage.output_tokens).toBe(70)
+    expect(events.some(e => e.type === 'budget_exceeded')).toBe(true)
+  })
+})

tests/trace.test.ts

@@ -186,7 +186,7 @@ describe('AgentRunner trace events', () => {
     })
 
     const runOptions: RunOptions = {
-      onTrace: (e) => traces.push(e),
+      onTrace: (e) => { traces.push(e) },
       runId: 'run-1',
       traceAgent: 'test-agent',
     }
@@ -234,7 +234,7 @@ describe('AgentRunner trace events', () => {
 
     await runner.run(
       [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
-      { onTrace: (e) => traces.push(e), runId: 'run-2', traceAgent: 'tooler' },
+      { onTrace: (e) => { traces.push(e) }, runId: 'run-2', traceAgent: 'tooler' },
     )
 
     const toolTraces = traces.filter(t => t.type === 'tool_call')
@@ -273,7 +273,7 @@ describe('AgentRunner trace events', () => {
 
     await runner.run(
       [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
-      { onTrace: (e) => traces.push(e), runId: 'run-3', traceAgent: 'err-agent' },
+      { onTrace: (e) => { traces.push(e) }, runId: 'run-3', traceAgent: 'err-agent' },
     )
 
     const toolTraces = traces.filter(t => t.type === 'tool_call')
@@ -316,7 +316,7 @@ describe('Agent trace events', () => {
     const agent = buildMockAgent(config, [textResponse('Hello world')])
 
     const runOptions: Partial<RunOptions> = {
-      onTrace: (e) => traces.push(e),
+      onTrace: (e) => { traces.push(e) },
       runId: 'run-agent-1',
       traceAgent: 'my-agent',
     }
@@ -367,7 +367,7 @@ describe('Agent trace events', () => {
 
     const runId = 'shared-run-id'
     await agent.run('test', {
-      onTrace: (e) => traces.push(e),
+      onTrace: (e) => { traces.push(e) },
       runId,
       traceAgent: 'multi-trace-agent',
     })
@@ -436,7 +436,7 @@ describe('Agent trace events', () => {
 
     await runner.run(
       [{ role: 'user', content: [{ type: 'text', text: 'go' }] }],
-      { onTrace: (e) => traces.push(e), runId: 'run-tok', traceAgent: 'token-agent' },
+      { onTrace: (e) => { traces.push(e) }, runId: 'run-tok', traceAgent: 'token-agent' },
     )
 
     const llmTraces = traces.filter(t => t.type === 'llm_call')