@namzu/sdk 2.0.0 → 3.0.0

package/src/manager/agent/__tests__/lifecycle.test.ts

@@ -1,5 +1,7 @@
 import { describe, expect, it } from 'vitest'
+import { AGENT_MANAGER_DEFAULTS } from '../../../constants/agent/index.js'
 import { EMPTY_TOKEN_USAGE } from '../../../constants/limits.js'
+import { LocalTaskGateway } from '../../../gateway/local.js'
 import { AgentRegistry } from '../../../registry/agent/definitions.js'
 import {
 	DefaultCapacityValidator,
@@ -302,6 +304,53 @@ describe('AgentManager.sendMessage — Phase 6 SubSession spawn', () => {
 		).rejects.toBeInstanceOf(DelegationCapacityExceeded)
 	})
 
+	it('width: two concurrent spawns cannot both slip past the last slot', async () => {
+		const childAgent = makeAgent('child-1', async () => successResult())
+		const harness = await buildHarness(childAgent)
+
+		// Fill to one below the cap, so exactly one more child fits.
+		for (let i = 0; i < 7; i++) {
+			const sibling = await harness.store.createSession(
+				{
+					threadId: harness.threadId,
+					projectId: harness.projectId,
+					currentActor: agentActor('sibling'),
+				},
+				tenant,
+			)
+			await harness.store.createSubSession(
+				{
+					parentSessionId: harness.parentSession.id,
+					childSessionId: sibling.id,
+					kind: 'agent_spawn',
+					spawnedBy: user(),
+				},
+				tenant,
+			)
+		}
+
+		// Both read the count before either writes. The check and the write
+		// that invalidates it used to have every other provisioning step
+		// between them, so a cap of 8 admitted 9.
+		const attempts = await Promise.allSettled([
+			harness.manager.sendMessage(
+				buildOptions('child-1', harness.parentSession.id, harness.projectId),
+				buildContext(harness.parentSession.id, harness.projectId, harness.threadId),
+			),
+			harness.manager.sendMessage(
+				buildOptions('child-1', harness.parentSession.id, harness.projectId),
+				buildContext(harness.parentSession.id, harness.projectId, harness.threadId),
+			),
+		])
+
+		const rejected = attempts.filter((a) => a.status === 'rejected')
+		expect(rejected).toHaveLength(1)
+		expect((rejected[0] as PromiseRejectedResult).reason).toBeInstanceOf(DelegationCapacityExceeded)
+
+		const children = await harness.store.getChildren(harness.parentSession.id, tenant)
+		expect(children.length).toBe(8)
+	})
+
 	it('depth: ancestry chain exceeding maxDelegationDepth (4) rejects with DelegationCapacityExceeded', async () => {
 		const childAgent = makeAgent('child-1', async () => successResult())
 		const harness = await buildHarness(childAgent)
@@ -471,3 +520,251 @@ describe('AgentManager.sendMessage — Phase 6 SubSession spawn', () => {
 // unconditional required. Prior `describe('AgentManager.sendMessage — legacy
 // mode (no session deps)')` block deleted; every spawn now produces a
 // SubSession + Session + WorkspaceRef triple (Convention #0).
+
+/**
+ * Two unit errors in the delegation path that no test covered, because the
+ * numbers involved stay plausible-looking until you check their units and
+ * their identity.
+ */
+describe('AgentManager.sendMessage — budget and deadline arithmetic', () => {
+	it('does not derive the child deadline from the TOKEN budget', async () => {
+		// The fallback used to be `context.budgetTracker.remaining` — a token
+		// count read as milliseconds. It hid for so long because a six-figure
+		// token budget lands in a plausible range of milliseconds.
+		const seen: BaseAgentConfig[] = []
+		const childAgent = makeAgent('child-deadline', async (_input, config) => {
+			seen.push(config)
+			return successResult()
+		})
+		const harness = await buildHarness(childAgent)
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		context.budgetTracker = { total: 100_000, remaining: 100_000 }
+
+		const task = await harness.manager.sendMessage(
+			buildOptions('child-deadline', harness.parentSession.id, harness.projectId),
+			context,
+		)
+		await waitForTask(harness.manager, task.taskId)
+
+		expect(seen).toHaveLength(1)
+		// Before the fix this was the post-debit token remainder (50_000).
+		expect(seen[0]?.timeoutMs).toBe(AGENT_MANAGER_DEFAULTS.childTimeoutMs)
+	})
+
+	it('a nearly-exhausted parent still gives its child a real deadline', async () => {
+		// The edge where the unit error bit hardest: a tiny token remainder
+		// read as milliseconds produced a child that was out of time on
+		// arrival.
+		const seen: BaseAgentConfig[] = []
+		const childAgent = makeAgent('child-small', async (_input, config) => {
+			seen.push(config)
+			return successResult()
+		})
+		const harness = await buildHarness(childAgent)
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		context.budgetTracker = { total: 100_000, remaining: 20 }
+
+		const task = await harness.manager.sendMessage(
+			buildOptions('child-small', harness.parentSession.id, harness.projectId),
+			context,
+		)
+		await waitForTask(harness.manager, task.taskId)
+
+		// Old code: floor(20 * 0.5) = 10 tokens, read as 10 MILLISECONDS.
+		expect(seen[0]?.timeoutMs).toBe(AGENT_MANAGER_DEFAULTS.childTimeoutMs)
+	})
+
+	it('refuses to spawn when the allocation would floor to zero', async () => {
+		// Because `tokenBudget: 0` means UNLIMITED downstream
+		// (`LimitChecker`: `tokenBudget > 0 && total >= tokenBudget`), the
+		// most depleted parent in the tree was the one that spawned an
+		// uncapped child. Budget exhaustion must not invert into no budget.
+		const childAgent = makeAgent('child-broke', async () => successResult())
+		const harness = await buildHarness(childAgent)
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		context.budgetTracker = { total: 100_000, remaining: 1 }
+
+		await expect(
+			harness.manager.sendMessage(
+				buildOptions('child-broke', harness.parentSession.id, harness.projectId),
+				context,
+			),
+		).rejects.toThrow(/allocates 0 to the child/)
+	})
+
+	it('siblings divide ONE budget pool when spawned THROUGH THE GATEWAY', async () => {
+		// `spawn` debits the shared tracker. A caller that hands each spawn a
+		// cloned tracker makes the debit land on a throwaway object, so N
+		// children are each allocated maxBudgetFraction of the SAME number.
+		const seen: BaseAgentConfig[] = []
+		const childAgent = makeAgent('child-budget', async (_input, config) => {
+			seen.push(config)
+			return successResult()
+		})
+		const harness = await buildHarness(childAgent)
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		context.budgetTracker = { total: 100_000, remaining: 100_000 }
+
+		// Go through the GATEWAY, which is where the clone was: calling
+		// `manager.sendMessage` directly always shared the tracker, so a test
+		// at that level proves nothing about the defect.
+		const gateway = new LocalTaskGateway(harness.manager, context)
+
+		for (let i = 0; i < 3; i++) {
+			const handle = await gateway.createTask({
+				agentId: 'child-budget',
+				prompt: 'work',
+				workingDirectory: '/tmp',
+			})
+			await waitForTask(harness.manager, handle.taskId)
+		}
+
+		const allocated = seen.map((c) => c.tokenBudget ?? 0)
+		expect(allocated).toHaveLength(3)
+		// Never over-committed: whatever is outstanding at any moment fits
+		// in the pool.
+		expect(allocated[0]).toBeLessThanOrEqual(100_000)
+
+		// These children settle before the next one spawns and spend
+		// nothing, so each gets the same allocation from a pool that was
+		// restored. This assertion used to require a STRICT decrease, which
+		// only held because the reservation was never returned — it was
+		// measuring the leak, not the rule.
+		expect(allocated[1]).toBe(allocated[0])
+		expect(allocated[2]).toBe(allocated[0])
+	})
+
+	it('returns what a settled child did not spend', async () => {
+		// The debit is a reservation so siblings cannot each be promised the
+		// same headroom. Nothing returned it, so a pool shrank by the full
+		// allocation on every spawn no matter what the child used: at a
+		// half-pool fraction, ten delegations left a parent with a
+		// thousandth of its budget and the next spawn was refused for a
+		// budget that had barely been spent.
+		const childAgent = makeAgent('child-thrifty', async () => ({
+			...successResult(),
+			usage: { ...EMPTY_TOKEN_USAGE, totalTokens: 1_000 },
+		}))
+		const harness = await buildHarness(childAgent)
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		context.budgetTracker = { total: 100_000, remaining: 100_000 }
+
+		const gateway = new LocalTaskGateway(harness.manager, context)
+		const handle = await gateway.createTask({
+			agentId: 'child-thrifty',
+			prompt: 'work',
+			workingDirectory: '/tmp',
+		})
+		await waitForTask(harness.manager, handle.taskId)
+
+		// Reserved 50_000, spent 1_000: the pool is down by what was used,
+		// not by what was set aside.
+		expect(context.budgetTracker.remaining).toBe(99_000)
+	})
+
+	it('keeps a concurrent sibling reservation until it settles', async () => {
+		// The reservation exists for exactly this: two children in flight
+		// must not each be promised the same headroom.
+		let release: (() => void) | undefined
+		const gate = new Promise<void>((resolve) => {
+			release = resolve
+		})
+		const childAgent = makeAgent('child-slow', async () => {
+			await gate
+			return successResult()
+		})
+		const harness = await buildHarness(childAgent)
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		context.budgetTracker = { total: 100_000, remaining: 100_000 }
+
+		const gateway = new LocalTaskGateway(harness.manager, context)
+		const first = await gateway.createTask({
+			agentId: 'child-slow',
+			prompt: 'work',
+			workingDirectory: '/tmp',
+		})
+
+		expect(context.budgetTracker.remaining).toBe(50_000)
+
+		release?.()
+		await waitForTask(harness.manager, first.taskId)
+		expect(context.budgetTracker.remaining).toBe(100_000)
+	})
+})
+
+/**
+ * A supervisor that fans out N tasks and watches one die had no way to say
+ * the other N-1 were now pointless. The primitive to stop them existed —
+ * every child holds an abort controller chained to the parent's — but
+ * nothing connected a failure to it.
+ */
+describe('LocalTaskGateway — what a failed child means for its siblings', () => {
+	async function fanOut(policy?: 'continue' | 'cancel-siblings') {
+		// One agent that fails, one that would run long enough to be worth
+		// cancelling.
+		let releaseSlow: (() => void) | undefined
+		const slowDone = new Promise<void>((resolve) => {
+			releaseSlow = resolve
+		})
+
+		const registry = new AgentRegistry()
+		registry.register(makeDefinition(makeAgent('fails', async () => failureResult('boom'))))
+		registry.register(
+			makeDefinition(
+				makeAgent('slow', async () => {
+					await slowDone
+					return successResult()
+				}),
+			),
+		)
+
+		const harness = await buildHarness(makeAgent('unused', async () => successResult()))
+		// Swap in the two-agent registry.
+		const manager = new AgentManager(registry, undefined, {
+			sessionStore: harness.store,
+			summaryMaterializer: harness.materializer,
+			workspaceRegistry: new WorkspaceBackendRegistry(),
+			capacity: new DefaultCapacityValidator(harness.store),
+			threadManager: harness.threadManager,
+		})
+
+		const context = buildContext(harness.parentSession.id, harness.projectId, harness.threadId)
+		const gateway = new LocalTaskGateway(
+			manager,
+			context,
+			undefined,
+			undefined,
+			policy ? { siblingFailurePolicy: policy } : undefined,
+		)
+
+		const slow = await gateway.createTask({
+			agentId: 'slow',
+			prompt: 'long work',
+			workingDirectory: '/tmp',
+		})
+		const failing = await gateway.createTask({
+			agentId: 'fails',
+			prompt: 'doomed work',
+			workingDirectory: '/tmp',
+		})
+
+		await manager.waitForCompletion(failing.taskId)
+		// Let the completion callback run.
+		await new Promise((r) => setTimeout(r, 0))
+
+		return { manager, slow, releaseSlow: releaseSlow as () => void }
+	}
+
+	it('leaves them alone by default — partial results are usually worth having', async () => {
+		const { manager, slow, releaseSlow } = await fanOut()
+		expect(manager.getInstance(slow.taskId)?.state).not.toBe('canceled')
+		releaseSlow()
+		await manager.waitForCompletion(slow.taskId)
+	})
+
+	it('cancels them when the fan-out only means something together', async () => {
+		const { manager, slow, releaseSlow } = await fanOut('cancel-siblings')
+		expect(manager.getInstance(slow.taskId)?.state).toBe('canceled')
+		releaseSlow()
+	})
+})

package/src/manager/agent/lifecycle.ts

@@ -20,6 +20,7 @@ import type {
 	SendMessageOptions,
 } from '../../types/agent/task.js'
 import { isTerminalAgentTaskState } from '../../types/agent/task.js'
+import { NamzuError } from '../../types/errors/index.js'
 import type { AgentId, RunId, SessionId, TaskId, TenantId } from '../../types/ids/index.js'
 import type { Message } from '../../types/message/index.js'
 import type { RunEvent, RunEventListener } from '../../types/run/events.js'
@@ -90,6 +91,8 @@ export class AgentManager {
 	private log: Logger
 	private config: Readonly<AgentManagerConfig>
 	private evictionTimers: Map<TaskId, ReturnType<typeof setTimeout>> = new Map()
+	/** One provisioning at a time per parent — see {@link provisionSpawn}. */
+	private spawnLocks: Map<SessionId, Promise<void>> = new Map()
 	private deps: AgentManagerDeps
 
 	constructor(
@@ -131,13 +134,38 @@ export class AgentManager {
 			options.budgetAllocation?.tokenBudget ?? maxAllocation,
 			maxAllocation,
 		)
-		context.budgetTracker.remaining -= allocatedTokens
+
+		// Budget exhaustion must not INVERT into no budget at all. Downstream,
+		// `tokenBudget: 0` means "uncapped" (`LimitChecker`: `tokenBudget > 0
+		// && total >= tokenBudget`), and `maxAllocation` floors to 0 as soon as
+		// the parent's remaining drops below `1 / maxBudgetFraction`. So the
+		// most depleted parent in the tree was the one that spawned an
+		// unlimited child. Refuse instead: a caller that wants an uncapped
+		// child can say so explicitly with its own `budgetAllocation`.
+		if (allocatedTokens <= 0) {
+			throw new NamzuError({
+				code: 'invalid_config',
+				message: `Cannot spawn "${options.agentId}": the parent has ${context.budgetTracker.remaining} tokens remaining, which allocates 0 to the child — and a token budget of 0 means UNLIMITED downstream.`,
+				details: {
+					agentId: options.agentId,
+					parentRemaining: context.budgetTracker.remaining,
+					maxBudgetFraction: this.config.maxBudgetFraction,
+				},
+			})
+		}
 
 		// Phase 6: SubSession + child Session + WorkspaceRef triple. Happens
 		// before taskId minting so a capacity failure short-circuits cleanly
 		// with no observable state change.
+		//
+		// The budget debit follows it for the same reason. It used to come
+		// first, so a spawn this call rejected still burned its allocation
+		// from a pool nobody credited back — the one state change the
+		// comment above promised there would not be.
 		const spawnRecord = await this.provisionSpawn(options, context)
 
+		context.budgetTracker.remaining -= allocatedTokens
+
 		const taskId = generateTaskId()
 
 		const childParentActor: ActorRef = {
@@ -168,6 +196,7 @@ export class AgentManager {
 			childAbortController,
 			context: childContext,
 			state: 'pending',
+			budgetReservation: allocatedTokens,
 			pendingMessages: [],
 			createdAt: Date.now(),
 			runEventListener: listener,
@@ -215,14 +244,14 @@ export class AgentManager {
 		let childConfig: BaseAgentConfig
 		if (definition.configBuilder) {
 			// Call the configBuilder regardless of whether factoryOptions were
-			// supplied. BYO-provider flows (Bedrock IAM, custom ProviderRegistry)
+			// supplied. BYO-provider flows (ambient cloud credentials, a custom registry)
 			// commonly omit factoryOptions because the provider resolves its own
 			// credentials; the builder still needs to run to wire provider+tools.
 			// Defaults: empty factoryOptions when omitted; configOverrides win.
 			childConfig = await definition.configBuilder({
 				...(context.factoryOptions ?? {}),
 				tokenBudget: allocatedTokens,
-				timeoutMs: options.budgetAllocation?.timeoutMs ?? context.budgetTracker.remaining,
+				timeoutMs: options.budgetAllocation?.timeoutMs ?? this.config.childTimeoutMs,
 				parentRunId: context.parentRunId as string | undefined,
 				depth: context.depth + 1,
 				...options.configOverrides,
@@ -239,6 +268,11 @@ export class AgentManager {
 			childConfig.threadId = context.threadId
 			childConfig.projectId = context.projectId
 			childConfig.tenantId = context.tenantId
+			// Stamp the trace parent the same way, rather than trusting every
+			// configBuilder to forward an option it may not know about.
+			if (options.configOverrides?.parentSpan) {
+				childConfig.parentSpan = options.configOverrides.parentSpan
+			}
 		} else {
 			this.log.warn('No configBuilder, using bare config', {
 				agentId: options.agentId,
@@ -246,8 +280,9 @@ export class AgentManager {
 			childConfig = {
 				model: options.configOverrides?.model ?? 'default',
 				tokenBudget: allocatedTokens,
-				timeoutMs: options.budgetAllocation?.timeoutMs ?? context.budgetTracker.remaining,
+				timeoutMs: options.budgetAllocation?.timeoutMs ?? this.config.childTimeoutMs,
 				temperature: options.configOverrides?.temperature,
+				parentSpan: options.configOverrides?.parentSpan,
 				maxIterations: options.configOverrides?.maxIterations,
 				maxResponseTokens: options.configOverrides?.maxResponseTokens,
 				env: options.configOverrides?.env,
@@ -367,15 +402,61 @@ export class AgentManager {
 		for (const taskId of this.instances.keys()) {
 			this.clearEvictionTimer(taskId)
 		}
-		this.cancelAll('' as RunId)
+		// Every live child, not the children of one parent. This used to call
+		// `cancelAll('' as RunId)`, and `cancelAll` filters by parent run —
+		// no task has an empty parent, so it matched nothing and the lines
+		// below then dropped every reference to work that was still running.
+		for (const taskId of [...this.instances.keys()]) {
+			this.cancel(taskId)
+		}
 		this.instances.clear()
 		this.spawnRecords.clear()
 		this.listeners.length = 0
 	}
 
+	/**
+	 * Serializes provisioning per parent session.
+	 *
+	 * The width cap counted existing children and then created one, with
+	 * every remaining provisioning step in between. Two concurrent spawns
+	 * under the same parent both read the same count, both saw room, and
+	 * both created — so a cap of N admitted N+1. The check and the write
+	 * that invalidates it have to be one critical section, and the parent
+	 * session is the narrowest key that makes them one: spawns under
+	 * different parents never contend.
+	 *
+	 * In-process only, which is the honest scope. Cross-process capacity is
+	 * the store's to enforce, and no store here spans processes.
+	 */
 	private async provisionSpawn(
 		options: SendMessageOptions,
 		context: AgentTaskContext,
+	): Promise<ChildSpawnRecord> {
+		const key = options.parentSessionId
+		const queued = (this.spawnLocks.get(key) ?? Promise.resolve()).then(
+			() => this.provisionSpawnUnlocked(options, context),
+			// A failed predecessor releases the lock rather than wedging every
+			// later spawn under the same parent.
+			() => this.provisionSpawnUnlocked(options, context),
+		)
+		const barrier = queued.then(
+			() => undefined,
+			() => undefined,
+		)
+		this.spawnLocks.set(key, barrier)
+		try {
+			return await queued
+		} finally {
+			// Clear the slot only if nobody queued behind us — otherwise we
+			// would drop the barrier the next waiter is chained to, and the
+			// map would leak one entry per parent if we never cleared at all.
+			if (this.spawnLocks.get(key) === barrier) this.spawnLocks.delete(key)
+		}
+	}
+
+	private async provisionSpawnUnlocked(
+		options: SendMessageOptions,
+		context: AgentTaskContext,
 	): Promise<ChildSpawnRecord> {
 		// Phase 9: deps are unconditional required. Every spawn produces a
 		// SubSession + Session + WorkspaceRef triple (Convention #0: no
@@ -467,9 +548,9 @@ export class AgentManager {
 
 		// Compensating rollback wraps every mutation after createSession so a
 		// mid-flight failure (status flip, subsession insert, workspace driver)
-		// leaves no orphan child session. Codex SPAWN-ROLLBACK critique (Phase
+		// leaves no orphan child session. spawn-rollback critique (Phase
 		// 2 review, 2026-04-18): without this, `workspaceRegistry.get().create`
-		// throwing — or a concurrent `updateSession` race — strands an
+		// throwing — or a concurrent `updateSession` race — leaves stranded an
 		// `active` child session with no subsession edge, invisible to the
 		// parent but counted against `maxDelegationWidth`.
 		let subSession: Awaited<ReturnType<typeof store.createSubSession>> | undefined
@@ -678,10 +759,44 @@ export class AgentManager {
 		this.markCompleted(agentTask.taskId, result)
 	}
 
+	/**
+	 * Return the unspent part of a settled child's reservation.
+	 *
+	 * The debit at spawn reserves headroom so siblings cannot each be
+	 * promised the same tokens. Nothing returned it, so a pool shrank by
+	 * the full allocation on every spawn no matter what the child used: at
+	 * a half-pool fraction, ten delegations left a parent with a
+	 * thousandth of its budget and the next spawn was refused for a budget
+	 * that had barely been spent.
+	 *
+	 * Idempotent — the reservation is cleared as it is returned, so a
+	 * second terminal transition for the same task cannot credit twice.
+	 * Spend above the reservation returns nothing rather than going
+	 * negative; the child was capped at its allocation, so that case means
+	 * the accounting was already wrong and inventing headroom would hide
+	 * it.
+	 */
+	private releaseBudget(agentTask: AgentTask, spent: number): void {
+		const reserved = agentTask.budgetReservation
+		if (reserved === undefined) return
+		agentTask.budgetReservation = undefined
+
+		const unused = Math.max(0, reserved - Math.max(0, spent))
+		if (unused === 0) return
+		agentTask.context.budgetTracker.remaining += unused
+		this.log.debug('Returned unspent child budget', {
+			taskId: agentTask.taskId,
+			reserved,
+			spent,
+			returned: unused,
+		})
+	}
+
 	private markCompleted(taskId: TaskId, result: BaseAgentResult): void {
 		const agentTask = this.instances.get(taskId)
 		if (!agentTask || isTerminalAgentTaskState(agentTask.state)) return
 
+		this.releaseBudget(agentTask, result.usage?.totalTokens ?? 0)
 		agentTask.result = result
 		agentTask.completedAt = Date.now()
 		this.updateState(taskId, 'completed')
@@ -701,6 +816,11 @@ export class AgentManager {
 		const agentTask = this.instances.get(taskId)
 		if (!agentTask || isTerminalAgentTaskState(agentTask.state)) return
 
+		// A child that failed still spent whatever it spent, and holding the
+		// rest of its reservation would punish the parent twice for one
+		// failure.
+		this.releaseBudget(agentTask, agentTask.result?.usage?.totalTokens ?? 0)
+
 		agentTask.result = {
 			runId: agentTask.context.parentRunId,
 			status: 'failed',