@namzu/sdk 6.0.0 → 6.1.0

package/src/runtime/query/__tests__/completion-notification.test.ts

@@ -0,0 +1,196 @@
+import { mkdtemp, rm } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { afterEach, describe, expect, it } from 'vitest'
+import { z } from 'zod'
+
+import { CompletionInbox } from '../../../gateway/completion-inbox.js'
+import { ToolRegistry } from '../../../registry/tool/execute.js'
+import { defineTool } from '../../../tools/defineTool.js'
+import type { TaskHandle } from '../../../types/agent/gateway.js'
+import type { SessionId, TaskId, TenantId } from '../../../types/ids/index.js'
+import { createUserMessage } from '../../../types/message/index.js'
+import type { LLMProvider, StreamChunk } from '../../../types/provider/index.js'
+import type { ProjectId, ThreadId } from '../../../types/session/ids.js'
+import { drainQuery } from '../index.js'
+
+/**
+ * The last link, and the one most likely to be built and never wired.
+ *
+ * `CompletionInbox` can queue an unclaimed completion and format it, and the
+ * coordinator tools can decline to claim one — but none of that reaches the
+ * model unless the iteration loop actually drains the inbox into the
+ * transcript. A mechanism that is declared, threaded through types, and
+ * driven by nothing is precisely the shape of the defect this whole change
+ * exists to fix, so the drain gets its own test at the loop level rather
+ * than being assumed from its parts.
+ */
+
+const ZERO_USAGE = {
+	promptTokens: 0,
+	completionTokens: 0,
+	totalTokens: 0,
+	cachedTokens: 0,
+	cacheWriteTokens: 0,
+}
+
+/** Calls a tool once, then answers. Two turns is all the drain needs. */
+class ToolThenAnswerProvider implements LLMProvider {
+	readonly id = 'tool-then-answer'
+	readonly name = 'Tool Then Answer Provider'
+	calls = 0
+
+	async *chatStream(): AsyncIterable<StreamChunk> {
+		this.calls += 1
+
+		if (this.calls === 1) {
+			yield {
+				id: 'msg_1',
+				delta: {
+					toolCalls: [
+						{
+							index: 0,
+							id: 'toolu_noop_1',
+							type: 'function',
+							function: { name: 'noop', arguments: '{}' },
+						},
+					],
+				},
+			}
+			yield {
+				id: 'msg_1',
+				delta: {},
+				finishReason: 'tool_calls',
+				usage: ZERO_USAGE,
+			}
+			return
+		}
+
+		yield { id: 'msg_2', delta: { content: 'Done.' } }
+		yield { id: 'msg_2', delta: {}, finishReason: 'stop', usage: ZERO_USAGE }
+	}
+}
+
+const noop = defineTool({
+	name: 'noop',
+	description: 'does nothing',
+	inputSchema: z.object({}),
+	category: 'analysis',
+	permissions: [],
+	readOnly: true,
+	destructive: false,
+	concurrencySafe: true,
+	async execute() {
+		return { success: true, output: 'ok' }
+	},
+})
+
+function completed(taskId: string, result: string): TaskHandle {
+	return {
+		taskId: taskId as TaskId,
+		agentId: 'reviewer',
+		state: 'completed',
+		createdAt: 1_000,
+		completedAt: 4_000,
+		result: { status: 'completed', result } as TaskHandle['result'],
+	}
+}
+
+const workdirs: string[] = []
+afterEach(async () => {
+	await Promise.all(workdirs.map((dir) => rm(dir, { recursive: true, force: true })))
+	workdirs.length = 0
+})
+
+async function runWith(inbox: CompletionInbox | undefined): Promise<string[]> {
+	const workingDirectory = await mkdtemp(join(tmpdir(), 'namzu-completion-'))
+	workdirs.push(workingDirectory)
+
+	const tools = new ToolRegistry()
+	tools.register(noop)
+
+	const run = await drainQuery({
+		provider: new ToolThenAnswerProvider(),
+		tools,
+		...(inbox ? { completionInbox: inbox } : {}),
+		agentId: 'agent_test',
+		agentName: 'Test Agent',
+		messages: [createUserMessage('delegate and report')],
+		workingDirectory,
+		runConfig: {
+			model: 'mock-model',
+			timeoutMs: 10_000,
+			tokenBudget: 100_000,
+			maxIterations: 4,
+			maxResponseTokens: 256,
+		},
+		sessionId: 'ses_completion' as SessionId,
+		threadId: 'thd_completion' as ThreadId,
+		projectId: 'prj_completion' as ProjectId,
+		tenantId: 'tnt_completion' as TenantId,
+	})
+
+	return run.messages
+		.filter((m) => m.role === 'user')
+		.map((m) => (typeof m.content === 'string' ? m.content : JSON.stringify(m.content)))
+}
+
+describe('an unclaimed completion reaches the transcript', () => {
+	it('injects the notification as a user message the next turn can read', async () => {
+		const inbox = new CompletionInbox()
+		// Settled by a gateway before the turn ended, with nothing waiting on
+		// it — the abandoned-launch case.
+		inbox.attach({
+			onTaskCompleted: (cb: (h: TaskHandle) => void) => {
+				cb(completed('tsk_late', 'the worker finished after the wait was abandoned'))
+				return () => {}
+			},
+		} as never)
+
+		const userMessages = await runWith(inbox)
+		const notification = userMessages.find((m) => m.includes('task-notification'))
+
+		expect(notification).toBeDefined()
+		expect(notification).toContain('tsk_late')
+		expect(notification).toContain('the worker finished after the wait was abandoned')
+	})
+
+	it('drains it exactly once, however many turns follow', async () => {
+		const inbox = new CompletionInbox()
+		inbox.attach({
+			onTaskCompleted: (cb: (h: TaskHandle) => void) => {
+				cb(completed('tsk_late', 'only once'))
+				return () => {}
+			},
+		} as never)
+
+		const userMessages = await runWith(inbox)
+
+		expect(userMessages.filter((m) => m.includes('task-notification'))).toHaveLength(1)
+	})
+
+	it('says nothing when every completion was already delivered', async () => {
+		// The `dc16d58` regression at the loop level: a blocking `create_task`
+		// claims its own completion, so the transcript must stay clean.
+		const inbox = new CompletionInbox()
+		inbox.attach({
+			onTaskCompleted: (cb: (h: TaskHandle) => void) => {
+				cb(completed('tsk_awaited', 'delivered as a tool_result'))
+				return () => {}
+			},
+		} as never)
+		inbox.claim('tsk_awaited' as TaskId)
+
+		const userMessages = await runWith(inbox)
+
+		expect(userMessages.some((m) => m.includes('task-notification'))).toBe(false)
+	})
+
+	it('runs unchanged with no inbox at all', async () => {
+		// The kernel must not require one: a host on the old wiring keeps
+		// working, it just never hears about abandoned completions.
+		const userMessages = await runWith(undefined)
+
+		expect(userMessages.some((m) => m.includes('task-notification'))).toBe(false)
+	})
+})

