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>
This commit is contained in:
Ibrahim Kazimov
GitHub
parent
1e3bd1013e
commit
60fb2b142e
|
|
@ -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 } : {}),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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'
|
||||
}
|
||||
}
|
||||
|
|
@ -107,6 +107,7 @@ export {
|
|||
|
||||
export { createAdapter } from './llm/adapter.js'
|
||||
export type { SupportedProvider } from './llm/adapter.js'
|
||||
export { TokenBudgetExceededError } from './errors.js'
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Memory
|
||||
|
|
|
|||
|
|
@ -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'>
|
||||
> & Pick<OrchestratorConfig, 'onApproval' | 'onProgress' | 'onTrace' | 'defaultBaseURL' | 'defaultApiKey'>
|
||||
Omit<OrchestratorConfig, 'onApproval' | 'onProgress' | 'onTrace' | 'defaultBaseURL' | 'defaultApiKey' | 'maxTokenBudget'>
|
||||
> & 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)
|
||||
|
|
|
|||
10
src/types.ts
10
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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
})
|
||||
|
|
@ -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')
|
||||
|
|
|
|||
Loading…
Reference in New Issue