@namzu/sdk 0.4.1 → 0.4.3

package/src/registry/tool/execute.test.ts

@@ -0,0 +1,330 @@
+/**
+ * Current-code invariants asserted (2026-04-21, ses_006 Phase 3):
+ *
+ *   - `ToolRegistry` extends `ManagedRegistry<ToolDefinition>` with
+ *     `idField: 'name'` and an additional `availability` map keyed by
+ *     tool name.
+ *   - `register(tool)` defaults availability to 'active'; the array
+ *     overload uses the second arg (if it's a string) as the initial
+ *     state for all tools.
+ *   - Tier validation: if `tool.tier` is set AND `tierConfig` is
+ *     configured, the tier id must be in `tierConfig.tiers`; otherwise
+ *     register throws. Tools without a tier (even with tierConfig
+ *     configured) are always accepted.
+ *   - `unregister` + `clear` drop both the Registry entry and the
+ *     availability entry.
+ *   - `activate` / `defer` take a list of names; each throws (via
+ *     `getOrThrow`) if the name is unknown, then toggles availability.
+ *   - `suspendAll` flips every `active` tool to `suspended`;
+ *     `deferred` / `suspended` tools are left alone. `hasSuspended`
+ *     reports true iff at least one tool is suspended.
+ *   - `getAvailability(name)` returns 'active' as a default even for
+ *     unknown names (this is non-obvious but is the current behavior).
+ *   - `searchDeferred(q)` is a case-insensitive filter against name OR
+ *     description of every DEFERRED tool.
+ *   - `assignTiers(mapping)` mutates `tool.tier` on existing tools;
+ *     throws via `getOrThrow` on unknown name; throws if the tier id
+ *     is not in `tierConfig.tiers`.
+ *   - `toTierGuidance` returns null without a guidanceTemplate; calls
+ *     the template with every defined tier otherwise.
+ *   - `toPromptSection`: returns '' when no active or deferred tools;
+ *     otherwise produces `<available_tools>` + `<deferred_tools>`
+ *     fragments.
+ *   - `toLLMTools`: converts active + suspended tools (filtered by
+ *     `toolNames` if provided) into LLM tool schemas via
+ *     `zodToJsonSchema`. When `tierConfig.labelInDescription` is true,
+ *     the description is prefixed with the tier label.
+ *   - `execute`: validates availability → permissionMode plan gate →
+ *     zod.safeParse → calls tool.execute inside an OTEL span; wraps
+ *     thrown errors as `{success:false, error}`.
+ */
+
+import { describe, expect, it, vi } from 'vitest'
+import { z } from 'zod'
+
+import type { RunId } from '../../types/ids/index.js'
+import type { ToolContext, ToolDefinition, ToolTierConfig } from '../../types/tool/index.js'
+
+import { ToolRegistry } from './execute.js'
+
+function makeTool(name: string, overrides: Partial<ToolDefinition> = {}): ToolDefinition {
+	return {
+		name,
+		description: `${name} tool`,
+		inputSchema: z.object({ k: z.string().optional() }),
+		async execute() {
+			return { success: true, output: `${name}-ran` }
+		},
+		...overrides,
+	}
+}
+
+function makeContext(overrides: Partial<ToolContext> = {}): ToolContext {
+	return {
+		runId: 'run_1' as RunId,
+		workingDirectory: '/tmp',
+		abortSignal: new AbortController().signal,
+		env: {},
+		log: () => {},
+		...overrides,
+	}
+}
+
+describe('ToolRegistry — register + availability', () => {
+	it('single-tool register: default availability is active', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('read'))
+		expect(r.getAvailability('read')).toBe('active')
+	})
+
+	it('array-register: second arg sets availability for all', () => {
+		const r = new ToolRegistry()
+		r.register([makeTool('a'), makeTool('b')], 'deferred')
+		expect(r.getAvailability('a')).toBe('deferred')
+		expect(r.getAvailability('b')).toBe('deferred')
+	})
+
+	it('getAvailability returns active for unknown names (current default)', () => {
+		const r = new ToolRegistry()
+		expect(r.getAvailability('never-registered')).toBe('active')
+	})
+
+	it('unregister drops both the tool + the availability entry', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		expect(r.unregister('a')).toBe(true)
+		expect(r.get('a')).toBeUndefined()
+		expect(r.listIds()).toEqual([])
+	})
+
+	it('clear empties both maps', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		r.clear()
+		expect(r.listIds()).toEqual([])
+	})
+})
+
+describe('ToolRegistry — tier validation', () => {
+	const tierConfig: ToolTierConfig = {
+		tiers: [
+			{ id: 'safe', label: 'Safe', priority: 1 },
+			{ id: 'danger', label: 'Danger', priority: 10 },
+		],
+	}
+
+	it('accepts tools without a tier even when tierConfig is configured', () => {
+		const r = new ToolRegistry({ tierConfig })
+		expect(() => r.register(makeTool('a'))).not.toThrow()
+	})
+
+	it('accepts tools whose tier is in the config', () => {
+		const r = new ToolRegistry({ tierConfig })
+		expect(() => r.register(makeTool('a', { tier: 'safe' }))).not.toThrow()
+	})
+
+	it('throws when tier id is unknown', () => {
+		const r = new ToolRegistry({ tierConfig })
+		expect(() => r.register(makeTool('a', { tier: 'bogus' }))).toThrow(/not defined/)
+	})
+})
+
+describe('ToolRegistry — activate / defer / suspend', () => {
+	it('activate + defer throw via getOrThrow on unknown name', () => {
+		const r = new ToolRegistry()
+		expect(() => r.activate(['unknown'])).toThrow(/Not found/)
+		expect(() => r.defer(['unknown'])).toThrow(/Not found/)
+	})
+
+	it('activate sets state to active', () => {
+		const r = new ToolRegistry()
+		r.register([makeTool('a')], 'deferred')
+		r.activate(['a'])
+		expect(r.getAvailability('a')).toBe('active')
+	})
+
+	it('defer sets state to deferred', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		r.defer(['a'])
+		expect(r.getAvailability('a')).toBe('deferred')
+	})
+
+	it('suspendAll flips active → suspended; leaves deferred / suspended alone', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		r.register(makeTool('b'))
+		r.register([makeTool('c')], 'deferred')
+		r.suspendAll()
+		expect(r.getAvailability('a')).toBe('suspended')
+		expect(r.getAvailability('b')).toBe('suspended')
+		expect(r.getAvailability('c')).toBe('deferred')
+		expect(r.hasSuspended()).toBe(true)
+	})
+
+	it('hasSuspended returns false when none are suspended', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		expect(r.hasSuspended()).toBe(false)
+	})
+})
+
+describe('ToolRegistry — searchDeferred', () => {
+	it('filters only deferred tools by name OR description (case-insensitive)', () => {
+		const r = new ToolRegistry()
+		r.register([makeTool('alpha', { description: 'does alpha' })], 'deferred')
+		r.register([makeTool('beta', { description: 'does BETA' })], 'deferred')
+		r.register(makeTool('gamma'))
+		expect(r.searchDeferred('alpha').map((t) => t.name)).toEqual(['alpha'])
+		expect(r.searchDeferred('beta').map((t) => t.name)).toEqual(['beta'])
+		expect(r.searchDeferred('does').map((t) => t.name)).toEqual(['alpha', 'beta'])
+		expect(r.searchDeferred('gamma')).toEqual([])
+	})
+})
+
+describe('ToolRegistry — tier mutation + guidance', () => {
+	const tierConfig: ToolTierConfig = {
+		tiers: [
+			{ id: 'safe', label: 'Safe', priority: 1 },
+			{ id: 'danger', label: 'Danger', priority: 10 },
+		],
+		guidanceTemplate: (tiers) => `tiers: ${tiers.map((t) => t.id).join(',')}`,
+		labelInDescription: true,
+	}
+
+	it('assignTiers sets the tier on each mapped tool', () => {
+		const r = new ToolRegistry({ tierConfig })
+		r.register(makeTool('a'))
+		r.assignTiers({ a: 'safe' })
+		expect(r.get('a')?.tier).toBe('safe')
+	})
+
+	it('assignTiers throws when the tier is unknown', () => {
+		const r = new ToolRegistry({ tierConfig })
+		r.register(makeTool('a'))
+		expect(() => r.assignTiers({ a: 'bogus' })).toThrow(/not defined/)
+	})
+
+	it('toTierGuidance returns null without a template', () => {
+		const r = new ToolRegistry()
+		expect(r.toTierGuidance()).toBeNull()
+	})
+
+	it('toTierGuidance renders via the template', () => {
+		const r = new ToolRegistry({ tierConfig })
+		expect(r.toTierGuidance()).toBe('tiers: safe,danger')
+	})
+})
+
+describe('ToolRegistry — toPromptSection + toLLMTools', () => {
+	it('toPromptSection returns empty when the registry is empty', () => {
+		const r = new ToolRegistry()
+		expect(r.toPromptSection()).toBe('')
+	})
+
+	it('toPromptSection renders available + deferred fragments', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		r.register([makeTool('b')], 'deferred')
+		const s = r.toPromptSection()
+		expect(s).toContain('<available_tools>')
+		expect(s).toContain('- a: a tool')
+		expect(s).toContain('<deferred_tools>')
+		expect(s).toContain('- b')
+	})
+
+	it('toLLMTools: converts active + suspended tools', () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('a'))
+		r.register([makeTool('b')], 'deferred')
+		r.register([makeTool('c')], 'suspended')
+
+		const schemas = r.toLLMTools()
+		const names = schemas.map((t) => t.function.name).sort()
+		expect(names).toEqual(['a', 'c'])
+	})
+
+	it('toLLMTools: prefixes description with tier label when labelInDescription is true', () => {
+		const tierConfig: ToolTierConfig = {
+			tiers: [{ id: 'safe', label: 'Safe', priority: 1 }],
+			labelInDescription: true,
+		}
+		const r = new ToolRegistry({ tierConfig })
+		r.register(makeTool('a', { tier: 'safe' }))
+		const schemas = r.toLLMTools()
+		expect(schemas[0]?.function.description).toBe('[Safe] a tool')
+	})
+})
+
+describe('ToolRegistry — execute', () => {
+	it('returns error when tool is not active (e.g. deferred)', async () => {
+		const r = new ToolRegistry()
+		r.register([makeTool('a')], 'deferred')
+		const result = await r.execute('a', {}, makeContext())
+		expect(result.success).toBe(false)
+		expect(result.error).toMatch(/deferred and cannot be executed/)
+	})
+
+	it('blocks non-read-only tools in plan mode', async () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('write', { isReadOnly: () => false }))
+		const result = await r.execute(
+			'write',
+			{},
+			makeContext({
+				permissionContext: { mode: 'plan', runId: 'run_1', workingDirectory: '/tmp' },
+			}),
+		)
+		expect(result.success).toBe(false)
+		expect(result.permissionDenied).toBe(true)
+	})
+
+	it('allows read-only tools in plan mode', async () => {
+		const r = new ToolRegistry()
+		const execute = vi.fn(async () => ({ success: true, output: 'ok' }))
+		r.register(makeTool('read', { isReadOnly: () => true, execute }))
+		const result = await r.execute(
+			'read',
+			{},
+			makeContext({
+				permissionContext: { mode: 'plan', runId: 'run_1', workingDirectory: '/tmp' },
+			}),
+		)
+		expect(result.success).toBe(true)
+		expect(execute).toHaveBeenCalled()
+	})
+
+	it('returns error when input fails zod validation', async () => {
+		const r = new ToolRegistry()
+		r.register(
+			makeTool('strict', {
+				inputSchema: z.object({ required: z.string() }),
+			}),
+		)
+		const result = await r.execute('strict', { required: 123 }, makeContext())
+		expect(result.success).toBe(false)
+		expect(result.error).toMatch(/Invalid input/)
+	})
+
+	it('wraps thrown errors in the execute function', async () => {
+		const r = new ToolRegistry()
+		r.register(
+			makeTool('bad', {
+				async execute() {
+					throw new Error('boom')
+				},
+			}),
+		)
+		const result = await r.execute('bad', {}, makeContext())
+		expect(result.success).toBe(false)
+		expect(result.error).toMatch(/execution failed: boom/)
+	})
+
+	it('returns the tool result on happy path', async () => {
+		const r = new ToolRegistry()
+		r.register(makeTool('good'))
+		const result = await r.execute('good', {}, makeContext())
+		expect(result.success).toBe(true)
+		expect(result.output).toBe('good-ran')
+	})
+})

