diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 904b5bc..8b02641 100644 --- a/src/agent/agent.ts +++ b/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 } : {}), } } diff --git a/src/agent/runner.ts b/src/agent/runner.ts index f7f1d6a..4ca0975 100644 --- a/src/agent/runner.ts +++ b/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 diff --git a/src/errors.ts b/src/errors.ts new file mode 100644 index 0000000..be7ac67 --- /dev/null +++ b/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' + } +} diff --git a/src/index.ts b/src/index.ts index 115ab2e..20d8d1a 100644 --- a/src/index.ts +++ b/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 diff --git a/src/orchestrator/orchestrator.ts b/src/orchestrator/orchestrator.ts index 1526827..8b0073d 100644 --- a/src/orchestrator/orchestrator.ts +++ b/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 { */ export class OpenMultiAgent { private readonly config: Required< - Omit - > & Pick + Omit + > & Pick private readonly teams: Map = 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 { + 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() 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 | 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) diff --git a/src/types.ts b/src/types.ts index 3ffbfee..e43bfbd 100644 --- a/src/types.ts +++ b/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 diff --git a/tests/agent-hooks.test.ts b/tests/agent-hooks.test.ts index 13044a3..285da4c 100644 --- a/tests/agent-hooks.test.ts +++ b/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) } diff --git a/tests/token-budget.test.ts b/tests/token-budget.test.ts new file mode 100644 index 0000000..72cde15 --- /dev/null +++ b/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 { + 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) + }) +}) diff --git a/tests/trace.test.ts b/tests/trace.test.ts index fbeb78c..00e8330 100644 --- a/tests/trace.test.ts +++ b/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 = { - 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')