@namzu/sdk 34.2.0 → 36.0.0

package/src/agents/PipelineAgent.ts

@@ -26,7 +26,7 @@
  * That is the bar a reopen has to clear.
  */
 
-import { EMPTY_TOKEN_USAGE } from '../constants/limits.js'
+import { withTokenBudget } from '../provider/token-budget.js'
 import type {
 	AgentInput,
 	AgentMetadata,
@@ -40,6 +40,7 @@ import { ZERO_COST } from '../utils/cost.js'
 import { toErrorMessage } from '../utils/error.js'
 import type { Logger } from '../utils/logger.js'
 import { AbstractAgent } from './AbstractAgent.js'
+import { resolveAgentBudget } from './budget.js'
 
 export class PipelineAgent extends AbstractAgent<PipelineAgentConfig, PipelineAgentResult> {
 	readonly type = 'pipeline' as const
@@ -92,127 +93,146 @@ export class PipelineAgent extends AbstractAgent<PipelineAgentConfig, PipelineAg
 		const startTime = Date.now()
 		const runId = this.createRunId()
 		this.bindRun(runId, config.logger)
-		const stepResults: PipelineStepResult[] = []
-		const previousResults = new Map<string, unknown>()
-		let completedSteps = 0
-
-		await this.emitEvent({ type: 'run_started', runId }, listener)
-
-		let currentInput: unknown = input.messages
-			.filter((m) => m.role === 'user')
-			.map((m) => m.content)
-			.filter((c): c is string => c !== null)
-			.join('\n')
-
-		for (let i = 0; i < config.steps.length; i++) {
-			const step = config.steps[i]
-			if (!step) throw new Error(`Pipeline step at index ${i} is undefined`)
-			const stepStart = Date.now()
-
-			if (this.abortController.signal.aborted) {
-				stepResults.push({
-					stepName: step.name,
-					status: 'skipped',
-					durationMs: 0,
-				})
-				continue
-			}
-
-			await this.emitEvent({ type: 'iteration_started', runId, iteration: i + 1 }, listener)
-
-			const context: StepContext = {
-				runId,
-				stepIndex: i,
-				totalSteps: config.steps.length,
-				previousResults,
-				provider: config.provider,
-				signal: this.abortController.signal,
-				env: config.env ?? {},
-			}
-
-			try {
-				if (step.validate && !step.validate(currentInput)) {
-					throw new Error(`Validation failed for step "${step.name}"`)
+		const budget = await resolveAgentBudget(input, config, runId)
+		const provider = config.provider ? withTokenBudget(config.provider, budget) : undefined
+		try {
+			const stepResults: PipelineStepResult[] = []
+			const previousResults = new Map<string, unknown>()
+			let completedSteps = 0
+
+			await this.emitEvent({ type: 'run_started', runId }, listener)
+
+			let currentInput: unknown = input.messages
+				.filter((m) => m.role === 'user')
+				.map((m) => m.content)
+				.filter((c): c is string => c !== null)
+				.join('\n')
+
+			for (let i = 0; i < config.steps.length; i++) {
+				const step = config.steps[i]
+				if (!step) throw new Error(`Pipeline step at index ${i} is undefined`)
+				const stepStart = Date.now()
+
+				if (this.abortController.signal.aborted) {
+					stepResults.push({
+						stepName: step.name,
+						status: 'skipped',
+						durationMs: 0,
+					})
+					continue
 				}
 
-				const output = await step.execute(currentInput, context)
-				previousResults.set(step.name, output)
-				currentInput = output
-				completedSteps++
-
-				stepResults.push({
-					stepName: step.name,
-					status: 'completed',
-					output,
-					durationMs: Date.now() - stepStart,
-				})
-			} catch (err) {
-				const errorMsg = toErrorMessage(err)
-
-				stepResults.push({
-					stepName: step.name,
-					status: 'failed',
-					error: errorMsg,
-					durationMs: Date.now() - stepStart,
-				})
-
-				if (step.rollback) {
-					try {
-						await step.rollback(currentInput, context)
-					} catch (rollbackErr) {
-						this.log.error('Rollback failed for a step', {
-							'namzu.pipeline.step_name': step.name,
-							'exception.message': toErrorMessage(rollbackErr),
-						})
-					}
+				await this.emitEvent({ type: 'iteration_started', runId, iteration: i + 1 }, listener)
+
+				const context: StepContext = {
+					runId,
+					stepIndex: i,
+					totalSteps: config.steps.length,
+					previousResults,
+					provider,
+					budget,
+					signal: this.abortController.signal,
+					env: config.env ?? {},
 				}
 
-				if (!config.continueOnError) {
-					await this.emitEvent({ type: 'run_failed', runId, error: errorMsg }, listener)
+				try {
+					if (step.validate && !step.validate(currentInput)) {
+						throw new Error(`Validation failed for step "${step.name}"`)
+					}
 
-					return {
-						runId,
+					const output = await step.execute(currentInput, context)
+					previousResults.set(step.name, output)
+					currentInput = output
+					completedSteps++
+
+					stepResults.push({
+						stepName: step.name,
+						status: 'completed',
+						output,
+						durationMs: Date.now() - stepStart,
+					})
+				} catch (err) {
+					const errorMsg = toErrorMessage(err)
+
+					stepResults.push({
+						stepName: step.name,
 						status: 'failed',
-						stopReason: 'error',
-						usage: { ...EMPTY_TOKEN_USAGE },
-						cost: { ...ZERO_COST },
-						iterations: i + 1,
-						durationMs: Date.now() - startTime,
-						messages: input.messages,
-						lastError: errorMsg,
-						stepResults,
-						completedSteps,
-						totalSteps: config.steps.length,
+						error: errorMsg,
+						durationMs: Date.now() - stepStart,
+					})
+
+					if (step.rollback) {
+						try {
+							await step.rollback(currentInput, context)
+						} catch (rollbackErr) {
+							this.log.error('Rollback failed for a step', {
+								'namzu.pipeline.step_name': step.name,
+								'exception.message': toErrorMessage(rollbackErr),
+							})
+						}
+					}
+
+					if (!config.continueOnError) {
+						await this.emitEvent(
+							{
+								type: 'run_failed',
+								runId,
+								error: errorMsg,
+								budget: budget.summary(),
+							},
+							listener,
+						)
+
+						return {
+							runId,
+							status: 'failed',
+							stopReason: 'error',
+							usage: budget.ownUsage,
+							budget: budget.summary(),
+							cost: { ...ZERO_COST, unpricedTokens: budget.ownTokens },
+							iterations: i + 1,
+							durationMs: Date.now() - startTime,
+							messages: input.messages,
+							lastError: errorMsg,
+							stepResults,
+							completedSteps,
+							totalSteps: config.steps.length,
+						}
 					}
 				}
 			}
-		}
 
-		const finalStatus = completedSteps === config.steps.length ? 'completed' : 'failed'
-		const lastOutput = stepResults[stepResults.length - 1]?.output
+			const finalStatus = completedSteps === config.steps.length ? 'completed' : 'failed'
+			const lastOutput = stepResults[stepResults.length - 1]?.output
 
-		await this.emitEvent(
-			{
-				type: 'run_completed',
+			await this.emitEvent(
+				{
+					type: 'run_completed',
+					budget: budget.summary(),
+					runId,
+					result: typeof lastOutput === 'string' ? lastOutput : JSON.stringify(lastOutput),
+				},
+				listener,
+			)
+
+			return {
 				runId,
+				status: finalStatus,
+				stopReason: 'end_turn',
+				usage: budget.ownUsage,
+				budget: budget.summary(),
+				cost: { ...ZERO_COST, unpricedTokens: budget.ownTokens },
+				iterations: config.steps.length,
+				durationMs: Date.now() - startTime,
+				messages: input.messages,
 				result: typeof lastOutput === 'string' ? lastOutput : JSON.stringify(lastOutput),
-			},
-			listener,
-		)
-
-		return {
-			runId,
-			status: finalStatus,
-			stopReason: 'end_turn',
-			usage: { ...EMPTY_TOKEN_USAGE },
-			cost: { ...ZERO_COST },
-			iterations: config.steps.length,
-			durationMs: Date.now() - startTime,
-			messages: input.messages,
-			result: typeof lastOutput === 'string' ? lastOutput : JSON.stringify(lastOutput),
-			stepResults,
-			completedSteps,
-			totalSteps: config.steps.length,
+				stepResults,
+				completedSteps,
+				totalSteps: config.steps.length,
+			}
+		} finally {
+			budget.settle()
+			await budget.flush()
 		}
 	}
 }

package/src/agents/ReactiveAgent.ts

@@ -70,6 +70,7 @@ export class ReactiveAgent extends AbstractAgent<ReactiveAgentConfig, ReactiveAg
 				skills: config.skills,
 				basePrompt: config.basePrompt,
 				provider: config.provider,
+				...(config.budget ? { budget: config.budget } : {}),
 				tools: config.tools,
 				...(input.attachmentStore ? { attachmentStore: input.attachmentStore } : {}),
 				...(config.attachmentResolveTimeoutMs !== undefined
@@ -186,6 +187,7 @@ export class ReactiveAgent extends AbstractAgent<ReactiveAgentConfig, ReactiveAg
 			status: run.status,
 			stopReason: run.stopReason,
 			usage: run.tokenUsage,
+			...(run.budget ? { budget: run.budget } : {}),
 			cost: run.costInfo,
 			iterations: run.currentIteration,
 			durationMs: Date.now() - startTime,

package/src/agents/RouterAgent.ts

@@ -1,6 +1,7 @@
 import { EMPTY_TOKEN_USAGE } from '../constants/limits.js'
 import { collectChatCompletion } from '../provider/collect-chat-completion.js'
 import { resolveStreamIdleTimeoutMs, withStreamIdleTimeout } from '../provider/idle-timeout.js'
+import { withTokenBudget } from '../provider/token-budget.js'
 import { FallbackResolver } from '../runtime/decision/fallback.js'
 import { DecisionParser } from '../runtime/decision/parser.js'
 import type {
@@ -18,6 +19,7 @@ import type { RunEventListener } from '../types/run/index.js'
 import { ZERO_COST } from '../utils/cost.js'
 import type { Logger } from '../utils/logger.js'
 import { AbstractAgent } from './AbstractAgent.js'
+import { resolveAgentBudget } from './budget.js'
 
 export class RouterAgent extends AbstractAgent<RouterAgentConfig, RouterAgentResult> {
 	readonly type = 'router' as const
@@ -71,88 +73,124 @@ export class RouterAgent extends AbstractAgent<RouterAgentConfig, RouterAgentRes
 		const startTime = Date.now()
 		const runId = this.createRunId()
 		this.bindRun(runId, config.logger)
+		const budget = await resolveAgentBudget(input, config, runId)
+		const budgetConfig = { ...config, budget }
+		try {
+			await this.emitEvent({ type: 'run_started', runId }, listener)
 
-		await this.emitEvent({ type: 'run_started', runId }, listener)
+			const decision = await this.route(input, budgetConfig, streamIdleTimeoutMs, signal)
 
-		const decision = await this.route(input, config, streamIdleTimeoutMs, signal)
+			let targetRoute = config.routes.find((r) => r.agentId === decision.agentId)
 
-		let targetRoute = config.routes.find((r) => r.agentId === decision.agentId)
+			if (!targetRoute) {
+				const fallback = config.fallbackAgentId
+					? config.routes.find((r) => r.agentId === config.fallbackAgentId)
+					: undefined
 
-		if (!targetRoute) {
-			const fallback = config.fallbackAgentId
-				? config.routes.find((r) => r.agentId === config.fallbackAgentId)
-				: undefined
+				if (!fallback) {
+					const errorMsg = `No route found for "${decision.agentId}"`
 
-			if (!fallback) {
-				const errorMsg = `No route found for "${decision.agentId}"`
+					await this.emitEvent(
+						{
+							type: 'run_failed',
+							runId,
+							error: errorMsg,
+							budget: budget.summary(),
+						},
+						listener,
+					)
 
-				await this.emitEvent({ type: 'run_failed', runId, error: errorMsg }, listener)
-
-				return {
-					runId,
-					status: 'failed',
-					stopReason: 'error',
-					usage: { ...EMPTY_TOKEN_USAGE },
-					cost: { ...ZERO_COST },
-					iterations: 1,
-					durationMs: Date.now() - startTime,
-					messages: input.messages,
-					lastError: errorMsg,
-					selectedRoute: decision.agentId,
-					routingDecision: decision,
-					delegateResult: {
+					return {
 						runId,
 						status: 'failed',
-						usage: { ...EMPTY_TOKEN_USAGE },
-						cost: { ...ZERO_COST },
-						iterations: 0,
-						durationMs: 0,
-						messages: [],
-					},
+						stopReason: 'error',
+						usage: budget.ownUsage,
+						budget: budget.summary(),
+						cost: { ...ZERO_COST, unpricedTokens: budget.ownTokens },
+						iterations: 1,
+						durationMs: Date.now() - startTime,
+						messages: input.messages,
+						lastError: errorMsg,
+						selectedRoute: decision.agentId,
+						routingDecision: decision,
+						delegateResult: {
+							runId,
+							status: 'failed',
+							usage: { ...EMPTY_TOKEN_USAGE },
+							cost: { ...ZERO_COST },
+							iterations: 0,
+							durationMs: 0,
+							messages: [],
+						},
+					}
 				}
+
+				decision.agentId = fallback.agentId
+				targetRoute = fallback
 			}
 
-			decision.agentId = fallback.agentId
-			targetRoute = fallback
-		}
+			if (!config.invocationState) {
+				throw new Error(
+					'RouterAgent requires invocationState with tenantId in config (session-hierarchy.md §12.1).',
+				)
+			}
+			const childInvocationState = deriveChildState(config.invocationState, this.metadata.id)
+			const allocation = Number.isFinite(budget.remaining)
+				? Math.floor(budget.remaining)
+				: config.tokenBudget || 200_000
+			const childBudget = budget.reserve(allocation)
+			await childBudget.flush()
+			let delegateResult: RouterAgentResult['delegateResult']
+			try {
+				delegateResult = await targetRoute.agent.run(
+					input,
+					{
+						...config,
+						budget: childBudget,
+						tokenBudget: allocation,
+						parentRunId: runId,
+						depth: (config.depth ?? 0) + 1,
+						invocationState: childInvocationState,
+					},
+					listener,
+				)
+				childBudget.bindRun(delegateResult.runId)
+				childBudget.settle(delegateResult.usage.totalTokens)
+			} finally {
+				// A thrown custom delegate supplies no final usage receipt.
+				// Its reservation stays held until an authoritative settlement.
+				await childBudget.flush()
+			}
 
-		if (!config.invocationState) {
-			throw new Error(
-				'RouterAgent requires invocationState with tenantId in config (session-hierarchy.md §12.1).',
+			await this.emitEvent(
+				{
+					type: 'run_completed',
+					budget: budget.summary(),
+					runId,
+					result: delegateResult.result ?? '',
+				},
+				listener,
 			)
-		}
-		const childInvocationState = deriveChildState(config.invocationState, this.metadata.id)
-		const delegateResult = await targetRoute.agent.run(
-			input,
-			{ ...config, invocationState: childInvocationState },
-			listener,
-		)
 
-		await this.emitEvent(
-			{
-				type: 'run_completed',
+			return {
 				runId,
-				result: delegateResult.result ?? '',
-			},
-			listener,
-		)
-
-		return {
-			runId,
-			status: delegateResult.status,
-			stopReason: delegateResult.stopReason,
-			// Routing is a model call the run paid for; reporting only the
-			// delegate's usage silently under-reports every routed run.
-			usage: accumulateTokenUsage(delegateResult.usage, decision.usage ?? EMPTY_TOKEN_USAGE),
-			cost: delegateResult.cost,
-			iterations: delegateResult.iterations + 1,
-			durationMs: Date.now() - startTime,
-			messages: delegateResult.messages,
-			result: delegateResult.result,
-			lastError: delegateResult.lastError,
-			selectedRoute: decision.agentId,
-			routingDecision: decision,
-			delegateResult,
+				status: delegateResult.status,
+				stopReason: delegateResult.stopReason,
+				usage: budget.ownUsage,
+				budget: budget.summary(),
+				cost: { ...ZERO_COST, unpricedTokens: budget.ownTokens },
+				iterations: delegateResult.iterations + 1,
+				durationMs: Date.now() - startTime,
+				messages: delegateResult.messages,
+				result: delegateResult.result,
+				lastError: delegateResult.lastError,
+				selectedRoute: decision.agentId,
+				routingDecision: decision,
+				delegateResult,
+			}
+		} finally {
+			budget.settle()
+			await budget.flush()
 		}
 	}
 
@@ -197,7 +235,10 @@ export class RouterAgent extends AbstractAgent<RouterAgentConfig, RouterAgentRes
 		}
 
 		if (config.fallbackAgentId) {
-			fallbackStrategies.push({ type: 'fixed', agentId: config.fallbackAgentId })
+			fallbackStrategies.push({
+				type: 'fixed',
+				agentId: config.fallbackAgentId,
+			})
 		}
 
 		fallbackStrategies.push({ type: 'first_route' })
@@ -221,10 +262,13 @@ export class RouterAgent extends AbstractAgent<RouterAgentConfig, RouterAgentRes
 		// Every routing attempt is a billed model call; a fallback after
 		// three failed parses still cost three calls.
 		let routingUsage: TokenUsage = { ...EMPTY_TOKEN_USAGE }
-		const routingProvider = withStreamIdleTimeout(config.provider, {
-			idleTimeoutMs: streamIdleTimeoutMs,
-			log,
-		})
+		const routingProvider = withStreamIdleTimeout(
+			config.budget ? withTokenBudget(config.provider, config.budget) : config.provider,
+			{
+				idleTimeoutMs: streamIdleTimeoutMs,
+				log,
+			},
+		)
 
 		for (let attempt = 0; attempt < maxRetries; attempt++) {
 			try {
@@ -297,6 +341,9 @@ export class RouterAgent extends AbstractAgent<RouterAgentConfig, RouterAgentRes
 			}
 		}
 
-		return { ...fallbackResolver.resolve(userContent, validAgentIds), usage: routingUsage }
+		return {
+			...fallbackResolver.resolve(userContent, validAgentIds),
+			usage: routingUsage,
+		}
 	}
 }

package/src/agents/SupervisorAgent.ts

@@ -21,6 +21,7 @@ import type { ActorRef } from '../types/session/actor.js'
 import { ZERO_COST } from '../utils/cost.js'
 import type { Logger } from '../utils/logger.js'
 import { AbstractAgent } from './AbstractAgent.js'
+import { resolveAgentBudget } from './budget.js'
 
 /**
  * Build the authoritative per-task ledger from the gateway's task handles.
@@ -131,6 +132,14 @@ export class SupervisorAgent extends AbstractAgent<SupervisorAgentConfig, Superv
 		// on one path and `undefined` on another — a half-migration that
 		// fails silently, which is worse than not renaming the field.
 		const configuredScheduler = config.scheduler
+		const budget = await resolveAgentBudget(
+			input,
+			{ ...config, budget: config.budget ?? configuredScheduler?.budget },
+			runId,
+		)
+		if (configuredScheduler && configuredScheduler.budget !== budget) {
+			throw new Error('Injected task scheduler must share the supervisor token budget authority')
+		}
 
 		let gateway: TaskScheduler
 		if (configuredScheduler) {
@@ -155,10 +164,7 @@ export class SupervisorAgent extends AbstractAgent<SupervisorAgentConfig, Superv
 				// the increment when it constructs the child. Resetting a delegated
 				// supervisor to zero here made every grandchild depth one again.
 				depth: config.depth ?? 0,
-				budgetTracker: {
-					total: config.tokenBudget,
-					remaining: config.tokenBudget,
-				},
+				budget,
 				factoryOptions: mergedFactoryOptions,
 				// The supervisor already hands this to its OWN run. Handing it to
 				// the spawn context makes a worker's REVIEW-tier calls reach the
@@ -365,6 +371,7 @@ export class SupervisorAgent extends AbstractAgent<SupervisorAgentConfig, Superv
 					runtimeToolOverrides: input.runtimeToolOverrides,
 					runtimeContext: input.runtimeContext,
 					taskScheduler: gateway,
+					budget,
 					completionInbox,
 					advisory: config.advisory,
 					invocationState: childInvocationState,
@@ -414,6 +421,7 @@ export class SupervisorAgent extends AbstractAgent<SupervisorAgentConfig, Superv
 				status: run.status === 'completed' ? 'completed' : 'failed',
 				stopReason: run.stopReason,
 				usage: run.tokenUsage,
+				budget: budget.summary(),
 				cost: run.costInfo,
 				iterations: run.currentIteration,
 				durationMs: Date.now() - startTime,

package/src/agents/budget.ts

@@ -0,0 +1,42 @@
+import { TokenBudget } from '../run/token-budget.js'
+import { openTokenBudget } from '../store/run/token-budget-disk.js'
+import type { AgentInput, BaseAgentConfig } from '../types/agent/base.js'
+import type { RunId } from '../types/ids/index.js'
+
+/** Resolve one authority before a composite agent constructs its scheduler. */
+export async function resolveAgentBudget(
+	input: AgentInput,
+	config: BaseAgentConfig,
+	runId: RunId,
+): Promise<TokenBudget> {
+	if (!Number.isSafeInteger(config.tokenBudget) || config.tokenBudget < 0) {
+		throw new Error('Agent token budget must be a nonnegative safe integer')
+	}
+	if (config.budget) {
+		config.budget.bindRun(runId)
+		if (
+			config.tokenBudget > 0 &&
+			(config.budget.limit === 0 || config.tokenBudget < config.budget.limit)
+		) {
+			config.budget.narrow(config.tokenBudget)
+		}
+		await config.budget.flush()
+		return config.budget
+	}
+	if (config.parentRunId) {
+		throw new Error('A delegated agent requires its inherited token budget authority')
+	}
+	const { tenantId, projectId, sessionId } = config
+	const budget =
+		tenantId && projectId && sessionId
+			? await openTokenBudget({
+					scope: { tenantId, projectId, sessionId, runId },
+					limit: config.tokenBudget,
+					pathBuilder: config.pathBuilder,
+					workingDirectory: input.workingDirectory,
+				})
+			: TokenBudget.create(config.tokenBudget, runId)
+	budget.bindRun(runId)
+	await budget.flush()
+	return budget
+}

package/src/agents/explore.ts

@@ -0,0 +1,29 @@
+/**
+ * The delegate that can only look.
+ *
+ * "Find where X is defined", "which files reference Y", "how does Z work" are
+ * the delegations a parent makes most, and a child built from the parent's
+ * whole working set answers them holding `write` and `bash` — so a lookup
+ * needs the same approvals as a change. A host that offers this delegate
+ * builds its roster with `filterReadOnlyTools` and its prompt from the text
+ * here; the two are exported separately because a host adds its own
+ * doctrine and environment around the prompt.
+ */
+
+export const EXPLORE_AGENT_ID = 'explore'
+
+export const EXPLORE_AGENT_DESCRIPTION =
+	'A read-only sub-agent for finding files, symbols and answers.'
+
+/**
+ * What an explore delegate is told. Names its limits in the first person so
+ * a model does not describe a change as made, and asks for file:line
+ * references because that is what the parent acts on.
+ */
+export const EXPLORE_AGENT_PROMPT = [
+	'You are a read-only sub-agent dispatched to find things out: where something is defined, which files reference it, how a piece of code works, what a directory contains.',
+	'You cannot see the parent conversation — work only from the prompt you were given.',
+	'You have reading and searching tools only. You cannot edit, write or run commands, and you must not describe a change as made.',
+	'Search broadly first, then read what matters. Report file paths with line numbers, quote the relevant lines, and say plainly what you did not find.',
+	'End with a concise answer the parent can act on; do not ask the parent questions.',
+].join('\n')