package/src/run/LimitChecker.ts

@@ -13,15 +13,6 @@ export type LimitCheckResult =
 	| { type: 'warning'; reason: StopReason }
 	| { type: 'hard_stop'; reason: StopReason }
 
-export function checkLimits(
-	config: LimitCheckerConfig,
-	state: LimitCheckerState,
-): StopReason | null {
-	const result = checkLimitsDetailed(config, state)
-	if (result.type === 'ok') return null
-	return result.reason
-}
-
 export function checkLimitsDetailed(
 	config: LimitCheckerConfig,
 	state: LimitCheckerState,

package/src/run/reporter.ts

@@ -1,4 +1,4 @@
-import type { AgentRun, RunEvent, RunEventListener } from '../types/run/index.js'
+import type { Run, RunEvent, RunEventListener } from '../types/run/index.js'
 import { formatCost } from '../utils/cost.js'
 import { type Logger, getRootLogger } from '../utils/logger.js'
 
@@ -13,7 +13,7 @@ function formatDuration(ms: number): string {
 
 export interface RunReporter {
 	listener: RunEventListener
-	summary(run: AgentRun): void
+	summary(run: Run): void
 }
 
 export function createRunReporter(parentLogger?: Logger): RunReporter {
@@ -229,7 +229,7 @@ export function createRunReporter(parentLogger?: Logger): RunReporter {
 		}
 	}
 
-	function summary(run: AgentRun): void {
+	function summary(run: Run): void {
 		const elapsed = (run.endedAt ?? Date.now()) - run.startedAt
 		const { tokenUsage, costInfo, currentIteration, stopReason } = run
 

package/src/runtime/query/__tests__/checkpoint.test.ts

@@ -0,0 +1,169 @@
+import { describe, expect, it } from 'vitest'
+import type { RunDiskStore } from '../../../store/run/disk.js'
+import type { CheckpointId, IterationCheckpoint } from '../../../types/hitl/index.js'
+import type { EmergencySaveId, RunId } from '../../../types/ids/index.js'
+import type { EmergencySaveData } from '../../../types/run/emergency.js'
+import { CheckpointManager, projectEmergencyToCheckpoint } from '../checkpoint.js'
+
+function makeCheckpoint(overrides: Partial<IterationCheckpoint> = {}): IterationCheckpoint {
+	return {
+		id: 'cp_test_a' as CheckpointId,
+		runId: 'run_test' as RunId,
+		iteration: 1,
+		messages: [{ role: 'user', content: 'hello' }],
+		tokenUsage: {
+			promptTokens: 0,
+			completionTokens: 0,
+			totalTokens: 0,
+			cachedTokens: 0,
+			cacheWriteTokens: 0,
+		},
+		costInfo: {
+			inputCostPer1M: 0,
+			outputCostPer1M: 0,
+			totalCost: 0,
+			cacheDiscount: 0,
+		},
+		guardState: { iterationCount: 1, elapsedMs: 100 },
+		createdAt: Date.now(),
+		...overrides,
+	}
+}
+
+function makeStoreStub(checkpoints: IterationCheckpoint[]): RunDiskStore {
+	return {
+		listCheckpoints: async () => checkpoints,
+	} as unknown as RunDiskStore
+}
+
+describe('CheckpointManager.listEntries', () => {
+	it('projects stored checkpoints to CheckpointListEntry', async () => {
+		const store = makeStoreStub([
+			makeCheckpoint({
+				id: 'cp_a' as CheckpointId,
+				iteration: 1,
+				createdAt: 1000,
+				messages: [
+					{ role: 'user', content: 'hi' },
+					{ role: 'assistant', content: 'hello' },
+				],
+			}),
+			makeCheckpoint({
+				id: 'cp_b' as CheckpointId,
+				iteration: 2,
+				createdAt: 2000,
+				messages: [
+					{ role: 'user', content: 'hi' },
+					{ role: 'assistant', content: 'hello' },
+					{ role: 'user', content: 'more' },
+				],
+			}),
+		])
+
+		const mgr = new CheckpointManager(store)
+		const entries = await mgr.listEntries()
+
+		expect(entries).toHaveLength(2)
+		expect(entries[0]).toEqual({
+			id: 'cp_a',
+			runId: 'run_test',
+			iteration: 1,
+			createdAt: 1000,
+			messageCount: 2,
+		})
+		expect(entries[1]).toEqual({
+			id: 'cp_b',
+			runId: 'run_test',
+			iteration: 2,
+			createdAt: 2000,
+			messageCount: 3,
+		})
+	})
+
+	it('returns empty array when no checkpoints exist', async () => {
+		const mgr = new CheckpointManager(makeStoreStub([]))
+		expect(await mgr.listEntries()).toEqual([])
+	})
+
+	it('does not include full checkpoint payload fields', async () => {
+		const mgr = new CheckpointManager(
+			makeStoreStub([
+				makeCheckpoint({
+					toolResultHashes: { call_x: 'hash' },
+					branchStack: [{ agentId: 'a', decision: 'd', confidence: 1, timestamp: 0 }],
+				}),
+			]),
+		)
+		const [entry] = await mgr.listEntries()
+		expect(entry).not.toHaveProperty('tokenUsage')
+		expect(entry).not.toHaveProperty('costInfo')
+		expect(entry).not.toHaveProperty('toolResultHashes')
+		expect(entry).not.toHaveProperty('branchStack')
+		expect(entry).not.toHaveProperty('guardState')
+		expect(entry).not.toHaveProperty('messages')
+	})
+})
+
+function makeEmergencyDump(overrides: Partial<EmergencySaveData> = {}): EmergencySaveData {
+	return {
+		id: 'esave_dump_a' as EmergencySaveId,
+		runId: 'run_victim' as RunId,
+		messages: [
+			{ role: 'user', content: 'before the crash' },
+			{ role: 'assistant', content: 'working' },
+		],
+		tokenUsage: {
+			promptTokens: 10,
+			completionTokens: 5,
+			totalTokens: 15,
+			cachedTokens: 0,
+			cacheWriteTokens: 0,
+		},
+		currentIteration: 7,
+		startedAt: 1_000,
+		savedAt: 2_500,
+		processSignal: 'SIGTERM',
+		...overrides,
+	}
+}
+
+describe('projectEmergencyToCheckpoint', () => {
+	it('produces an IterationCheckpoint with all required fields', () => {
+		const dump = makeEmergencyDump()
+		const cp = projectEmergencyToCheckpoint(dump)
+
+		expect(cp.runId).toBe('run_victim')
+		expect(cp.iteration).toBe(7)
+		expect(cp.messages).toBe(dump.messages)
+		expect(cp.tokenUsage).toBe(dump.tokenUsage)
+		expect(cp.createdAt).toBe(2_500)
+		expect(cp.guardState).toEqual({ iterationCount: 7, elapsedMs: 1_500 })
+		expect(cp.costInfo).toEqual({
+			inputCostPer1M: 0,
+			outputCostPer1M: 0,
+			totalCost: 0,
+			cacheDiscount: 0,
+		})
+	})
+
+	it('derives a deterministic CheckpointId from the emergency save id', () => {
+		const dump = makeEmergencyDump({ id: 'esave_xyz123' as EmergencySaveId })
+		const cp1 = projectEmergencyToCheckpoint(dump)
+		const cp2 = projectEmergencyToCheckpoint(dump)
+
+		expect(cp1.id).toBe('cp_emergency_xyz123')
+		expect(cp1.id).toBe(cp2.id)
+	})
+
+	it('clamps guardState.elapsedMs to 0 when savedAt precedes startedAt', () => {
+		const cp = projectEmergencyToCheckpoint(makeEmergencyDump({ startedAt: 2_000, savedAt: 1_000 }))
+		expect(cp.guardState.elapsedMs).toBe(0)
+	})
+
+	it('leaves optional fields (toolResultHashes, branchStack, activeNode) unset', () => {
+		const cp = projectEmergencyToCheckpoint(makeEmergencyDump())
+		expect(cp.toolResultHashes).toBeUndefined()
+		expect(cp.branchStack).toBeUndefined()
+		expect(cp.activeNode).toBeUndefined()
+	})
+})

package/src/runtime/query/checkpoint.ts

@@ -8,9 +8,52 @@ import type {
 	IterationCheckpoint,
 } from '../../types/hitl/index.js'
 import type { AssistantMessage } from '../../types/message/index.js'
+import type { EmergencySaveData } from '../../types/run/emergency.js'
+import type { CheckpointListEntry } from '../../types/run/replay.js'
+import { ZERO_COST } from '../../utils/cost.js'
 import { buildToolResultHashes } from '../../utils/hash.js'
 import { generateCheckpointId } from '../../utils/id.js'
 
+function toCheckpointListEntry(cp: IterationCheckpoint): CheckpointListEntry {
+	return {
+		id: cp.id,
+		runId: cp.runId,
+		iteration: cp.iteration,
+		createdAt: cp.createdAt,
+		messageCount: cp.messages.length,
+	}
+}
+
+/**
+ * Project an {@link EmergencySaveData} dump to an {@link IterationCheckpoint}
+ * shape so `replay({ fromCheckpoint: 'emergency' })` can consume it through
+ * the same restore path as any other checkpoint.
+ *
+ * The projection is lossy: `costInfo`, `guardState.elapsedMs`,
+ * `toolResultHashes`, `branchStack`, and `activeNode` are not captured at
+ * emergency-save time and default to zero/empty values. The synthetic
+ * checkpoint id is derived deterministically from the emergency save id so
+ * re-projecting the same dump yields the same {@link CheckpointId}.
+ *
+ * See ses_005-deterministic-replay design §2 + §5.2.
+ */
+export function projectEmergencyToCheckpoint(dump: EmergencySaveData): IterationCheckpoint {
+	const emergencySuffix = dump.id.replace(/^esave_/, '')
+	return {
+		id: `cp_emergency_${emergencySuffix}` as CheckpointId,
+		runId: dump.runId,
+		iteration: dump.currentIteration,
+		messages: dump.messages,
+		tokenUsage: dump.tokenUsage,
+		costInfo: { ...ZERO_COST },
+		guardState: {
+			iterationCount: dump.currentIteration,
+			elapsedMs: Math.max(0, dump.savedAt - dump.startedAt),
+		},
+		createdAt: dump.savedAt,
+	}
+}
+
 export class CheckpointManager {
 	private store: RunDiskStore
 
@@ -60,6 +103,17 @@ export class CheckpointManager {
 		return this.store.listCheckpoints()
 	}
 
+	/**
+	 * Listing projection used by the public `listCheckpoints` API. Returns
+	 * only the fields a consumer needs to pick a fork point for
+	 * {@link import('./replay/prepare.js').prepareReplayState} — not the
+	 * full checkpoint payload. See ses_005-deterministic-replay design §3.1.
+	 */
+	async listEntries(): Promise<CheckpointListEntry[]> {
+		const checkpoints = await this.store.listCheckpoints()
+		return checkpoints.map(toCheckpointListEntry)
+	}
+
 	async prune(keepLast: number): Promise<void> {
 		const all = await this.list()
 		if (all.length <= keepLast) return

package/src/runtime/query/context.ts

@@ -99,11 +99,6 @@ export interface RunContext {
  */
 const migrationPromises = new Map<string, Promise<FilesystemMigrationResult>>()
 
-/** Testing hook — clears the first-call guard cache. */
-export function __resetMigrationGuardForTests(): void {
-	migrationPromises.clear()
-}
-
 export class RunContextFactory {
 	/**
 	 * Run the boot-time filesystem migration for `rootDir` at most once per

package/src/runtime/query/index.ts

@@ -27,7 +27,7 @@ import { type Message, createSystemMessage } from '../../types/message/index.js'
 import type { AgentPersona } from '../../types/persona/index.js'
 import type { LLMProvider } from '../../types/provider/index.js'
 import type { TaskRouterConfig } from '../../types/router/index.js'
-import type { AgentRun, AgentRunConfig, RunEvent, RunEventListener } from '../../types/run/index.js'
+import type { AgentRunConfig, Run, RunEvent, RunEventListener } from '../../types/run/index.js'
 import type { Sandbox, SandboxProvider } from '../../types/sandbox/index.js'
 import type { ProjectId, ThreadId } from '../../types/session/ids.js'
 import type { Skill } from '../../types/skills/index.js'
@@ -136,7 +136,7 @@ export interface QueryParams {
 	invocationState?: InvocationState
 }
 
-export async function* query(params: QueryParams): AsyncGenerator<RunEvent, AgentRun> {
+export async function* query(params: QueryParams): AsyncGenerator<RunEvent, Run> {
 	// Boot-time filesystem migration (session-hierarchy.md §13.4.1). First
 	// call per process per root actually runs; subsequent calls short-circuit
 	// via the in-memory guard in `context.ts`. Kept here rather than inside
@@ -335,7 +335,7 @@ export async function* query(params: QueryParams): AsyncGenerator<RunEvent, Agen
 
 	const tracer = getTracer()
 
-	return yield* (async function* (): AsyncGenerator<RunEvent, AgentRun> {
+	return yield* (async function* (): AsyncGenerator<RunEvent, Run> {
 		const rootSpan = tracer.startSpan(agentRunSpanName(params.agentName))
 		rootSpan.setAttributes({
 			[NAMZU.RUN_ID]: ctx.runMgr.id,
@@ -515,7 +515,7 @@ export async function* query(params: QueryParams): AsyncGenerator<RunEvent, Agen
 export async function drainQuery(
 	params: Omit<QueryParams, 'resumeHandler'> & { resumeHandler?: ResumeHandler },
 	listener?: RunEventListener,
-): Promise<AgentRun> {
+): Promise<Run> {
 	const fullParams: QueryParams = {
 		...params,
 		resumeHandler: params.resumeHandler ?? autoApproveHandler,