From 126b4234f3e233d1df2d48227b24ddd2807d76bf Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Wed, 29 Jul 2026 23:41:02 -0600 Subject: [PATCH 1/3] fix(supervise): preserve Pi RPC receipts --- docs/api/runtime.md | 10 ++ src/runtime/supervise/pi-executor.ts | 128 ++++++++++++++++++++------ src/runtime/supervise/trace-source.ts | 5 + tests/runtime/pi-executor.test.ts | 98 ++++++++++++++++++-- 4 files changed, 207 insertions(+), 34 deletions(-) diff --git a/docs/api/runtime.md b/docs/api/runtime.md index 7bbea49f..9177e8f9 100644 --- a/docs/api/runtime.md +++ b/docs/api/runtime.md @@ -12146,6 +12146,12 @@ How the settled ledger becomes the run's output (both arms). Default `bestDelive > `readonly` **args**: `unknown` +##### argsCaptured? + +> `readonly` `optional` **argsCaptured?**: `boolean` + +False when the call was observed but its original arguments were unavailable. + ##### status? > `readonly` `optional` **status?**: `"error"` \| `"ok"` @@ -12154,6 +12160,10 @@ How the settled ledger becomes the run's output (both arms). Default `bestDelive > `readonly` `optional` **result?**: `unknown` +##### error? + +> `readonly` `optional` **error?**: `string` + ##### callId? > `readonly` `optional` **callId?**: `string` diff --git a/src/runtime/supervise/pi-executor.ts b/src/runtime/supervise/pi-executor.ts index a8bcc3a1..4017e3e8 100644 --- a/src/runtime/supervise/pi-executor.ts +++ b/src/runtime/supervise/pi-executor.ts @@ -24,9 +24,10 @@ * are read structurally off JSON lines, so a pi that adds commands stays compatible and a pi that * is not installed simply fails loud at spawn instead of at import. * - * Usage accounting: pi reports token usage on its assistant messages when the provider supplies - * it. Nothing is fabricated — a turn whose usage pi does not report contributes an `iteration` - * event and zero tokens, exactly like the other honest executors. + * Usage accounting: pi repeats the same assistant receipt in `message_end` and `turn_end`. + * Only `turn_end` is counted, once per model call. Pi reports cache traffic separately from fresh + * input, so all three input classes are folded into Runtime's two-field token total. Subscription + * usage whose dollar price is absent or zero stays explicitly unknown on the terminal artifact. * * @experimental */ @@ -78,6 +79,11 @@ interface PiEvent { message?: unknown } +interface PendingPiTool { + args: unknown + startedAt: number +} + /** Build the `Executor` for one pi worker. Registered as runtime `'pi'`. */ export const piExecutor: ExecutorFactory = (spec, ctx) => { const seam = readPiSeam(ctx) @@ -181,6 +187,8 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { state.note = 'connected' const events: PiEvent[] = [] + const pendingTools = new Map() + let usdKnown = true let idle = false let failure: Error | undefined let wake: (() => void) | undefined @@ -239,7 +247,9 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { // Drain what pi has emitted so far, projecting usage + activity. while (events.length > 0) { const ev = events.shift() as PiEvent - for (const usage of projectPiEvent(ev, args, tokens)) { + const projected = projectPiEvent(ev, args, tokens, pendingTools) + if (projected.usdKnown === false) usdKnown = false + for (const usage of projected.events) { if (usage.kind === 'cost') usd += usage.usd yield usage } @@ -282,7 +292,13 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { await killPi(proc, 2_000).catch(() => ({ destroyed: false })) } - const spent: Spend = { iterations: state.turns, tokens, usd, ms: Date.now() - started } + const spent: Spend = { + iterations: state.turns, + tokens, + usd, + ...(usdKnown ? {} : { usdKnown: false }), + ms: Date.now() - started, + } state.artifact = { outRef: `pi:${hash(state.lastText)}`, out: { content: state.lastText, turns: state.turns }, @@ -321,59 +337,96 @@ function projectPiEvent( ev: PiEvent, args: StreamPiArgs, tokens: { input: number; output: number }, -): UsageEvent[] { + pendingTools: Map, +): { events: UsageEvent[]; usdKnown?: false } { const out: UsageEvent[] = [] const at = Date.now() if (ev.type === 'tool_execution_start' && typeof ev.toolName === 'string') { args.activity.push({ at, kind: 'tool', label: ev.toolName, detail: describeArgs(ev.args) }) - return out + if (typeof ev.toolCallId === 'string') { + pendingTools.set(ev.toolCallId, { args: ev.args ?? {}, startedAt: at }) + } + return { events: out } } if (ev.type === 'tool_execution_end' && typeof ev.toolName === 'string') { const status = ev.isError === true ? 'error' : 'ok' + const callId = typeof ev.toolCallId === 'string' ? ev.toolCallId : undefined + const started = callId ? pendingTools.get(callId) : undefined + if (callId) pendingTools.delete(callId) + const error = status === 'error' ? describeToolError(ev.result) : undefined args.activity.push({ at, kind: 'tool', label: ev.toolName, status }) args.record({ toolName: ev.toolName, - args: ev.args ?? {}, + args: started?.args ?? {}, + ...(started ? {} : { argsCaptured: false }), status, - ...(typeof ev.toolCallId === 'string' ? { callId: ev.toolCallId } : {}), + ...(ev.result !== undefined ? { result: ev.result } : {}), + ...(error !== undefined ? { error } : {}), + ...(callId ? { callId } : {}), + ...(started ? { startedAt: started.startedAt } : {}), + endedAt: at, }) - return out + return { events: out } + } + if (ev.type === 'message_end') { + const text = readText(ev.message) + if (text) args.state.lastText = text + return { events: out } } - if (ev.type === 'turn_end' || ev.type === 'message_end') { + if (ev.type === 'turn_end') { const usage = readUsage(ev.message) if (usage && (usage.input || usage.output)) { tokens.input += usage.input tokens.output += usage.output out.push({ kind: 'tokens', input: usage.input, output: usage.output }) } - if (usage?.usd) out.push({ kind: 'cost', usd: usage.usd }) + if (usage?.usd !== undefined) out.push({ kind: 'cost', usd: usage.usd }) const text = readText(ev.message) if (text) args.state.lastText = text - if (ev.type === 'turn_end') { - args.state.turns += 1 - args.activity.push({ at, kind: 'turn', label: `turn ${args.state.turns}` }) - out.push({ kind: 'iteration' }) + args.state.turns += 1 + args.activity.push({ at, kind: 'turn', label: `turn ${args.state.turns}` }) + out.push({ kind: 'iteration' }) + return { + events: out, + ...(!usage || usage.usdKnown === false ? { usdKnown: false } : {}), } } - return out + return { events: out } } -/** pi's assistant message carries provider usage when the provider reported it. Absent, nothing - * is fabricated — the turn still counts as an iteration with zero tokens. */ -function readUsage(message: unknown): { input: number; output: number; usd: number } | undefined { +/** Pi's fresh input excludes cache reads and writes. Runtime's input channel includes all model + * input, so combine them once at the assistant receipt. A missing or zero price is unknown because + * subscription-backed providers report zero even when compute was not free. */ +function readUsage( + message: unknown, +): { input: number; output: number; usd?: number; usdKnown: boolean } | undefined { if (!message || typeof message !== 'object') return undefined const usage = (message as { usage?: unknown }).usage if (!usage || typeof usage !== 'object') return undefined const u = usage as Record - const input = num(u.input) ?? num(u.inputTokens) ?? num(u.prompt_tokens) ?? 0 + const promptTokens = num(u.prompt_tokens) + const input = + promptTokens ?? + (num(u.input) ?? num(u.inputTokens) ?? 0) + + (num(u.cacheRead) ?? num(u.cache_read_input_tokens) ?? num(u.cacheReadInputTokens) ?? 0) + + (num(u.cacheWrite) ?? + num(u.cache_creation_input_tokens) ?? + num(u.cacheCreationInputTokens) ?? + 0) const output = num(u.output) ?? num(u.outputTokens) ?? num(u.completion_tokens) ?? 0 const costRaw = u.cost - const usd = + const reportedUsd = num(costRaw) ?? (costRaw && typeof costRaw === 'object' - ? (num((costRaw as Record).totalCost) ?? 0) - : 0) - return { input, output, usd } + ? (num((costRaw as Record).total) ?? + num((costRaw as Record).totalCost)) + : undefined) + return { + input, + output, + ...(reportedUsd !== undefined && reportedUsd > 0 ? { usd: reportedUsd } : {}), + usdKnown: reportedUsd !== undefined && reportedUsd > 0, + } } function readText(message: unknown): string | undefined { @@ -392,7 +445,7 @@ function readText(message: unknown): string | undefined { } function num(v: unknown): number | undefined { - return typeof v === 'number' && Number.isFinite(v) ? v : undefined + return typeof v === 'number' && Number.isFinite(v) && v >= 0 ? v : undefined } function describeArgs(argsValue: unknown): string | undefined { @@ -405,6 +458,29 @@ function describeArgs(argsValue: unknown): string | undefined { return undefined } +function describeToolError(result: unknown): string | undefined { + if (typeof result === 'string' && result.length > 0) return result + if (result && typeof result === 'object') { + const value = result as Record + const direct = value.error ?? value.message + if (typeof direct === 'string' && direct.length > 0) return direct + if (Array.isArray(value.content)) { + const text = value.content + .map((block) => + block && + typeof block === 'object' && + typeof (block as { text?: unknown }).text === 'string' + ? (block as { text: string }).text + : '', + ) + .filter(Boolean) + .join('\n') + if (text.length > 0) return text + } + } + return undefined +} + function taskText(task: unknown): string { if (typeof task === 'string') return task try { diff --git a/src/runtime/supervise/trace-source.ts b/src/runtime/supervise/trace-source.ts index a6aa4ec9..d5749a2b 100644 --- a/src/runtime/supervise/trace-source.ts +++ b/src/runtime/supervise/trace-source.ts @@ -24,8 +24,11 @@ import type { ToolPart, ToolState } from '@tangle-network/agent-interface' export interface ToolStepInput { readonly toolName: string readonly args: unknown + /** False when the call was observed but its original arguments were unavailable. */ + readonly argsCaptured?: boolean readonly status?: 'ok' | 'error' readonly result?: unknown + readonly error?: string /** Stable id of the tool call — used to de-duplicate the repeated state transitions a harness * streams for one call (opencode emits pending→running→completed, plus a `raw`-wrapped copy). */ readonly callId?: string @@ -56,7 +59,9 @@ export function toToolSpan(input: ToolStepInput, runId: string, seq: number, at: name: input.toolName, toolName: input.toolName, args: input.args, + ...(input.argsCaptured === false ? { argsCaptured: false } : {}), status: input.status ?? 'ok', + ...(input.error !== undefined ? { error: input.error } : {}), startedAt, endedAt, ...(input.result !== undefined ? { result: input.result } : {}), diff --git a/tests/runtime/pi-executor.test.ts b/tests/runtime/pi-executor.test.ts index 3ae65058..d019b441 100644 --- a/tests/runtime/pi-executor.test.ts +++ b/tests/runtime/pi-executor.test.ts @@ -1,5 +1,5 @@ /** - * `piExecutor` — pi wrapped behind `Executor`, driven against a FAKE pi that speaks pi 0.80.2's + * `piExecutor` — pi wrapped behind `Executor`, driven against a FAKE pi that speaks pi 0.83's * real RPC wire (`--mode rpc`, JSON lines on stdin/stdout, `AgentEvent` shapes). * * The point of the wrapper is that it delegates rather than reimplements: a message becomes pi's @@ -29,6 +29,7 @@ const FAKE_PI = `#!/usr/bin/env node const fs = require('node:fs') const log = process.env.PI_COMMAND_LOG const emit = (o) => process.stdout.write(JSON.stringify(o) + '\\n') +const result = (text) => ({ content: [{ type: 'text', text }], details: {} }) let buf = '' let turn = 0 process.stdin.on('data', (c) => { @@ -45,12 +46,37 @@ process.stdin.on('data', (c) => { if (cmd.type === 'abort') { process.exit(0) } if (cmd.type !== 'prompt') continue const target = String(cmd.message).includes('right.ts') ? 'right.ts' : 'wrong.ts' + const subscription = String(cmd.message).includes('subscription usage') + const parallel = String(cmd.message).includes('parallel tools') const myTurn = turn++ + const assistant = { + role: 'assistant', + content: [{ type: 'text', text: 'edited ' + target }], + usage: { + input: 30, + output: 12, + cacheRead: 5, + cacheWrite: 3, + totalTokens: 50, + cost: subscription + ? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } + : { input: 0.001, output: 0.001, cacheRead: 0, cacheWrite: 0, total: 0.002 }, + }, + } emit({ type: 'agent_start' }) emit({ type: 'turn_start' }) - emit({ type: 'tool_execution_start', toolCallId: 't' + myTurn, toolName: 'edit', args: { path: target } }) - emit({ type: 'tool_execution_end', toolCallId: 't' + myTurn, toolName: 'edit', result: 'ok', isError: false, args: { path: target } }) - emit({ type: 'turn_end', message: { role: 'assistant', content: 'edited ' + target, usage: { input: 30, output: 12, cost: 0.002 } }, toolResults: [] }) + emit({ type: 'message_start', message: assistant }) + emit({ type: 'message_end', message: assistant }) + if (parallel) { + emit({ type: 'tool_execution_start', toolCallId: 'read-' + myTurn, toolName: 'read', args: { path: 'alpha.ts' } }) + emit({ type: 'tool_execution_start', toolCallId: 'bash-' + myTurn, toolName: 'bash', args: { command: 'pnpm test' } }) + emit({ type: 'tool_execution_end', toolCallId: 'bash-' + myTurn, toolName: 'bash', result: result('tests passed'), isError: false }) + emit({ type: 'tool_execution_end', toolCallId: 'read-' + myTurn, toolName: 'read', result: result('permission denied'), isError: true }) + } else { + emit({ type: 'tool_execution_start', toolCallId: 't' + myTurn, toolName: 'edit', args: { path: target } }) + emit({ type: 'tool_execution_end', toolCallId: 't' + myTurn, toolName: 'edit', result: result('updated ' + target), isError: false }) + } + emit({ type: 'turn_end', message: assistant, toolResults: [] }) emit({ type: 'agent_end', messages: [] }) } }) @@ -109,7 +135,9 @@ describe('piExecutor — pi wrapped, not forked', () => { ) // REAL usage only — the numbers the fake pi reported, not a fabricated estimate. - expect(events).toContainEqual({ kind: 'tokens', input: 30, output: 12 }) + expect(events.filter((event) => event.kind === 'tokens')).toEqual([ + { kind: 'tokens', input: 38, output: 12 }, + ]) expect(events).toContainEqual({ kind: 'cost', usd: 0.002 }) expect(events.filter((e) => e.kind === 'iteration')).toHaveLength(1) @@ -117,12 +145,60 @@ describe('piExecutor — pi wrapped, not forked', () => { expect(progress?.turns).toBe(1) expect(progress?.recentActivity?.some((a) => a.label === 'edit')).toBe(true) - const spans = await ex.traceSource?.()?.collect() - expect((spans as ToolSpan[]).map((s) => s.toolName)).toContain('edit') + const spans = (await ex.traceSource?.()?.collect()) as ToolSpan[] + expect(spans).toHaveLength(1) + expect(spans[0]).toMatchObject({ + toolName: 'edit', + args: { path: 'wrong.ts' }, + status: 'ok', + result: { content: [{ type: 'text', text: 'updated wrong.ts' }], details: {} }, + }) const artifact = ex.resultArtifact() expect(String((artifact.out as { content: string }).content)).toContain('edited wrong.ts') - expect(artifact.spent.tokens).toEqual({ input: 30, output: 12 }) + expect(artifact.spent.tokens).toEqual({ input: 38, output: 12 }) + expect(artifact.spent.usd).toBe(0.002) + await ex.teardown('brutalKill') + }) + + it('joins parallel Pi 0.83 tool completions to their starts by toolCallId', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + + await drain(ex.execute('run parallel tools', ctx.signal) as AsyncIterable) + + const spans = (await ex.traceSource?.()?.collect()) as ToolSpan[] + expect(spans).toHaveLength(2) + expect(spans[0]).toMatchObject({ + toolName: 'bash', + args: { command: 'pnpm test' }, + status: 'ok', + result: { content: [{ type: 'text', text: 'tests passed' }], details: {} }, + }) + expect(spans[1]).toMatchObject({ + toolName: 'read', + args: { path: 'alpha.ts' }, + status: 'error', + error: 'permission denied', + result: { content: [{ type: 'text', text: 'permission denied' }], details: {} }, + }) + await ex.teardown('brutalKill') + }) + + it('marks subscription-priced Pi usage as unknown instead of known zero', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + + const events = await drain( + ex.execute('subscription usage', ctx.signal) as AsyncIterable, + ) + + // UsageEvent cannot currently carry an unknown-dollar marker, so the stream emits no fake + // $0 receipt and the direct terminal artifact retains the unknown state. + expect(events.filter((event) => event.kind === 'cost')).toEqual([]) + expect(ex.resultArtifact().spent).toMatchObject({ usd: 0, usdKnown: false }) await ex.teardown('brutalKill') }) @@ -167,8 +243,14 @@ describe('piExecutor — pi wrapped, not forked', () => { } // Two turns ran: the original and the steered one. expect(seen.filter((e) => e.kind === 'iteration').length).toBe(2) + expect(seen.filter((e) => e.kind === 'tokens')).toEqual([ + { kind: 'tokens', input: 38, output: 12 }, + { kind: 'tokens', input: 38, output: 12 }, + ]) const artifact = ex.resultArtifact() expect(String((artifact.out as { content: string }).content)).toContain('edited right.ts') + expect(artifact.spent.tokens).toEqual({ input: 76, output: 24 }) + expect(artifact.spent.usd).toBe(0.004) await ex.teardown('brutalKill') }) From 9091af519dcb287adcefb9ffbc1019f3dbc1a5c0 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 30 Jul 2026 00:47:43 -0600 Subject: [PATCH 2/3] fix(supervise): honor Pi terminal semantics --- docs/api/runtime.md | 46 ++- src/runtime/supervise/budget.ts | 59 +++- src/runtime/supervise/pi-executor.ts | 140 ++++++--- src/runtime/supervise/progress.ts | 4 + src/runtime/supervise/scope.ts | 19 +- src/runtime/supervise/types.ts | 7 +- tests/kernel/supervise.test.ts | 134 +++++++++ tests/runtime/pi-executor.test.ts | 423 ++++++++++++++++++++++++--- 8 files changed, 724 insertions(+), 108 deletions(-) diff --git a/docs/api/runtime.md b/docs/api/runtime.md index 9177e8f9..947eedcb 100644 --- a/docs/api/runtime.md +++ b/docs/api/runtime.md @@ -10263,7 +10263,7 @@ Extra args appended after `--mode rpc`. `--provider` / `--model` are added from > `optional` **turnTimeoutMs?**: `number` -Wall-clock ceiling for one `prompt` (the wait for `agent_end`). Omit = no timeout. +Wall-clock ceiling for one `prompt` (the wait for `agent_settled`). Omit = no timeout. ##### activityWindow? @@ -10407,6 +10407,12 @@ Metered iterations so far (the executor's own count when it reports one). > `readonly` **usd**: `number` +##### usdKnown? + +> `readonly` `optional` **usdKnown?**: `boolean` + +False when observed dollar spend is only a known subtotal, not a complete total. + ##### pendingMessages > `readonly` **pendingMessages**: `number` @@ -10519,6 +10525,10 @@ The scope-side facts about a child, independent of whether its executor cooperat > `readonly` **usd**: `number` +##### usdKnown? + +> `readonly` `optional` **usdKnown?**: `boolean` + *** ### InMemoryRunContextOptions @@ -16086,12 +16096,44 @@ How to run a sandboxed harness as the DRIVER, with the coordination verbs mounte ### UsageEvent -> **UsageEvent** = \{ `kind`: `"tokens"`; `input`: `number`; `output`: `number`; \} \| \{ `kind`: `"cost"`; `usd`: `number`; \} \| \{ `kind`: `"iteration"`; \} +> **UsageEvent** = \{ `kind`: `"tokens"`; `input`: `number`; `output`: `number`; \} \| \{ `kind`: `"cost"`; `usdKnown?`: `false`; `usd`: `number`; \} \| \{ `kind`: `"iteration"`; \} Normalized usage event — the single channel every executor reports through, so the conserved pool meters all runtimes identically. `tokens` carries `LoopTokenUsage`'s `{ input, output }`; `usd` is a SEPARATE channel (never folded into tokens). +#### Union Members + +##### Type Literal + +\{ `kind`: `"tokens"`; `input`: `number`; `output`: `number`; \} + +*** + +##### Type Literal + +\{ `kind`: `"cost"`; `usdKnown?`: `false`; `usd`: `number`; \} + +###### kind + +> **kind**: `"cost"` + +###### usdKnown? + +> `optional` **usdKnown?**: `false` + +Known dollar subtotal. When false, `usd` must not be treated as total cost. + +###### usd + +> **usd**: `number` + +*** + +##### Type Literal + +\{ `kind`: `"iteration"`; \} + *** ### Runtime diff --git a/src/runtime/supervise/budget.ts b/src/runtime/supervise/budget.ts index b3c88eba..7c8775b2 100644 --- a/src/runtime/supervise/budget.ts +++ b/src/runtime/supervise/budget.ts @@ -5,7 +5,7 @@ * quantities (tokens, usd, iterations) plus an absolute deadline. Children reserve * atomically at spawn and reconcile at settle: * - * total ≡ free + reserved + committed (invariant, always) + * total ≡ free + reserved + committed (for every known quantity) * * `reserve` moves a child's whole ceiling from `free` → `reserved` and fails closed * when `free` can't cover it (never read-then-spawn overcommit, so `Σk(treatment) ≡ @@ -15,7 +15,9 @@ * * Pure and deterministic: `now()` is injected, there is no I/O, and no wall-clock or * RNG read. A `reserve`/`reconcile` ticket is single-use (fail-loud on double or - * unknown reconcile) so a child can never refund twice. + * unknown reconcile) so a child can never refund twice. If dollar cost is unknowable under a + * dollar limit, reconciliation closes the known token/iteration work, marks the dollar channel + * unusable, and refuses later reservations; inventing a numeric dollar total would be worse. * * @experimental */ @@ -93,34 +95,50 @@ export interface BudgetPool { export function spendFromUsageEvents(events: UsageEvent[]): Spend { const tokens = zeroTokenUsage() let usd = 0 + let usdKnown = true let iterations = 0 for (const ev of events) { if (ev.kind === 'tokens') { addTokenUsage(tokens, { input: ev.input, output: ev.output }) } else if (ev.kind === 'cost') { usd += ev.usd + if (ev.usdKnown === false) usdKnown = false } else { iterations += 1 } } - return { iterations, tokens, usd, ms: 0 } + return { + iterations, + tokens, + usd, + ...(usdKnown ? {} : { usdKnown: false }), + ms: 0, + } } async function foldUsage(events: AsyncIterable | UsageEvent[]): Promise { if (Array.isArray(events)) return spendFromUsageEvents(events) const tokens = zeroTokenUsage() let usd = 0 + let usdKnown = true let iterations = 0 for await (const ev of events) { if (ev.kind === 'tokens') { addTokenUsage(tokens, { input: ev.input, output: ev.output }) } else if (ev.kind === 'cost') { usd += ev.usd + if (ev.usdKnown === false) usdKnown = false } else { iterations += 1 } } - return { iterations, tokens, usd, ms: 0 } + return { + iterations, + tokens, + usd, + ...(usdKnown ? {} : { usdKnown: false }), + ms: 0, + } } function totalTokens(usage: LoopTokenUsage): number { @@ -143,6 +161,7 @@ export function createBudgetPool(root: Budget, now: () => number = Date.now): Bu let freeUsd = root.maxUsd ?? 0 let reservedUsd = 0 let committedUsd = 0 + let usdTainted = false let freeIterations = root.maxIterations let reservedIterations = 0 @@ -159,6 +178,7 @@ export function createBudgetPool(root: Budget, now: () => number = Date.now): Bu const wantTokens = b.maxTokens const wantUsd = b.maxUsd ?? 0 const wantIterations = b.maxIterations + if (usdCapped && usdTainted) return { ok: false, reason: 'budget-exhausted' } // Fail-closed admission: every requested channel must fit the free balance. A // usd request against an uncapped root is unsatisfiable (the root declared no $). if (wantTokens > freeTokens) return { ok: false, reason: 'budget-exhausted' } @@ -188,14 +208,8 @@ export function createBudgetPool(root: Budget, now: () => number = Date.now): Bu if (!open.has(ticket.id)) { throw new Error(`budget pool: reconcile of unknown or already-settled ticket ${ticket.id}`) } - open.delete(ticket.id) - const { tokens: rTokens, usd: rUsd, iterations: rIterations } = ticket.reserved - if (usdCapped && spent.usdKnown === false) { - throw new Error( - `budget pool: ticket ${ticket.id} reported unknown dollar cost under a dollar-capped budget`, - ) - } + const unknownUnderCap = usdCapped && spent.usdKnown === false // Clamp actual spend to the reservation: a child must never commit more than it // reserved (that would overdraw the conserved pool). Over-spend is a fail-loud bug. @@ -218,6 +232,10 @@ export function createBudgetPool(root: Budget, now: () => number = Date.now): Bu throw new Error(`budget pool: ticket ${ticket.id} spent $${spent.usd} > reserved $${rUsd}`) } + // Ordinary validation errors leave the ticket open. Unknown dollars are different: the known + // channels still settle, then the dollar channel is permanently tainted and admission closes. + open.delete(ticket.id) + // Release the whole reservation, then commit actual spend; the difference is the // refund that flows back to `free`. reservedTokens -= rTokens @@ -228,15 +246,24 @@ export function createBudgetPool(root: Budget, now: () => number = Date.now): Bu committedIterations += spent.iterations freeIterations += rIterations - spent.iterations - if (usdCapped && rUsd > 0) { + if (usdCapped) { reservedUsd -= rUsd committedUsd += spent.usd - freeUsd += rUsd - spent.usd + if (unknownUnderCap) { + usdTainted = true + freeUsd = 0 + } else { + freeUsd += rUsd - spent.usd + } } else { - // Uncapped (or a zero-ceiling child under a capped root): record the observed spend - // without touching the reservation channel — usd is accounted, not conserved here. + // With no root dollar limit, record observed spend without a reservation channel. committedUsd += spent.usd } + if (unknownUnderCap) { + throw new Error( + `budget pool: ticket ${ticket.id} reported unknown dollar cost under a dollar-capped budget`, + ) + } } function observe(spend: Spend): void { @@ -261,7 +288,7 @@ export function createBudgetPool(root: Budget, now: () => number = Date.now): Bu function readout(): BudgetReadout { return { tokensLeft: freeTokens, - usdLeft: usdCapped ? freeUsd : 0, + usdLeft: usdCapped ? (usdTainted ? 0 : freeUsd) : 0, usdCapped, deadlineMs: absoluteDeadlineMs, reservedTokens, diff --git a/src/runtime/supervise/pi-executor.ts b/src/runtime/supervise/pi-executor.ts index 4017e3e8..32094e86 100644 --- a/src/runtime/supervise/pi-executor.ts +++ b/src/runtime/supervise/pi-executor.ts @@ -9,13 +9,13 @@ * already maintains. So this module is a thin protocol adapter, and every capability maps onto a * verb pi already has: * - * `execute` → `prompt`, draining pi's event stream until `agent_end` + * `execute` → `prompt`, draining pi's event stream until `agent_settled` * `deliver` → `prompt` with `streamingBehavior` — pi owns the queue, we do not * `teardown` → `abort`, then close stdin and reap the process * `progress` → pi's `tool_execution_start`/`_end` + `turn_end` events, plus `get_state`'s * `pendingMessageCount` mirrored locally so the read stays synchronous * `traceSource` → the same tool events decoded into the shared `ToolSpan` currency - * `resultArtifact` → the last assistant text collected off the stream + * `resultArtifact` → the terminal successful assistant turn * * It is registered through the DOCUMENTED extension point (`ExecutorRegistry.register('pi', …)`), * so nothing in the resolver switches on it and a consumer can replace it wholesale. @@ -27,13 +27,15 @@ * Usage accounting: pi repeats the same assistant receipt in `message_end` and `turn_end`. * Only `turn_end` is counted, once per model call. Pi reports cache traffic separately from fresh * input, so all three input classes are folded into Runtime's two-field token total. Subscription - * usage whose dollar price is absent or zero stays explicitly unknown on the terminal artifact. + * usage whose dollar price is absent or zero stays explicitly unknown through the live usage + * stream and terminal artifact. * * @experimental */ import { type ChildProcess, spawn } from 'node:child_process' import { ValidationError } from '../../errors' +import { abortError, throwIfAborted } from '../util' import { createInbox, type Inbox, type InboxMessage } from './inbox' import { type ActivityLog, createActivityLog, type ExecutorProgress } from './progress' import { createPushTraceSource, type ToolStepInput, type TraceSource } from './trace-source' @@ -62,7 +64,7 @@ export interface PiSeam { model?: string cwd?: string env?: Record - /** Wall-clock ceiling for one `prompt` (the wait for `agent_end`). Omit = no timeout. */ + /** Wall-clock ceiling for one `prompt` (the wait for `agent_settled`). Omit = no timeout. */ turnTimeoutMs?: number /** Newest-last activity window `progress()` reports. Default 12. */ activityWindow?: number @@ -84,6 +86,12 @@ interface PendingPiTool { startedAt: number } +interface PiAssistantOutcome { + text: string + stopReason?: string + errorMessage?: string +} + /** Build the `Executor` for one pi worker. Registered as runtime `'pi'`. */ export const piExecutor: ExecutorFactory = (spec, ctx) => { const seam = readPiSeam(ctx) @@ -171,16 +179,19 @@ interface StreamPiArgs { } /** - * One pi RPC session, run to `agent_end`. Every steer delivered while the turn is in flight is + * One pi RPC session, run to `agent_settled`. Every steer delivered while the turn is in flight is * forwarded to pi IMMEDIATELY (pi queues and delivers it at its own turn boundary — the queue we - * deliberately do not reimplement), and any steer still unread when pi goes idle re-prompts it, - * so a pi worker cannot settle while a message it never saw is pending. + * deliberately do not reimplement), and any steer still unread when pi settles re-prompts it. + * `agent_end` is deliberately not terminal: Pi may auto-retry or compact after emitting it. */ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { const { seam, inbox, activity, state } = args const started = Date.now() const tokens = { input: 0, output: 0 } let usd = 0 + let usdKnown = true + throwIfAborted(args.signal) + throwIfAborted(args.controller.signal) const proc = spawnPi(seam) state.proc = proc @@ -188,9 +199,11 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { const events: PiEvent[] = [] const pendingTools = new Map() - let usdKnown = true - let idle = false + let settled = false + let lastAssistant: PiAssistantOutcome | undefined let failure: Error | undefined + let abortDeadline: number | undefined + let exited = false let wake: (() => void) | undefined const notify = () => { const w = wake @@ -200,73 +213,84 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { const stdoutLines = readJsonLines(proc, (value) => { const ev = value as PiEvent - if (ev.type === 'agent_start') idle = false - if (ev.type === 'agent_end') idle = true + if (ev.type === 'agent_start') settled = false + if (ev.type === 'agent_settled') settled = true events.push(ev) notify() }) proc.once('exit', (code) => { - if (!idle && code !== 0 && code !== null) { - failure = new ValidationError(`piExecutor: pi exited with code ${code}`) + exited = true + if (!settled && abortDeadline === undefined) { + failure = new ValidationError( + `piExecutor: pi exited before agent_settled (code ${code ?? 'unknown'})`, + ) } - idle = true notify() }) proc.once('error', (e) => { failure = new ValidationError(`piExecutor: pi failed to start: ${e.message}`) - idle = true notify() }) const abortAll = () => { + if (abortDeadline !== undefined) return sendCommand(proc, { type: 'abort' }) - idle = true + abortDeadline = Date.now() + PI_ABORT_RECEIPT_MS + state.note = 'aborting' notify() } - if (args.signal.aborted || args.controller.signal.aborted) abortAll() - else { - args.signal.addEventListener('abort', abortAll, { once: true }) - args.controller.signal.addEventListener('abort', abortAll, { once: true }) - } + args.signal.addEventListener('abort', abortAll, { once: true }) + args.controller.signal.addEventListener('abort', abortAll, { once: true }) const system = args.spec.profile.prompt?.systemPrompt const opening = system ? `${system}\n\n${taskText(args.task)}` : taskText(args.task) - const deadline = seam.turnTimeoutMs ? Date.now() + seam.turnTimeoutMs : undefined try { + // Close the check→listener race without ever dispatching a cancelled task. + if (args.signal.aborted || args.controller.signal.aborted) throw abortError() sendCommand(proc, { type: 'prompt', message: opening }) state.note = 'turn 0' for (;;) { // Forward anything the driver delivered — pi's own queue is the single source of truth // for ordering, so this is a route, not a second queue. - if (!idle) forwardPending(proc, inbox, activity) + if (!settled && abortDeadline === undefined) forwardPending(proc, inbox, activity) // Drain what pi has emitted so far, projecting usage + activity. while (events.length > 0) { const ev = events.shift() as PiEvent const projected = projectPiEvent(ev, args, tokens, pendingTools) - if (projected.usdKnown === false) usdKnown = false + if (projected.assistant) lastAssistant = projected.assistant for (const usage of projected.events) { - if (usage.kind === 'cost') usd += usage.usd + if (usage.kind === 'cost') { + usd += usage.usd + if (usage.usdKnown === false) usdKnown = false + } yield usage } } if (failure) throw failure - if (deadline !== undefined && Date.now() > deadline) { + if (abortDeadline !== undefined && (settled || exited)) throw abortError() + // Once cancellation starts, Pi's terminal receipt gets its own bounded drain window. + if (abortDeadline === undefined && deadline !== undefined && Date.now() > deadline) { throw new ValidationError('piExecutor: turn exceeded turnTimeoutMs') } + if (abortDeadline !== undefined && Date.now() > abortDeadline) { + const error = abortError() + error.message = `piExecutor: abort did not settle within ${PI_ABORT_RECEIPT_MS}ms` + throw error + } - if (idle) { - // pi went idle. A steer delivered in the gap re-prompts it: a worker must not settle - // while an unread instruction is pending (the same contract every steerable runtime - // in this package honors). + if (settled) { + // A steer delivered in the settle gap starts a new Pi prompt. Pi has declared the prior + // session activity fully quiet, so no retry/compaction can race this transition. const pending = inbox.drain() if (pending.length === 0) break - idle = false + settled = false + lastAssistant = undefined sendCommand(proc, { type: 'prompt', message: inbox.fold(pending) }) state.note = `turn ${state.turns}` continue @@ -292,6 +316,18 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { await killPi(proc, 2_000).catch(() => ({ destroyed: false })) } + if (args.signal.aborted || args.controller.signal.aborted) throw abortError() + if (!lastAssistant) { + throw new ValidationError('piExecutor: agent_settled without an assistant turn') + } + if (lastAssistant.stopReason === 'aborted') throw abortError() + if (lastAssistant.stopReason === 'error') { + throw new ValidationError( + `piExecutor: Pi assistant failed: ${lastAssistant.errorMessage ?? 'unknown provider error'}`, + ) + } + state.lastText = lastAssistant.text + const spent: Spend = { iterations: state.turns, tokens, @@ -338,7 +374,7 @@ function projectPiEvent( args: StreamPiArgs, tokens: { input: number; output: number }, pendingTools: Map, -): { events: UsageEvent[]; usdKnown?: false } { +): { events: UsageEvent[]; assistant?: PiAssistantOutcome } { const out: UsageEvent[] = [] const at = Date.now() if (ev.type === 'tool_execution_start' && typeof ev.toolName === 'string') { @@ -368,11 +404,9 @@ function projectPiEvent( }) return { events: out } } - if (ev.type === 'message_end') { - const text = readText(ev.message) - if (text) args.state.lastText = text - return { events: out } - } + // Pi emits this for user, assistant, and toolResult messages. `turn_end.message` is the + // authoritative assistant receipt; a generic message can never become the result artifact. + if (ev.type === 'message_end') return { events: out } if (ev.type === 'turn_end') { const usage = readUsage(ev.message) if (usage && (usage.input || usage.output)) { @@ -380,16 +414,16 @@ function projectPiEvent( tokens.output += usage.output out.push({ kind: 'tokens', input: usage.input, output: usage.output }) } - if (usage?.usd !== undefined) out.push({ kind: 'cost', usd: usage.usd }) - const text = readText(ev.message) - if (text) args.state.lastText = text + out.push( + usage?.usd !== undefined + ? { kind: 'cost', usd: usage.usd } + : { kind: 'cost', usd: 0, usdKnown: false }, + ) args.state.turns += 1 args.activity.push({ at, kind: 'turn', label: `turn ${args.state.turns}` }) out.push({ kind: 'iteration' }) - return { - events: out, - ...(!usage || usage.usdKnown === false ? { usdKnown: false } : {}), - } + const assistant = readAssistantOutcome(ev.message) + return { events: out, ...(assistant ? { assistant } : {}) } } return { events: out } } @@ -397,9 +431,7 @@ function projectPiEvent( /** Pi's fresh input excludes cache reads and writes. Runtime's input channel includes all model * input, so combine them once at the assistant receipt. A missing or zero price is unknown because * subscription-backed providers report zero even when compute was not free. */ -function readUsage( - message: unknown, -): { input: number; output: number; usd?: number; usdKnown: boolean } | undefined { +function readUsage(message: unknown): { input: number; output: number; usd?: number } | undefined { if (!message || typeof message !== 'object') return undefined const usage = (message as { usage?: unknown }).usage if (!usage || typeof usage !== 'object') return undefined @@ -425,7 +457,17 @@ function readUsage( input, output, ...(reportedUsd !== undefined && reportedUsd > 0 ? { usd: reportedUsd } : {}), - usdKnown: reportedUsd !== undefined && reportedUsd > 0, + } +} + +function readAssistantOutcome(message: unknown): PiAssistantOutcome | undefined { + if (!message || typeof message !== 'object') return undefined + const value = message as Record + if (value.role !== 'assistant') return undefined + return { + text: readText(message) ?? '', + ...(typeof value.stopReason === 'string' ? { stopReason: value.stopReason } : {}), + ...(typeof value.errorMessage === 'string' ? { errorMessage: value.errorMessage } : {}), } } @@ -549,6 +591,8 @@ function readJsonLines(proc: ChildProcess, onValue: (value: unknown) => void): ( * Without it the SIGTERM races the command down the pipe and pi never sees the abort at all — * which defeats the reason to wrap pi rather than kill it (a clean abort finalizes its session). */ const PI_ABORT_GRACE_MS = 500 +/** Maximum wait for Pi's aborted receipt and `agent_settled` before forced teardown. */ +const PI_ABORT_RECEIPT_MS = 2_000 async function killPi( proc: ChildProcess, diff --git a/src/runtime/supervise/progress.ts b/src/runtime/supervise/progress.ts index 441865e9..6b26b654 100644 --- a/src/runtime/supervise/progress.ts +++ b/src/runtime/supervise/progress.ts @@ -72,6 +72,8 @@ export interface WorkerProgress { readonly turns: number readonly tokens: { readonly input: number; readonly output: number } readonly usd: number + /** False when observed dollar spend is only a known subtotal, not a complete total. */ + readonly usdKnown?: boolean /** Steers delivered but not yet read by the worker. */ readonly pendingMessages: number /** Newest-last window of tool/turn activity; empty when the executor exposes none. */ @@ -113,6 +115,7 @@ export interface ScopeProgressInput { readonly turns: number readonly tokens: { readonly input: number; readonly output: number } readonly usd: number + readonly usdKnown?: boolean } /** Fold the scope-derived facts and the executor's optional enrichment into one read. Pure: the @@ -146,6 +149,7 @@ export function readWorkerProgress( turns: executor?.turns ?? scope.turns, tokens: scope.tokens, usd: scope.usd, + ...(scope.usdKnown === false ? { usdKnown: false } : {}), pendingMessages: executor?.pendingMessages ?? 0, recentActivity: executor?.recentActivity ?? [], ...(note ? { note } : {}), diff --git a/src/runtime/supervise/scope.ts b/src/runtime/supervise/scope.ts index 902f4a5c..10358139 100644 --- a/src/runtime/supervise/scope.ts +++ b/src/runtime/supervise/scope.ts @@ -720,6 +720,7 @@ export function createScope(args: ScopeArgs): Scope { turns: child.spent.iterations, tokens: child.spent.tokens, usd: child.spent.usd, + ...(child.spent.usdKnown === false ? { usdKnown: false } : {}), }, fromExecutor, opts.now ?? now(), @@ -1171,6 +1172,7 @@ async function foldStream( ): Promise { const tokens = { input: 0, output: 0 } let usd = 0 + let usdKnown = true let iterations = 0 for await (const ev of stream) { if (ev.kind === 'tokens') { @@ -1178,12 +1180,25 @@ async function foldStream( tokens.output += ev.output } else if (ev.kind === 'cost') { usd += ev.usd + if (ev.usdKnown === false) usdKnown = false } else { iterations += 1 } - onProgress?.({ iterations, tokens: { ...tokens }, usd, ms: 0 }) + onProgress?.({ + iterations, + tokens: { ...tokens }, + usd, + ...(usdKnown ? {} : { usdKnown: false }), + ms: 0, + }) + } + return { + iterations, + tokens, + usd, + ...(usdKnown ? {} : { usdKnown: false }), + ms: 0, } - return { iterations, tokens, usd, ms: 0 } } /** Clamp a child's reported spend to its reservation so the pool's fail-loud over-spend diff --git a/src/runtime/supervise/types.ts b/src/runtime/supervise/types.ts index f382b129..e9ff024c 100644 --- a/src/runtime/supervise/types.ts +++ b/src/runtime/supervise/types.ts @@ -170,7 +170,12 @@ export interface ExecutorResult { */ export type UsageEvent = | { kind: 'tokens'; input: number; output: number } - | { kind: 'cost'; usd: number } + | { + kind: 'cost' + /** Known dollar subtotal. When false, `usd` must not be treated as total cost. */ + usdKnown?: false + usd: number + } | { kind: 'iteration' } /** The runtime tag of a `Executor` impl. Open by intent: custom runtimes use their own string name. diff --git a/tests/kernel/supervise.test.ts b/tests/kernel/supervise.test.ts index 9f2de820..b95683aa 100644 --- a/tests/kernel/supervise.test.ts +++ b/tests/kernel/supervise.test.ts @@ -231,6 +231,16 @@ describe('conserved budget pool', () => { ms: 0, }), ).toThrow(/unknown dollar cost/) + expect(pool.readout()).toMatchObject({ + tokensLeft: 900, + reservedTokens: 0, + usdLeft: 0, + }) + expect(() => pool.assertNoOpenTickets()).not.toThrow() + expect(pool.reserve({ maxIterations: 1, maxTokens: 10 } as Budget)).toEqual({ + ok: false, + reason: 'budget-exhausted', + }) }) it('spendFromUsageEvents folds tokens + usd on separate channels', () => { @@ -242,11 +252,135 @@ describe('conserved budget pool', () => { ]) expect(spend).toEqual({ iterations: 1, tokens: { input: 12, output: 8 }, usd: 0.01, ms: 0 }) }) + + it('preserves explicitly unknown dollar cost in sync and async usage folds', async () => { + const events: UsageEvent[] = [ + { kind: 'tokens', input: 12, output: 3 }, + { kind: 'cost', usd: 0, usdKnown: false }, + { kind: 'iteration' }, + ] + const expected: Spend = { + iterations: 1, + tokens: { input: 12, output: 3 }, + usd: 0, + usdKnown: false, + ms: 0, + } + + expect(spendFromUsageEvents(events)).toEqual(expected) + + const pool = createBudgetPool({ maxIterations: 2, maxTokens: 100 }, () => 0) + const stream = (async function* (): AsyncIterable { + yield* events + })() + await expect(pool.spendFrom(stream)).resolves.toEqual(expected) + }) + + it('keeps a dollar limit unusable after a concurrent known reservation refunds', () => { + const pool = createBudgetPool({ maxIterations: 2, maxTokens: 100, maxUsd: 1 }, () => 0) + const unknown = pool.reserve({ maxIterations: 1, maxTokens: 50, maxUsd: 0.5 } as Budget) + const known = pool.reserve({ maxIterations: 1, maxTokens: 50, maxUsd: 0.5 } as Budget) + if (!unknown.ok || !known.ok) throw new Error('both reservations should have succeeded') + + expect(() => + pool.reconcile(unknown.ticket, { + iterations: 1, + tokens: { input: 5, output: 5 }, + usd: 0, + usdKnown: false, + ms: 0, + }), + ).toThrow(/unknown dollar cost/) + pool.reconcile(known.ticket, { + iterations: 1, + tokens: { input: 5, output: 5 }, + usd: 0.1, + ms: 0, + }) + + expect(pool.readout()).toMatchObject({ usdLeft: 0, reservedTokens: 0 }) + expect(pool.reserve({ maxIterations: 1, maxTokens: 1 } as Budget)).toEqual({ + ok: false, + reason: 'budget-exhausted', + }) + expect(() => pool.assertNoOpenTickets()).not.toThrow() + }) }) // ── 2. equal-k by construction ────────────────────────────────────────────────── describe('equal-k by construction', () => { + it('preserves unknown dollar cost when Scope folds a streaming executor', async () => { + const { scope } = await beginScope() + const spawned = scope.spawn( + leafAgent('unknown-cost', { + out: 'complete', + events: [ + { kind: 'tokens', input: 12, output: 3 }, + { kind: 'cost', usd: 0, usdKnown: false }, + { kind: 'iteration' }, + ], + }), + 'task', + { label: 'unknown-cost', budget: { maxIterations: 1, maxTokens: 100 } }, + ) + expect(spawned.ok).toBe(true) + if (!spawned.ok) return + + const settled = await scope.next() + expect(settled?.kind).toBe('done') + expect(settled?.spent).toMatchObject({ + iterations: 1, + tokens: { input: 12, output: 3 }, + usd: 0, + usdKnown: false, + }) + expect(scope.progress(spawned.handle.id)).toMatchObject({ + usd: 0, + usdKnown: false, + }) + }) + + it('closes a dollar-limited reservation after unknown cost and refuses further spend', async () => { + const pool = createBudgetPool({ maxIterations: 2, maxTokens: 100, maxUsd: 1 }, () => 0) + const { scope } = await beginScope({ pool }) + const spawned = scope.spawn( + leafAgent('unknown-priced-child', { + out: 'complete', + events: [ + { kind: 'tokens', input: 12, output: 3 }, + { kind: 'cost', usd: 0, usdKnown: false }, + { kind: 'iteration' }, + ], + }), + 'task', + { + label: 'unknown-priced-child', + budget: { maxIterations: 1, maxTokens: 100, maxUsd: 0.5 }, + }, + ) + expect(spawned.ok).toBe(true) + + const settled = await scope.next() + expect(settled).toMatchObject({ + kind: 'down', + reason: expect.stringMatching(/unknown dollar cost/), + }) + expect(pool.readout()).toMatchObject({ + tokensLeft: 85, + reservedTokens: 0, + usdLeft: 0, + }) + expect(() => pool.assertNoOpenTickets()).not.toThrow() + + expect( + scope.spawn(leafAgent('after-unknown-cost', { out: 'should-not-run', events: [] }), 'task', { + label: 'after-unknown-cost', + budget: { maxIterations: 1, maxTokens: 1 }, + }), + ).toEqual({ ok: false, reason: 'budget-exhausted' }) + }) + it('two arms at equal per-child budget spend equal total iterations', async () => { // Each arm spawns 3 children at a fixed 1-iteration budget; both arms draw from a // pool sized for exactly 6, so the realized Σiterations is equal by the conserved diff --git a/tests/runtime/pi-executor.test.ts b/tests/runtime/pi-executor.test.ts index d019b441..35ca31f4 100644 --- a/tests/runtime/pi-executor.test.ts +++ b/tests/runtime/pi-executor.test.ts @@ -14,24 +14,83 @@ import { join } from 'node:path' import type { ToolSpan } from '@tangle-network/agent-eval' import type { AgentProfile } from '@tangle-network/agent-interface' import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { InMemoryResultBlobStore, InMemorySpawnJournal } from '../../src/durable/spawn-journal' +import { createBudgetPool } from '../../src/runtime/supervise/budget' +import { gateOnDeliverable } from '../../src/runtime/supervise/completion-gate' import { piExecutor, piSeamKey } from '../../src/runtime/supervise/pi-executor' -import type { AgentSpec, ExecutorContext, UsageEvent } from '../../src/runtime/supervise/types' +import { createExecutorRegistry } from '../../src/runtime/supervise/runtime' +import { createScope } from '../../src/runtime/supervise/scope' +import type { + Agent, + AgentSpec, + ExecutorContext, + UsageEvent, +} from '../../src/runtime/supervise/types' let dir: string let fakePi: string let commandLog: string /** - * A stand-in for `pi --mode rpc`: it records every command it receives, and scripts a worker that - * edits `wrong.ts` until told otherwise, then edits `right.ts`. Emits pi's own event shapes. + * A stand-in for `pi --mode rpc` using Pi 0.83's actual session event order. In particular: + * user and tool-result messages both emit `message_end`; each model call emits `turn_end`; + * `agent_end` may be followed by an automatic retry; only `agent_settled` means the session is + * idle. The scenarios below pin those distinctions instead of approximating the wire. */ const FAKE_PI = `#!/usr/bin/env node const fs = require('node:fs') const log = process.env.PI_COMMAND_LOG const emit = (o) => process.stdout.write(JSON.stringify(o) + '\\n') const result = (text) => ({ content: [{ type: 'text', text }], details: {} }) +const message = (m) => { + emit({ type: 'message_start', message: m }) + emit({ type: 'message_end', message: m }) +} +const pricedUsage = (input, output, cacheRead, cacheWrite, total) => ({ + input, + output, + cacheRead, + cacheWrite, + totalTokens: input + output + cacheRead + cacheWrite, + cost: { input: total / 2, output: total / 2, cacheRead: 0, cacheWrite: 0, total }, +}) +const freeUsage = (input, output, cacheRead, cacheWrite) => ({ + input, + output, + cacheRead, + cacheWrite, + totalTokens: input + output + cacheRead + cacheWrite, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}) +const assistant = (text, stopReason, usage, extra = []) => ({ + role: 'assistant', + content: [...(text ? [{ type: 'text', text }] : []), ...extra], + stopReason, + usage, +}) +const begin = (prompt) => { + const user = { role: 'user', content: [{ type: 'text', text: String(prompt) }] } + emit({ type: 'agent_start' }) + emit({ type: 'turn_start' }) + message(user) + return user +} +const finish = (messages, willRetry = false) => { + emit({ type: 'agent_end', messages, willRetry }) + if (!willRetry) emit({ type: 'agent_settled' }) +} +const finalTurn = (text, usage = pricedUsage(10, 4, 2, 1, 0.001)) => { + const final = assistant(text, 'stop', usage) + emit({ type: 'turn_start' }) + message(final) + emit({ type: 'turn_end', message: final, toolResults: [] }) + return final +} let buf = '' let turn = 0 +let pendingAbort +let abortScheduled = false +let exitOnAbort = false process.stdin.on('data', (c) => { buf += c.toString('utf8') for (;;) { @@ -43,41 +102,111 @@ process.stdin.on('data', (c) => { let cmd try { cmd = JSON.parse(line) } catch { continue } fs.appendFileSync(log, JSON.stringify(cmd) + '\\n') - if (cmd.type === 'abort') { process.exit(0) } + + if (cmd.type === 'abort') { + if (exitOnAbort) process.exit(0) + if (!pendingAbort || abortScheduled) continue + abortScheduled = true + const aborted = pendingAbort + pendingAbort = undefined + setTimeout(() => { + message(aborted) + emit({ type: 'turn_end', message: aborted, toolResults: [] }) + finish([aborted]) + }, 100) + continue + } if (cmd.type !== 'prompt') continue - const target = String(cmd.message).includes('right.ts') ? 'right.ts' : 'wrong.ts' - const subscription = String(cmd.message).includes('subscription usage') - const parallel = String(cmd.message).includes('parallel tools') + + const prompt = String(cmd.message) + const target = prompt.includes('right.ts') ? 'right.ts' : 'wrong.ts' const myTurn = turn++ - const assistant = { - role: 'assistant', - content: [{ type: 'text', text: 'edited ' + target }], - usage: { - input: 30, - output: 12, - cacheRead: 5, - cacheWrite: 3, - totalTokens: 50, - cost: subscription - ? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } - : { input: 0.001, output: 0.001, cacheRead: 0, cacheWrite: 0, total: 0.002 }, - }, + const user = begin(prompt) + + if (prompt.includes('delayed abort receipt')) { + pendingAbort = assistant( + 'cancelled after billed work', + 'aborted', + freeUsage(40, 10, 60, 5), + ) + continue + } + + if (prompt.includes('exit on abort')) { + exitOnAbort = true + continue + } + + if (prompt.includes('retryable provider error')) { + const failed = { + ...assistant('', 'error', freeUsage(0, 0, 0, 0)), + errorMessage: 'Stream ended without finish_reason', + } + message(failed) + emit({ type: 'turn_end', message: failed, toolResults: [] }) + emit({ type: 'agent_end', messages: [user, failed], willRetry: true }) + emit({ type: 'auto_retry_start', attempt: 1, maxAttempts: 3, delayMs: 25 }) + setTimeout(() => { + emit({ type: 'agent_start' }) + emit({ type: 'turn_start' }) + const recovered = assistant('recovered', 'stop', pricedUsage(10, 4, 2, 1, 0.001)) + message(recovered) + emit({ type: 'turn_end', message: recovered, toolResults: [] }) + finish([recovered]) + }, 25) + continue + } + + if (prompt.includes('terminal provider error')) { + const failed = { + ...assistant('', 'error', freeUsage(0, 0, 0, 0)), + errorMessage: 'Stream ended without finish_reason', + } + message(failed) + emit({ type: 'turn_end', message: failed, toolResults: [] }) + finish([user, failed]) + continue } - emit({ type: 'agent_start' }) - emit({ type: 'turn_start' }) - emit({ type: 'message_start', message: assistant }) - emit({ type: 'message_end', message: assistant }) + + if (prompt.includes('subscription usage')) { + const final = assistant('subscription result', 'stop', freeUsage(30, 12, 5, 3)) + message(final) + emit({ type: 'turn_end', message: final, toolResults: [] }) + finish([user, final]) + continue + } + + const parallel = prompt.includes('parallel tools') + const emptyFinal = prompt.includes('empty final assistant') + const calls = parallel + ? [ + { type: 'toolCall', id: 'read-' + myTurn, name: 'read', arguments: { path: 'alpha.ts' } }, + { type: 'toolCall', id: 'bash-' + myTurn, name: 'bash', arguments: { command: 'pnpm test' } }, + ] + : [{ type: 'toolCall', id: 'edit-' + myTurn, name: 'edit', arguments: { path: target } }] + const toolTurn = assistant('working', 'toolUse', pricedUsage(30, 12, 5, 3, 0.002), calls) + message(toolTurn) + + const toolResults = [] if (parallel) { emit({ type: 'tool_execution_start', toolCallId: 'read-' + myTurn, toolName: 'read', args: { path: 'alpha.ts' } }) emit({ type: 'tool_execution_start', toolCallId: 'bash-' + myTurn, toolName: 'bash', args: { command: 'pnpm test' } }) emit({ type: 'tool_execution_end', toolCallId: 'bash-' + myTurn, toolName: 'bash', result: result('tests passed'), isError: false }) emit({ type: 'tool_execution_end', toolCallId: 'read-' + myTurn, toolName: 'read', result: result('permission denied'), isError: true }) + toolResults.push( + { role: 'toolResult', toolCallId: 'read-' + myTurn, toolName: 'read', content: result('permission denied').content, isError: true }, + { role: 'toolResult', toolCallId: 'bash-' + myTurn, toolName: 'bash', content: result('tests passed').content, isError: false }, + ) } else { - emit({ type: 'tool_execution_start', toolCallId: 't' + myTurn, toolName: 'edit', args: { path: target } }) - emit({ type: 'tool_execution_end', toolCallId: 't' + myTurn, toolName: 'edit', result: result('updated ' + target), isError: false }) + emit({ type: 'tool_execution_start', toolCallId: 'edit-' + myTurn, toolName: 'edit', args: { path: target } }) + const toolText = emptyFinal ? 'TOOL_RESULT_MUST_NOT_BECOME_OUTPUT' : 'updated ' + target + emit({ type: 'tool_execution_end', toolCallId: 'edit-' + myTurn, toolName: 'edit', result: result(toolText), isError: false }) + toolResults.push({ role: 'toolResult', toolCallId: 'edit-' + myTurn, toolName: 'edit', content: result(toolText).content, isError: false }) } - emit({ type: 'turn_end', message: assistant, toolResults: [] }) - emit({ type: 'agent_end', messages: [] }) + for (const toolResult of toolResults) message(toolResult) + emit({ type: 'turn_end', message: toolTurn, toolResults }) + const final = finalTurn(emptyFinal ? '' : 'edited ' + target) + finish([user, toolTurn, ...toolResults, final]) } }) ` @@ -114,6 +243,19 @@ async function drain(iter: AsyncIterable): Promise { return out } +async function captureFailure(iter: AsyncIterable): Promise<{ + events: UsageEvent[] + error: Error | undefined +}> { + const events: UsageEvent[] = [] + try { + for await (const event of iter) events.push(event) + return { events, error: undefined } + } catch (error) { + return { events, error: error instanceof Error ? error : new Error(String(error)) } + } +} + async function readCommands(): Promise>> { const { readFile } = await import('node:fs/promises') const text = await readFile(commandLog, 'utf8') @@ -123,6 +265,17 @@ async function readCommands(): Promise>> { .map((l) => JSON.parse(l) as Record) } +async function waitForCommand( + predicate: (command: Record) => boolean, +): Promise { + const deadline = Date.now() + 1_000 + while (Date.now() < deadline) { + if ((await readCommands()).some(predicate)) return + await new Promise((resolve) => setTimeout(resolve, 10)) + } + throw new Error('timed out waiting for fake Pi command') +} + describe('piExecutor — pi wrapped, not forked', () => { it('runs a turn, reports REAL usage off pi events, and exposes live progress + tool spans', async () => { await writeFile(commandLog, '') @@ -137,12 +290,14 @@ describe('piExecutor — pi wrapped, not forked', () => { // REAL usage only — the numbers the fake pi reported, not a fabricated estimate. expect(events.filter((event) => event.kind === 'tokens')).toEqual([ { kind: 'tokens', input: 38, output: 12 }, + { kind: 'tokens', input: 13, output: 4 }, ]) expect(events).toContainEqual({ kind: 'cost', usd: 0.002 }) - expect(events.filter((e) => e.kind === 'iteration')).toHaveLength(1) + expect(events).toContainEqual({ kind: 'cost', usd: 0.001 }) + expect(events.filter((e) => e.kind === 'iteration')).toHaveLength(2) const progress = ex.progress?.() - expect(progress?.turns).toBe(1) + expect(progress?.turns).toBe(2) expect(progress?.recentActivity?.some((a) => a.label === 'edit')).toBe(true) const spans = (await ex.traceSource?.()?.collect()) as ToolSpan[] @@ -156,8 +311,8 @@ describe('piExecutor — pi wrapped, not forked', () => { const artifact = ex.resultArtifact() expect(String((artifact.out as { content: string }).content)).toContain('edited wrong.ts') - expect(artifact.spent.tokens).toEqual({ input: 38, output: 12 }) - expect(artifact.spent.usd).toBe(0.002) + expect(artifact.spent.tokens).toEqual({ input: 51, output: 16 }) + expect(artifact.spent.usd).toBe(0.003) await ex.teardown('brutalKill') }) @@ -195,13 +350,200 @@ describe('piExecutor — pi wrapped, not forked', () => { ex.execute('subscription usage', ctx.signal) as AsyncIterable, ) - // UsageEvent cannot currently carry an unknown-dollar marker, so the stream emits no fake - // $0 receipt and the direct terminal artifact retains the unknown state. - expect(events.filter((event) => event.kind === 'cost')).toEqual([]) + expect(events.filter((event) => event.kind === 'cost')).toEqual([ + { kind: 'cost', usd: 0, usdKnown: false }, + ]) expect(ex.resultArtifact().spent).toMatchObject({ usd: 0, usdKnown: false }) await ex.teardown('brutalKill') }) + it('waits for agent_settled so Pi can recover after agent_end requests a retry', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + + const events = await drain( + ex.execute('retryable provider error', ctx.signal) as AsyncIterable, + ) + + expect(events.filter((event) => event.kind === 'iteration')).toHaveLength(2) + expect(events.filter((event) => event.kind === 'tokens')).toEqual([ + { kind: 'tokens', input: 13, output: 4 }, + ]) + expect(ex.resultArtifact()).toMatchObject({ + out: { content: 'recovered', turns: 2 }, + spent: { usdKnown: false }, + }) + await ex.teardown('brutalKill') + }) + + it('rejects a terminal Pi error instead of saving the user prompt as a result', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + + const captured = await captureFailure( + ex.execute('terminal provider error', ctx.signal) as AsyncIterable, + ) + + expect(captured.error).toBeInstanceOf(Error) + expect(captured.error?.message).toContain('Stream ended without finish_reason') + expect(captured.events).toContainEqual({ kind: 'cost', usd: 0, usdKnown: false }) + expect(() => ex.resultArtifact()).toThrow(/before stream drained/) + await ex.teardown('brutalKill') + }) + + it('settles a terminal Pi error as down without checking or persisting a deliverable', async () => { + await writeFile(commandLog, '') + const root = 'pi-terminal-error' + const journal = new InMemorySpawnJournal() + await journal.beginTree(root, new Date(0).toISOString()) + const innerBlobs = new InMemoryResultBlobStore() + let blobWrites = 0 + const blobs = { + async put(outRef: string, artifact: unknown) { + blobWrites += 1 + await innerBlobs.put(outRef, artifact) + }, + get: (outRef: string) => innerBlobs.get(outRef), + } + let deliverableChecks = 0 + const executor = gateOnDeliverable(piExecutor(spec, piCtx()), { + check: () => { + deliverableChecks += 1 + return true + }, + }) + const agentSpec: AgentSpec = { ...spec, executor } + const agent = { + name: 'pi-error-worker', + act: async () => undefined, + executorSpec: agentSpec, + } as Agent & { executorSpec: AgentSpec } + const scope = createScope({ + parentId: root, + root, + pool: createBudgetPool({ maxIterations: 4, maxTokens: 100 }), + journal, + blobs, + executors: createExecutorRegistry(), + seams: {}, + depth: 0, + signal: new AbortController().signal, + now: () => 0, + }) + + const spawned = scope.spawn(agent, 'terminal provider error', { + budget: { maxIterations: 4, maxTokens: 100 }, + label: 'provider-error', + }) + expect(spawned.ok).toBe(true) + if (!spawned.ok) return + + const settled = await scope.next() + expect(settled).toMatchObject({ + kind: 'down', + reason: expect.stringContaining('Stream ended without finish_reason'), + }) + expect(deliverableChecks).toBe(0) + expect(blobWrites).toBe(0) + expect(scope.progress(spawned.handle.id)).toMatchObject({ + live: false, + status: 'failed', + turns: 1, + usd: 0, + usdKnown: false, + }) + const terminal = (await journal.loadTree(root))?.find((event) => event.kind === 'settled') + expect(terminal).toMatchObject({ + kind: 'settled', + status: 'down', + spent: { + iterations: 1, + tokens: { input: 0, output: 0 }, + usd: 0, + usdKnown: false, + }, + }) + expect(terminal).not.toHaveProperty('outRef') + }) + + it('uses only the terminal successful assistant turn as output, never a tool result', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + + await drain(ex.execute('empty final assistant', ctx.signal) as AsyncIterable) + + expect(ex.resultArtifact().out).toEqual({ content: '', turns: 2 }) + await ex.teardown('brutalKill') + }) + + it('refuses a pre-aborted execution before spawning Pi or sending a prompt', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + const controller = new AbortController() + controller.abort() + + const captured = await captureFailure( + ex.execute('must not run', controller.signal) as AsyncIterable, + ) + + expect(captured.error?.name).toBe('AbortError') + expect(captured.events).toEqual([]) + expect(await readCommands()).toEqual([]) + }) + + it('drains Pi’s delayed terminal receipt before surfacing an active abort', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + const controller = new AbortController() + const pending = captureFailure( + ex.execute( + 'wait for a delayed abort receipt', + controller.signal, + ) as AsyncIterable, + ) + await waitForCommand( + (command) => + command.type === 'prompt' && String(command.message).includes('delayed abort receipt'), + ) + + controller.abort() + const captured = await pending + + expect(captured.error?.name).toBe('AbortError') + expect(captured.events).toEqual([ + { kind: 'tokens', input: 105, output: 10 }, + { kind: 'cost', usd: 0, usdKnown: false }, + { kind: 'iteration' }, + ]) + expect((await readCommands()).map((command) => command.type)).toContain('abort') + expect(() => ex.resultArtifact()).toThrow(/before stream drained/) + }) + + it('keeps an active cancellation typed as AbortError when Pi exits without a receipt', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + const controller = new AbortController() + const pending = captureFailure( + ex.execute('exit on abort', controller.signal) as AsyncIterable, + ) + await waitForCommand( + (command) => command.type === 'prompt' && String(command.message).includes('exit on abort'), + ) + + controller.abort() + const captured = await pending + + expect(captured.error?.name).toBe('AbortError') + expect(captured.events).toEqual([]) + expect(() => ex.resultArtifact()).toThrow(/before stream drained/) + }) + it('uses pi’s state-safe prompt behavior for forceful and queued messages', async () => { await writeFile(commandLog, '') const ctx = piCtx() @@ -239,18 +581,21 @@ describe('piExecutor — pi wrapped, not forked', () => { if (!steered && e.kind === 'iteration') { steered = true ex.deliver?.({ steer: 'the change belongs in right.ts', interrupt: false }) + expect(ex.progress?.()?.pendingMessages).toBe(1) } } - // Two turns ran: the original and the steered one. - expect(seen.filter((e) => e.kind === 'iteration').length).toBe(2) + // Each prompt contains a tool turn and a terminal assistant turn. + expect(seen.filter((e) => e.kind === 'iteration').length).toBe(4) expect(seen.filter((e) => e.kind === 'tokens')).toEqual([ { kind: 'tokens', input: 38, output: 12 }, + { kind: 'tokens', input: 13, output: 4 }, { kind: 'tokens', input: 38, output: 12 }, + { kind: 'tokens', input: 13, output: 4 }, ]) const artifact = ex.resultArtifact() expect(String((artifact.out as { content: string }).content)).toContain('edited right.ts') - expect(artifact.spent.tokens).toEqual({ input: 76, output: 24 }) - expect(artifact.spent.usd).toBe(0.004) + expect(artifact.spent.tokens).toEqual({ input: 102, output: 32 }) + expect(artifact.spent.usd).toBe(0.006) await ex.teardown('brutalKill') }) From 460bd66798de256f54aca5bf1bec00995247f413 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 30 Jul 2026 01:02:13 -0600 Subject: [PATCH 3/3] fix(supervise): correlate Pi prompt settlement --- src/runtime/supervise/pi-executor.ts | 78 +++++++++++++++++++++------- tests/runtime/pi-executor.test.ts | 29 +++++++++++ 2 files changed, 89 insertions(+), 18 deletions(-) diff --git a/src/runtime/supervise/pi-executor.ts b/src/runtime/supervise/pi-executor.ts index 32094e86..08944b02 100644 --- a/src/runtime/supervise/pi-executor.ts +++ b/src/runtime/supervise/pi-executor.ts @@ -70,9 +70,13 @@ export interface PiSeam { activityWindow?: number } -/** Structural read of pi's `AgentEvent` union — only the fields this adapter consumes. */ +/** Structural read of Pi's stdout records: agent events plus correlated RPC responses. */ interface PiEvent { + id?: string type?: string + command?: string + success?: boolean + error?: string toolCallId?: string toolName?: string args?: unknown @@ -199,11 +203,15 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { const events: PiEvent[] = [] const pendingTools = new Map() + const awaitingPromptResponses = new Set() let settled = false + let acceptedPromptSinceSettlement = false + let promptSequence = 0 let lastAssistant: PiAssistantOutcome | undefined let failure: Error | undefined let abortDeadline: number | undefined let exited = false + let exitCode: number | null | undefined let wake: (() => void) | undefined const notify = () => { const w = wake @@ -213,19 +221,13 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { const stdoutLines = readJsonLines(proc, (value) => { const ev = value as PiEvent - if (ev.type === 'agent_start') settled = false - if (ev.type === 'agent_settled') settled = true events.push(ev) notify() }) proc.once('exit', (code) => { exited = true - if (!settled && abortDeadline === undefined) { - failure = new ValidationError( - `piExecutor: pi exited before agent_settled (code ${code ?? 'unknown'})`, - ) - } + exitCode = code notify() }) proc.once('error', (e) => { @@ -246,21 +248,52 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { const system = args.spec.profile.prompt?.systemPrompt const opening = system ? `${system}\n\n${taskText(args.task)}` : taskText(args.task) const deadline = seam.turnTimeoutMs ? Date.now() + seam.turnTimeoutMs : undefined + const sendPrompt = (message: string, streamingBehavior?: 'steer' | 'followUp'): void => { + const id = `agent-runtime-prompt-${++promptSequence}` + awaitingPromptResponses.add(id) + sendCommand(proc, { + id, + type: 'prompt', + message, + ...(streamingBehavior ? { streamingBehavior } : {}), + }) + } try { // Close the check→listener race without ever dispatching a cancelled task. if (args.signal.aborted || args.controller.signal.aborted) throw abortError() - sendCommand(proc, { type: 'prompt', message: opening }) + sendPrompt(opening) state.note = 'turn 0' for (;;) { // Forward anything the driver delivered — pi's own queue is the single source of truth // for ordering, so this is a route, not a second queue. - if (!settled && abortDeadline === undefined) forwardPending(proc, inbox, activity) + if (!settled && abortDeadline === undefined) forwardPending(inbox, activity, sendPrompt) // Drain what pi has emitted so far, projecting usage + activity. while (events.length > 0) { const ev = events.shift() as PiEvent + if ( + ev.type === 'response' && + ev.command === 'prompt' && + typeof ev.id === 'string' && + awaitingPromptResponses.delete(ev.id) + ) { + if (ev.success === true) { + // A settlement that preceded this acceptance belongs to an older prompt. Require a + // later `agent_settled` before ending the session. + acceptedPromptSinceSettlement = true + } else { + failure = new ValidationError( + `piExecutor: Pi rejected prompt: ${ev.error ?? 'unknown RPC error'}`, + ) + } + } + if (ev.type === 'agent_start') settled = false + if (ev.type === 'agent_settled') { + settled = true + acceptedPromptSinceSettlement = false + } const projected = projectPiEvent(ev, args, tokens, pendingTools) if (projected.assistant) lastAssistant = projected.assistant for (const usage of projected.events) { @@ -274,6 +307,15 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { if (failure) throw failure if (abortDeadline !== undefined && (settled || exited)) throw abortError() + if ( + abortDeadline === undefined && + exited && + (!settled || awaitingPromptResponses.size > 0 || acceptedPromptSinceSettlement) + ) { + throw new ValidationError( + `piExecutor: pi exited before agent_settled (code ${exitCode ?? 'unknown'})`, + ) + } // Once cancellation starts, Pi's terminal receipt gets its own bounded drain window. if (abortDeadline === undefined && deadline !== undefined && Date.now() > deadline) { throw new ValidationError('piExecutor: turn exceeded turnTimeoutMs') @@ -284,14 +326,14 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { throw error } - if (settled) { + if (settled && awaitingPromptResponses.size === 0 && !acceptedPromptSinceSettlement) { // A steer delivered in the settle gap starts a new Pi prompt. Pi has declared the prior // session activity fully quiet, so no retry/compaction can race this transition. const pending = inbox.drain() if (pending.length === 0) break settled = false lastAssistant = undefined - sendCommand(proc, { type: 'prompt', message: inbox.fold(pending) }) + sendPrompt(inbox.fold(pending)) state.note = `turn ${state.turns}` continue } @@ -343,14 +385,14 @@ async function* streamPiSession(args: StreamPiArgs): AsyncIterable { } /** Route messages through pi's state-safe prompt command; pi owns the queue and idle transition. */ -function forwardPending(proc: ChildProcess, inbox: Inbox, activity: ActivityLog): void { +function forwardPending( + inbox: Inbox, + activity: ActivityLog, + sendPrompt: (message: string, streamingBehavior?: 'steer' | 'followUp') => void, +): void { const pending = inbox.drain() for (const m of pending) { - sendCommand(proc, { - type: 'prompt', - message: renderOne(m), - streamingBehavior: m.interrupt ? 'steer' : 'followUp', - }) + sendPrompt(renderOne(m), m.interrupt ? 'steer' : 'followUp') activity.push({ at: Date.now(), kind: 'note', diff --git a/tests/runtime/pi-executor.test.ts b/tests/runtime/pi-executor.test.ts index 35ca31f4..e0789aae 100644 --- a/tests/runtime/pi-executor.test.ts +++ b/tests/runtime/pi-executor.test.ts @@ -104,6 +104,7 @@ process.stdin.on('data', (c) => { fs.appendFileSync(log, JSON.stringify(cmd) + '\\n') if (cmd.type === 'abort') { + emit({ id: cmd.id, type: 'response', command: 'abort', success: true }) if (exitOnAbort) process.exit(0) if (!pendingAbort || abortScheduled) continue abortScheduled = true @@ -119,6 +120,17 @@ process.stdin.on('data', (c) => { if (cmd.type !== 'prompt') continue const prompt = String(cmd.message) + if (prompt.includes('reject at preflight')) { + emit({ + id: cmd.id, + type: 'response', + command: 'prompt', + success: false, + error: 'No model selected', + }) + continue + } + emit({ id: cmd.id, type: 'response', command: 'prompt', success: true }) const target = prompt.includes('right.ts') ? 'right.ts' : 'wrong.ts' const myTurn = turn++ const user = begin(prompt) @@ -544,6 +556,23 @@ describe('piExecutor — pi wrapped, not forked', () => { expect(() => ex.resultArtifact()).toThrow(/before stream drained/) }) + it('fails explicitly instead of hanging when Pi rejects a prompt at preflight', async () => { + await writeFile(commandLog, '') + const ctx = piCtx() + const ex = piExecutor(spec, ctx) + + const captured = await captureFailure( + ex.execute('reject at preflight', ctx.signal) as AsyncIterable, + ) + + expect(captured.events).toEqual([]) + expect(captured.error).toMatchObject({ + name: 'ValidationError', + message: 'piExecutor: Pi rejected prompt: No model selected', + }) + expect(() => ex.resultArtifact()).toThrow(/before stream drained/) + }) + it('uses pi’s state-safe prompt behavior for forceful and queued messages', async () => { await writeFile(commandLog, '') const ctx = piCtx()