package/src/runtime/query/__tests__/step-allow-list.test.ts

@@ -0,0 +1,205 @@
+import { mkdtemp, rm } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { afterEach, describe, expect, it } from 'vitest'
+import { z } from 'zod'
+
+import { ToolRegistry } from '../../../registry/tool/execute.js'
+import { defineTool } from '../../../tools/defineTool.js'
+import type { SessionId, TenantId } from '../../../types/ids/index.js'
+import { createUserMessage } from '../../../types/message/index.js'
+import type { LLMProvider, StreamChunk } from '../../../types/provider/index.js'
+import type { ProjectId, ThreadId } from '../../../types/session/ids.js'
+import { drainQuery } from '../index.js'
+
+/**
+ * Narrowing a step has to narrow what RUNS, not only what is offered.
+ *
+ * `prepareStep.activeTools` says "restrict which tools the model may call
+ * this step, by name", and for a while it did nothing of the sort: it chose
+ * which schemas went into the request, was copied into the tool context, and
+ * was read by nobody on the execution path. A model that named a withheld
+ * tool had it run — and a model does name one, from a repeated call earlier
+ * in the context, from a gateway carrying its own tool list, or from a
+ * replayed cache prefix.
+ *
+ * These tests call the withheld tool deliberately, which is the only way to
+ * tell an enforced restriction from a presentational one. Every existing test
+ * asked the model nicely and so could not have caught this.
+ */
+
+const ZERO_USAGE = {
+	promptTokens: 0,
+	completionTokens: 0,
+	totalTokens: 0,
+	cachedTokens: 0,
+	cacheWriteTokens: 0,
+}
+
+let danger = 0
+
+const readOnly = defineTool({
+	name: 'read_only',
+	description: 'observes',
+	inputSchema: z.object({}),
+	category: 'analysis',
+	permissions: [],
+	readOnly: true,
+	destructive: false,
+	concurrencySafe: true,
+	async execute() {
+		return { success: true, output: 'observed' }
+	},
+})
+
+const dangerous = defineTool({
+	name: 'danger',
+	description: 'mutates',
+	inputSchema: z.object({}),
+	category: 'custom',
+	permissions: [],
+	readOnly: false,
+	destructive: true,
+	concurrencySafe: false,
+	async execute() {
+		danger += 1
+		return { success: true, output: 'the withheld tool ran' }
+	},
+})
+
+/** Calls whichever tool it is told to, then answers. */
+class NamesTool implements LLMProvider {
+	readonly id = 'names-tool'
+	readonly name = 'Names Tool Provider'
+	calls = 0
+	offered: string[][] = []
+
+	constructor(private readonly toolName: string) {}
+
+	async *chatStream(params: {
+		tools?: { function: { name: string } }[]
+	}): AsyncIterable<StreamChunk> {
+		this.calls += 1
+		this.offered.push((params.tools ?? []).map((t) => t.function.name))
+
+		if (this.calls === 1) {
+			yield {
+				id: 'msg_1',
+				delta: {
+					toolCalls: [
+						{
+							index: 0,
+							id: 'toolu_1',
+							type: 'function',
+							function: { name: this.toolName, arguments: '{}' },
+						},
+					],
+				},
+			}
+			yield { id: 'msg_1', delta: {}, finishReason: 'tool_calls', usage: ZERO_USAGE }
+			return
+		}
+
+		yield { id: 'msg_2', delta: { content: 'done' } }
+		yield { id: 'msg_2', delta: {}, finishReason: 'stop', usage: ZERO_USAGE }
+	}
+}
+
+const workdirs: string[] = []
+afterEach(async () => {
+	await Promise.all(workdirs.map((dir) => rm(dir, { recursive: true, force: true })))
+	workdirs.length = 0
+	danger = 0
+})
+
+async function run(opts: {
+	names: string
+	allowedTools?: string[]
+	activeTools?: string[]
+}): Promise<{ offered: string[][]; output: string }> {
+	const workingDirectory = await mkdtemp(join(tmpdir(), 'namzu-allow-'))
+	workdirs.push(workingDirectory)
+
+	const tools = new ToolRegistry()
+	tools.register(readOnly)
+	tools.register(dangerous)
+
+	const provider = new NamesTool(opts.names)
+	let output = ''
+
+	await drainQuery(
+		{
+			provider,
+			tools,
+			...(opts.allowedTools ? { allowedTools: opts.allowedTools } : {}),
+			...(opts.activeTools
+				? { prepareStep: () => ({ activeTools: opts.activeTools as string[] }) }
+				: {}),
+			agentId: 'agent_test',
+			agentName: 'Test Agent',
+			messages: [createUserMessage('go')],
+			workingDirectory,
+			runConfig: {
+				model: 'mock-model',
+				timeoutMs: 10_000,
+				tokenBudget: 100_000,
+				maxIterations: 3,
+				maxResponseTokens: 256,
+			},
+			sessionId: 'ses_allow' as SessionId,
+			threadId: 'thd_allow' as ThreadId,
+			projectId: 'prj_allow' as ProjectId,
+			tenantId: 'tnt_allow' as TenantId,
+		},
+		(event) => {
+			if (event.type === 'tool_completed') output += event.result ?? ''
+		},
+	)
+
+	return { offered: provider.offered, output }
+}
+
+describe('a narrowed step narrows what can run, not just what is shown', () => {
+	it('refuses a tool the step withheld, even when the model names it', async () => {
+		const { offered, output } = await run({ names: 'danger', activeTools: ['read_only'] })
+
+		// The request really was narrowed — otherwise this test proves nothing
+		// about enforcement, only about the model's manners.
+		expect(offered[0]).toEqual(['read_only'])
+		// And the call was refused rather than executed.
+		expect(danger, 'the withheld tool executed').toBe(0)
+		expect(output).toContain('not available on this step')
+	})
+
+	it('refuses a tool outside the run-level list too', async () => {
+		const { output } = await run({ names: 'danger', allowedTools: ['read_only'] })
+
+		expect(danger).toBe(0)
+		expect(output).toContain('not available on this step')
+	})
+
+	it('still runs a tool that is on the list', async () => {
+		const { output } = await run({ names: 'read_only', activeTools: ['read_only'] })
+
+		expect(output).toContain('observed')
+	})
+
+	it('leaves an unnarrowed run alone', async () => {
+		// Absent means unrestricted. A run that never narrows anything must not
+		// suddenly start refusing its own tools.
+		const { output } = await run({ names: 'danger' })
+
+		expect(danger).toBe(1)
+		expect(output).toContain('the withheld tool ran')
+	})
+
+	it('treats an empty list as "nothing", not as "no restriction"', async () => {
+		// The fail-open reading this codebase has already been bitten by once,
+		// in the delegate roster: an empty allow-list IS the answer, and
+		// degrading it to "unrestricted" is how a closed list becomes open.
+		const { output } = await run({ names: 'danger', activeTools: [] })
+
+		expect(danger).toBe(0)
+		expect(output).toContain('(none)')
+	})
+})

