From ae81ac22238165dc243f0e56d07ba3c9474e61e9 Mon Sep 17 00:00:00 2001 From: Amish Pratap Singh Date: Tue, 21 Jul 2026 19:34:38 +0530 Subject: [PATCH] fix(sdk): add Query.continue() to resume after a failed model turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, when the model returned stopReason "error"/"aborted", query() surfaced a system error message and ended the session with no way to retry or resume — Query.steer() was also a no-op stub. Both are now wired into pi-agent-core's Agent instance: - steer(message) queues a steering message via agent.steer(). - continue(message?) resumes a session that has gone idle. With no argument it retries the exact turn the model failed to answer (dropping the failed assistant message so the transcript ends on the user/tool-result message it never responded to, per pi-agent-core's continue() contract). With a message, it queues a follow-up turn. The internal message channel can now reopen after finishing, since a completed query's async iterator may be resumed by continue(). --- src/sdk-types.ts | 15 ++++++++++ src/sdk.ts | 74 ++++++++++++++++++++++++++++++++++++++++++++---- test/sdk.test.ts | 61 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+), 5 deletions(-) diff --git a/src/sdk-types.ts b/src/sdk-types.ts index 20dc060..82ee604 100644 --- a/src/sdk-types.ts +++ b/src/sdk-types.ts @@ -158,6 +158,21 @@ export interface QueryOptions { export interface Query extends AsyncGenerator { abort(): void; steer(message: string): void; + /** + * Resume a session after it has gone idle — most commonly after the model + * failed to return output (an assistant message with `stopReason: "error"` + * or `"aborted"`) ended the run early. + * + * - `continue()` with no argument retries the exact turn the model failed + * to answer: the failed assistant message is dropped so the transcript + * ends on the user/tool-result message it never responded to, then the + * turn is re-sent. + * - `continue(message)` instead queues `message` as a new follow-up turn + * and resumes the session with it. + * + * Throws if the agent hasn't started yet or is still processing a turn. + */ + continue(message?: string): Promise; sessionId(): string; manifest(): AgentManifest; messages(): GCMessage[]; diff --git a/src/sdk.ts b/src/sdk.ts index dfa9a80..f669ab3 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -1,5 +1,5 @@ import { Agent } from "@mariozechner/pi-agent-core"; -import type { AgentEvent, AgentTool } from "@mariozechner/pi-agent-core"; +import type { AgentEvent, AgentMessage, AgentTool } from "@mariozechner/pi-agent-core"; import type { AssistantMessage } from "@mariozechner/pi-ai"; import { loadAgent } from "./loader.js"; import type { AgentManifest } from "./loader.js"; @@ -35,6 +35,9 @@ import { interface Channel { push(v: T): void; finish(): void; + // Undoes finish() so a completed query can resume streaming — used by + // Query.continue() to reopen the channel after the run first went idle. + reopen(): void; pull(): Promise>; } @@ -59,6 +62,9 @@ function createChannel(): Channel { resolve = null; } }, + reopen() { + done = false; + }, pull(): Promise> { if (buffer.length) { return Promise.resolve({ value: buffer.shift()!, done: false }); @@ -73,6 +79,14 @@ function createChannel(): Channel { // ── Extract text/thinking from AssistantMessage ──────────────────────── +function toUserMessage(text: string): AgentMessage { + return { + role: "user", + content: [{ type: "text", text }], + timestamp: Date.now(), + } as AgentMessage; +} + function extractContent(msg: AssistantMessage): { text: string; thinking: string } { let text = ""; let thinking = ""; @@ -111,6 +125,9 @@ export function query(options: QueryOptions): Query { let sandboxCtx: SandboxContext | undefined; // Local session (hoisted for cleanup in catch) let localSession: LocalSession | undefined; + // Agent instance (hoisted so steer()/continue() can reach it once the + // agent is loaded, including after the initial run has gone idle) + let agent: Agent | undefined; // OpenTelemetry session span — opened immediately so it covers agent // load + prompt + cleanup. Closed in the IIFE's finally so every exit @@ -299,7 +316,7 @@ export function query(options: QueryOptions): Query { } // 8. Create Agent - const agent = new Agent({ + agent = new Agent({ initialState: { systemPrompt, model: loaded.model, @@ -307,6 +324,9 @@ export function query(options: QueryOptions): Query { ...modelOptions, }, }); + // Non-null alias: `agent` is a `let` captured by closures below, so TS + // can't narrow it past the assignment above once inside them. + const activeAgent = agent; // 9. Subscribe to events and map to GCMessage agent.subscribe((event: AgentEvent) => { @@ -508,7 +528,7 @@ export function query(options: QueryOptions): Query { } } await otelContext.with(_session.ctx, () => - agent.prompt(options.prompt as string), + activeAgent.prompt(options.prompt as string), ); } else { // Multi-turn: iterate the async iterable @@ -532,7 +552,7 @@ export function query(options: QueryOptions): Query { } } await otelContext.with(_session.ctx, () => - agent.prompt(userMsg.content), + activeAgent.prompt(userMsg.content), ); } } @@ -593,7 +613,51 @@ export function query(options: QueryOptions): Query { ac.abort(); }, - steer(_message: string) { + steer(message: string) { + if (!agent) { + throw new Error("Cannot steer: agent not yet initialized (query hasn't started)"); + } + agent.steer(toUserMessage(message)); + }, + + async continue(message?: string) { + if (!agent) { + throw new Error("Cannot continue: agent not yet initialized (query hasn't started)"); + } + if (agent.signal) { + throw new Error("Cannot continue: agent is still processing a turn"); + } + + // The initial run may have already reached agent_end and finished + // the channel — reopen it so further messages can stream out. + channel.reopen(); + + try { + if (message !== undefined) { + pushMsg({ type: "user", content: message }); + agent.followUp(toUserMessage(message)); + } else { + // No new input: retry the turn the model failed to answer. + // Drop the failed assistant message so the transcript ends + // on the user/tool-result message it never responded to — + // pi-agent-core's continue() requires that to retry in place. + const messages = agent.state.messages as any[]; + const last = messages[messages.length - 1]; + if (last?.role === "assistant" && (last.stopReason === "error" || last.stopReason === "aborted")) { + agent.state.messages = messages.slice(0, -1); + } + } + + await otelContext.with(_session.ctx, () => (agent as Agent).continue()); + } catch (err: any) { + pushMsg({ + type: "system", + subtype: "error", + content: err.message, + }); + channel.finish(); + throw err; + } }, sessionId() { diff --git a/test/sdk.test.ts b/test/sdk.test.ts index a71fcc6..eaba799 100644 --- a/test/sdk.test.ts +++ b/test/sdk.test.ts @@ -271,3 +271,64 @@ describe("query()", () => { assert.ok(collected.length > 0); }); }); + +// ── query().continue() ────────────────────────────────────────────────── + +describe("query().continue()", () => { + it("exposes a continue method on the Query object", () => { + const q = query({ + prompt: "hello", + dir: "/nonexistent", + }); + + assert.equal(typeof q.continue, "function"); + q.return(); + }); + + it("rejects when called before the agent has finished loading", async () => { + const q = query({ + prompt: "hello", + dir: "/nonexistent", + }); + + await assert.rejects( + () => q.continue(), + (err: Error) => { + assert.ok(err.message.includes("not yet initialized")); + return true; + }, + ); + q.return(); + }); + + it("steer() also rejects before the agent has finished loading", async () => { + const q = query({ + prompt: "hello", + dir: "/nonexistent", + }); + + assert.throws( + () => q.steer("nudge"), + (err: Error) => { + assert.ok(err.message.includes("not yet initialized")); + return true; + }, + ); + q.return(); + }); + + it("continue() after a failed load surfaces the same 'not yet initialized' error rather than hanging", async () => { + const messages: any[] = []; + const q = query({ + prompt: "hello", + dir: "/nonexistent/path/to/agent", + }); + for await (const msg of q) { + messages.push(msg); + } + + // The initial run already failed (agent dir invalid) and never got + // far enough to construct an Agent instance. + await assert.rejects(() => q.continue()); + }); +});