package/src/runtime/query/executor.ts

@@ -205,6 +205,8 @@ export class ToolExecutor {
 	private workingStateManager?: WorkingStateManager
 	private probes: ProbeRegistry
 	private parentSpan?: Span
+	/** Set per turn by the orchestrator; see {@link setStepAllowedTools}. */
+	private stepAllowedTools?: readonly string[]
 	private readonly readPaths: Set<string> = new Set()
 	private readonly readFingerprints: Map<string, string> = new Map()
 	private readonly fileReadTracker: FileReadTracker = {
@@ -251,6 +253,21 @@ export class ToolExecutor {
 		this.parentSpan = span
 	}
 
+	/**
+	 * Narrow what this turn may call, or clear the narrowing.
+	 *
+	 * Re-set each turn by the orchestrator for the same reason the parent span
+	 * is: `prepareStep` can hand a different list to every step, and the run's
+	 * own `allowedTools` is only the default when a step names none.
+	 *
+	 * Without this the executor could only ever see the RUN-level list, so a
+	 * per-step narrowing reached the request that was sent and nothing else —
+	 * the model was shown fewer tools and could still call all of them.
+	 */
+	setStepAllowedTools(names: readonly string[] | undefined): void {
+		this.stepAllowedTools = names
+	}
+
 	/**
 	 * Answer every `tool_use` block in `response` with exactly one
 	 * `tool_result`.
@@ -424,7 +441,10 @@ export class ToolExecutor {
 			},
 			invocationState: this.config.invocationState,
 			toolRegistry: this.config.tools,
-			allowedTools: this.config.allowedTools,
+			// The step's list wins where it has one; the run's is the default.
+			// Same precedence the request already uses when it decides which
+			// schemas to send, so the menu and the kitchen agree.
+			allowedTools: this.stepAllowedTools ?? this.config.allowedTools,
 			sandbox: this.config.sandbox,
 			fileReadTracker: this.fileReadTracker,
 			...(this.parentSpan ? { parentSpan: this.parentSpan } : {}),

package/src/runtime/query/index.ts

@@ -410,6 +410,16 @@ export interface QueryParams {
 
 	taskGateway?: import('../../types/agent/gateway.js').TaskGateway
 
+	/**
+	 * Where a worker completion goes when no tool call is waiting for it.
+	 *
+	 * Supplied by whoever built the coordinator tools, because the tools and
+	 * this loop have to share one inbox: the tools claim what they deliver,
+	 * and the loop delivers what is left. Omitted, the loop drains nothing and
+	 * the behaviour is exactly what it was before the inbox existed.
+	 */
+	completionInbox?: import('../../gateway/completion-inbox.js').CompletionInbox
+
 	launchedTasks?: Map<
 		import('../../types/ids/index.js').TaskId,
 		import('./iteration/phases/context.js').LaunchedTaskMeta
@@ -842,6 +852,7 @@ export async function* query(params: QueryParams): AsyncGenerator<RunEvent, Run>
 		checkpointMgr,
 		planManager: ctx.planManager,
 		taskGateway: params.taskGateway,
+		completionInbox: params.completionInbox,
 		taskStore: params.taskStore,
 		launchedTasks: params.launchedTasks ?? new Map(),
 		// Run-scoped. An approval is a statement about this run's work;

package/src/runtime/query/iteration/index.ts

@@ -5,6 +5,7 @@ import {
 	DEFAULT_STRUCTURED_OUTPUT_RETRIES,
 	STRUCTURED_OUTPUT_REPROMPT,
 } from '../../../constants/tools/index.js'
+import { formatCompletionNotification } from '../../../gateway/completion-inbox.js'
 import { renderSkillsSection } from '../../../persona/assembler.js'
 import { collect } from '../../../provider/collect.js'
 import {
@@ -61,6 +62,16 @@ export type { ToolReviewOutcome } from './phases/index.js'
  */
 const DEFAULT_ANSWER_REVIEW_LIMIT = 3
 
+/**
+ * How long a finishing run waits for a background worker it launched.
+ *
+ * Long enough to be worth having — a delegated worker doing real work takes
+ * minutes — and bounded because a worker that never finishes must not hold
+ * the run open forever. `maxIterations` bounds how many times this can
+ * happen, and the run's own timeout bounds the whole thing regardless.
+ */
+const BACKGROUND_TASK_GRACE_MS = 120_000
+
 export class IterationOrchestrator {
 	private ctx: IterationContext
 	/** Rejections so far. See {@link DEFAULT_ANSWER_REVIEW_LIMIT}. */
@@ -207,7 +218,13 @@ export class IterationOrchestrator {
 				// supplied no hook.
 				const step = await this.prepareStep(iterationNum)
 
-				const llmTools = this.ctx.tools.toLLMTools(step.allowedTools ?? this.ctx.allowedTools)
+				const stepAllowedTools = step.allowedTools ?? this.ctx.allowedTools
+				const llmTools = this.ctx.tools.toLLMTools(stepAllowedTools)
+				// The same list the request was built from now also bounds what
+				// may run. Narrowing only the request left the restriction
+				// presentational — the model was shown fewer tools and could
+				// still call any of them by name.
+				this.ctx.toolExecutor.setStepAllowedTools(stepAllowedTools)
 				const enforceToolInputSchema = enforcedModelInputToolNames(this.ctx.tools, llmTools)
 				const stepModel = step.model ?? model
 
@@ -551,6 +568,36 @@ export class IterationOrchestrator {
 						}
 					}
 
+					// A background worker is still out there, and this turn was
+					// about to end the run.
+					//
+					// Settling here would throw away the very thing the launch
+					// existed to produce: the supervisor said "launched", the
+					// worker had not finished, and the run closed over it. So
+					// the run is held open — bounded by the deadline below and
+					// by `maxIterations` above, so a worker that never finishes
+					// cannot keep it open forever — and the completion arrives
+					// as a notification the next turn reads.
+					if (!forceFinalize && this.ctx.completionInbox?.hasPendingWork) {
+						this.ctx.log.info('Holding the run open for a background task', {
+							runId: runMgr.id,
+							iteration: iterationNum,
+						})
+						await this.ctx.completionInbox.waitForArrival(BACKGROUND_TASK_GRACE_MS)
+						const arrived = this.ctx.completionInbox.drain()
+						if (arrived.length > 0) {
+							runMgr.pushMessage(createUserMessage(formatCompletionNotification(arrived)))
+							await this.ctx.emitEvent({
+								type: 'iteration_completed',
+								runId: runMgr.id,
+								iteration: iterationNum,
+								hasToolCalls: false,
+							})
+							yield* this.ctx.drainPending()
+							continue
+						}
+					}
+
 					if (!hasContent && !forceFinalize) {
 						this.ctx.log.warn('Empty completion detected — requesting final summary', {
 							iteration: iterationNum,
@@ -674,6 +721,35 @@ export class IterationOrchestrator {
 					return
 				}
 
+				// Workers that finished with nobody listening.
+				//
+				// A completion normally reaches the supervisor as the
+				// `tool_result` of the `create_task` that launched it. Two
+				// cases have no such call: a launch made in the background on
+				// purpose, and a blocking launch whose deadline passed — the
+				// model was told "timed out, it may still be running" and the
+				// worker then finished, holding a result nothing would read.
+				//
+				// This is the channel that was removed in `dc16d58` because it
+				// double-delivered: it fired for completions the blocking tool
+				// had already handed over, so the supervisor saw each result
+				// twice. The inbox restores it with the distinction that was
+				// missing — a tool that delivers a completion claims it, and
+				// only unclaimed ones arrive here.
+				//
+				// Placed beside the advisory phase deliberately: that is the
+				// established seam for putting a user message in after tool
+				// results and before the next turn.
+				const unheard = this.ctx.completionInbox?.drain() ?? []
+				if (unheard.length > 0) {
+					this.ctx.log.info('Delivering unawaited task completions', {
+						runId: runMgr.id,
+						iteration: iterationNum,
+						tasks: unheard.map((h) => h.taskId),
+					})
+					runMgr.pushMessage(createUserMessage(formatCompletionNotification(unheard)))
+				}
+
 				await runAdvisoryPhase(this.ctx, iterationNum, response)
 
 				if (this.ctx.pluginManager) {

package/src/runtime/query/iteration/phases/context.ts

@@ -3,6 +3,7 @@ import type { AgentBus } from '../../../../bus/index.js'
 import type { WorkingStateManager } from '../../../../compaction/manager.js'
 import type { ContextReducer } from '../../../../compaction/reducer.js'
 import type { CompactionConfig } from '../../../../config/runtime.js'
+import type { CompletionInbox } from '../../../../gateway/completion-inbox.js'
 import type { PlanManager } from '../../../../manager/plan/lifecycle.js'
 import type { RunPersistence } from '../../../../manager/run/persistence.js'
 import type { ActivityStore } from '../../../../store/activity/memory.js'
@@ -103,6 +104,15 @@ export interface IterationContext {
 
 	readonly taskGateway?: TaskGateway
 
+	/**
+	 * Completions no call is waiting for, on their way to the transcript.
+	 *
+	 * Absent means the loop behaves exactly as it did before this existed:
+	 * a blocking `create_task` still delivers its own result, and a
+	 * completion nobody awaited is simply never mentioned.
+	 */
+	readonly completionInbox?: CompletionInbox
+
 	readonly taskStore?: TaskStore
 
 	readonly launchedTasks: Map<TaskId, LaunchedTaskMeta>