From d30e13d6423bfa49445eb5ed0beb531bf63394f6 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 14 Sep 2026 17:54:54 -0700 Subject: [PATCH] Added a single-step run and a prompt record that calls no model. `Agent.step` advances the transcript by one model call and the tools it asks for, without adding a prompt and without looping, so a caller outside the loop can drive a turn one step at a time and checkpoint in between. `AgentSession.recordPrompt` does everything `prompt` does to build the turn, except run it; `AgentSession.step` runs one step and keeps the retry and compaction handling that `prompt` gets from the post-run pass. --- packages/agent/src/agent-loop.ts | 46 +++++ packages/agent/src/agent.ts | 43 ++++- packages/agent/test/agent-step.test.ts | 174 ++++++++++++++++++ .../coding-agent/src/core/agent-session.ts | 71 ++++++- .../test/resume-interrupted-turn.test.ts | 49 ++++- 5 files changed, 373 insertions(+), 10 deletions(-) create mode 100644 packages/agent/test/agent-step.test.ts diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 92876338f91..9b32fc65af3 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -143,6 +143,52 @@ export async function runAgentLoopContinue( return newMessages; } +export interface AgentStepOutcome { + messages: AgentMessage[]; + /** True when the turn ran tools whose results still need another step. */ + hasMoreToolCalls: boolean; +} + +/** + * Run exactly one iteration of the loop against the current context, adding no + * new message. Like agentLoopContinue, the last message must convert to a `user` + * or `toolResult` message. Used to drive a turn one step at a time from outside. + * + * A step is one turn, not a whole run, so it emits no `agent_start`. The run ends + * only when the turn itself ends it, on an error, an abort, or a stop decision. + */ +export async function runAgentStep( + context: AgentContext, + config: AgentLoopConfig, + emit: AgentEventSink, + signal: AbortSignal | undefined, + streamFn: StreamFn, +): Promise { + if (context.messages.length === 0) { + throw new Error("Cannot step: no messages in context"); + } + if (context.messages[context.messages.length - 1].role === "assistant") { + throw new Error("Cannot step from message role: assistant"); + } + + const newMessages: AgentMessage[] = []; + const currentContext: AgentContext = { ...context }; + + const outcome = await runSingleTurn({ + context: currentContext, + config, + newMessages, + pendingMessages: [], + emitTurnStart: true, + fetchNextPending: false, + signal, + emit, + streamFunction: streamFn ?? getDefaultStreamFn(), + }); + + return { messages: newMessages, hasMoreToolCalls: !outcome.done && outcome.hasMoreToolCalls }; +} + // What the transcript says about a call when nothing can say whether it ran. The tool can have had // its effect before the run stopped, so calling it a failure would invite a second run of something // that already happened. diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 0de7edd8302..abdef455b44 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -7,7 +7,7 @@ import type { ThinkingBudgets, Transport, } from "@earendil-works/pi-ai"; -import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts"; +import { type AgentStepOutcome, runAgentLoop, runAgentLoopContinue, runAgentStep } from "./agent-loop.ts"; import { getDefaultStreamFn } from "./stream-fn.ts"; import type { AfterToolCallContext, @@ -387,6 +387,31 @@ export class Agent { await this.runContinuation(); } + /** + * Advance the current transcript by exactly one step (one model call and the + * tools it requests). Adds no prompt and does not loop. The last message must + * be a user or tool-result message. `hasMoreToolCalls` tells the caller whether + * the turn still needs another step, so it does not repeat the loop's own + * termination logic. + */ + async step(): Promise { + if (this.activeRun) { + throw new Error("Agent is already processing. Wait for completion before stepping."); + } + + // Checked here as well as in the loop: the lifecycle below turns a thrown error into a failed + // assistant message, and a precondition failure should reach the caller instead. + const lastMessage = this._state.messages[this._state.messages.length - 1]; + if (!lastMessage) { + throw new Error("No messages to step from"); + } + if (lastMessage.role === "assistant") { + throw new Error("Cannot step from message role: assistant"); + } + + return await this.runSingleStep(); + } + private normalizePromptInput( input: string | AgentMessage | AgentMessage[], images?: ImageContent[], @@ -434,6 +459,22 @@ export class Agent { }); } + private async runSingleStep(): Promise { + // A failed run is handled inside the lifecycle, so the default stands and the + // caller stops stepping. + let outcome: AgentStepOutcome = { messages: [], hasMoreToolCalls: false }; + await this.runWithLifecycle(async (signal) => { + outcome = await runAgentStep( + this.createContextSnapshot(), + this.createLoopConfig(), + (event) => this.processEvents(event), + signal, + this.streamFunction, + ); + }); + return outcome; + } + private createContextSnapshot(): AgentContext { return { systemPrompt: this._state.systemPrompt, diff --git a/packages/agent/test/agent-step.test.ts b/packages/agent/test/agent-step.test.ts new file mode 100644 index 00000000000..f5ef819c12b --- /dev/null +++ b/packages/agent/test/agent-step.test.ts @@ -0,0 +1,174 @@ +import { + type AssistantMessage, + type AssistantMessageEvent, + EventStream, + type Message, + type Model, + type UserMessage, +} from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { describe, expect, it } from "vitest"; +import { runAgentStep } from "../src/agent-loop.ts"; +import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.ts"; + +class MockAssistantStream extends EventStream { + constructor() { + super( + (event) => event.type === "done" || event.type === "error", + (event) => { + if (event.type === "done") return event.message; + if (event.type === "error") return event.error; + throw new Error("Unexpected event type"); + }, + ); + } +} + +function createUsage() { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createModel(): Model<"openai-responses"> { + return { + id: "mock", + name: "mock", + api: "openai-responses", + provider: "openai", + baseUrl: "https://example.invalid", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 8192, + maxTokens: 2048, + }; +} + +function createAssistantMessage( + content: AssistantMessage["content"], + stopReason: AssistantMessage["stopReason"] = "stop", +): AssistantMessage { + return { + role: "assistant", + content, + api: "openai-responses", + provider: "openai", + model: "mock", + usage: createUsage(), + stopReason, + timestamp: Date.now(), + }; +} + +function createUserMessage(text: string): UserMessage { + return { role: "user", content: text, timestamp: Date.now() }; +} + +const LLM_ROLES = ["user", "assistant", "toolResult"]; + +function identityConverter(messages: AgentMessage[]): Message[] { + return messages.filter((m) => LLM_ROLES.includes(m.role)) as Message[]; +} + +describe("runAgentStep", () => { + it("runs exactly one model call and its tools, then stops", async () => { + const toolSchema = Type.Object({ value: Type.String() }); + const executed: string[] = []; + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_id, params) { + executed.push(params.value); + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + }; + }, + }; + + let callCount = 0; + const streamFn = () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + callCount++; + if (callCount === 1) { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantMessage( + [{ type: "toolCall", id: "t1", name: "echo", arguments: { value: "x" } }], + "toolUse", + ), + }); + } else { + // A whole-turn loop would take this branch; a single step must not. + stream.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "more" }]), + }); + } + }); + return stream; + }; + + const context: AgentContext = { + systemPrompt: "", + messages: [createUserMessage("go")], + tools: [tool], + }; + const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter }; + + const events: AgentEvent[] = []; + const outcome = await runAgentStep( + context, + config, + (event) => { + events.push(event); + }, + undefined, + streamFn, + ); + + // Exactly one model call, and the tool it requested ran once. + expect(callCount).toBe(1); + expect(executed).toEqual(["x"]); + + // One turn boundary. A step is a turn, so it opens no run of its own. + expect(events.filter((e) => e.type === "turn_start").length).toBe(1); + expect(events.filter((e) => e.type === "turn_end").length).toBe(1); + expect(events.filter((e) => e.type === "agent_start").length).toBe(0); + expect(events.filter((e) => e.type === "agent_end").length).toBe(0); + + // The tool result still needs an answer, so the caller has to step again. + expect(outcome.hasMoreToolCalls).toBe(true); + + // The step produced the assistant message and its tool result, and nothing more. + expect(outcome.messages.map((m) => m.role)).toEqual(["assistant", "toolResult"]); + }); + + it("throws when the last message is an assistant", async () => { + const context: AgentContext = { + systemPrompt: "", + messages: [createAssistantMessage([{ type: "text", text: "hi" }])], + tools: [], + }; + await expect( + runAgentStep( + context, + { model: createModel(), convertToLlm: identityConverter }, + () => {}, + undefined, + () => new MockAssistantStream(), + ), + ).rejects.toThrow(/Cannot step from message role: assistant/); + }); +}); diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 2c444d3f468..6aa1f3d4e36 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1129,8 +1129,8 @@ export class AgentSession { * result get one, and a trailing assistant message that holds no answer is dropped. * Returns whether the turn still has work, so false means it already has its answer. * - * resumeInterruptedTurn() is this plus a run to the end of the turn. A caller that drives - * the turn itself wants this one, so recovery stays one step at a time. + * resumeInterruptedTurn() is this plus a run to the end of the turn. A caller driving + * step() itself wants this one, so recovery stays one step at a time. */ prepareStep(): boolean { if (this._isAgentRunActive) { @@ -1197,6 +1197,32 @@ export class AgentSession { // Other roles (bashExecution, compactionSummary, branchSummary) are persisted elsewhere. } + /** + * Advance the current turn by exactly one step (one model call and the tools it + * requests). Unlike prompt(), it adds no message and does not loop; the caller + * drives successive steps. The last recorded message must be a user or tool-result + * message, which is what recordPrompt() and prepareStep() leave behind. + * + * `done` is false while the turn needs another step, which includes a retry or a + * compaction that post-run handling asked for. Steering and follow-up queues stay + * untouched, so the caller decides when a queued message enters the run. + */ + async step(): Promise<{ done: boolean }> { + if (this._isAgentRunActive) { + return { done: false }; + } + + this._isAgentRunActive = true; + try { + const outcome = await this.agent.step(); + // Retries and compaction live here, so a stepped run keeps both. + const needsAnotherPass = await this._handlePostAgentRun(); + return { done: !outcome.hasMoreToolCalls && !needsAnotherPass }; + } finally { + await this._settleRun(); + } + } + private async _handlePostAgentRun(): Promise { const msg = this._lastAssistantMessage; this._lastAssistantMessage = undefined; @@ -1257,6 +1283,40 @@ export class AgentSession { * @throws Error if no model selected or no API key available (when not streaming) */ async prompt(text: string, options?: PromptOptions): Promise { + const messages = await this._buildPromptMessages(text, options); + if (!messages) { + return; + } + + options?.preflightResult?.(true); + await this._runAgentPrompt(messages); + } + + /** + * Record a prompt without running it. Everything prompt() does to build the turn + * happens here (extension input, template expansion, the model and auth checks); the + * model call does not. step() picks the turn up from the transcript. + * + * For a caller that drives a turn one step at a time and checkpoints in between. + * Returns whether a prompt was recorded: an extension command handles its own text, + * and a prompt sent mid-stream is queued instead. + */ + async recordPrompt(text: string, options?: PromptOptions): Promise { + const messages = await this._buildPromptMessages(text, options); + if (!messages) { + return false; + } + + options?.preflightResult?.(true); + this._recordMessages(messages); + return true; + } + + /** + * All of prompt() except the running of it. Returns the turn's messages, or undefined + * when the text needed no run at all. + */ + private async _buildPromptMessages(text: string, options?: PromptOptions): Promise { const expandPromptTemplates = options?.expandPromptTemplates ?? true; const preflightResult = options?.preflightResult; let messages: AgentMessage[] | undefined; @@ -1401,12 +1461,7 @@ export class AgentSession { throw error; } - if (!messages) { - return; - } - - preflightResult?.(true); - await this._runAgentPrompt(messages); + return messages; } /** diff --git a/packages/coding-agent/test/resume-interrupted-turn.test.ts b/packages/coding-agent/test/resume-interrupted-turn.test.ts index 9b6d7903411..b75ef4e6ffc 100644 --- a/packages/coding-agent/test/resume-interrupted-turn.test.ts +++ b/packages/coding-agent/test/resume-interrupted-turn.test.ts @@ -102,7 +102,7 @@ function assertValidToolPairing(messages: AgentMessage[]): void { } } -describe("AgentSession: settling an interrupted turn", () => { +describe("AgentSession: settling and stepping an interrupted turn", () => { let session: AgentSession; let sessionManager: SessionManager; let tempDir: string; @@ -256,4 +256,51 @@ describe("AgentSession: settling an interrupted turn", () => { expect(session.prepareStep()).toBe(true); expect(session.agent.state.messages.filter((m) => m.role === "toolResult").length).toBe(1); }); + + it("records a prompt without calling the model, and step runs it", async () => { + await createSession([answer("answer")]); + + expect(await session.recordPrompt("go")).toBe(true); + expect(modelCalls).toBe(0); + expect(session.agent.state.messages.map((m) => m.role)).toEqual(["user"]); + + expect(await session.step()).toEqual({ done: true }); + expect(modelCalls).toBe(1); + + const messages = session.agent.state.messages; + expect(messages.filter((m) => m.role === "user").length).toBe(1); + const last = messages[messages.length - 1]; + expect(last.role === "assistant" && last.content[0]).toMatchObject({ type: "text", text: "answer" }); + + // The turn is on disk, so the next step can run in another process. + expect(persisted().map((m) => m.role)).toEqual(["user", "assistant"]); + }); + + it("settles an interrupted turn without calling the model, then steps", async () => { + await createSession(); + seed([user("go"), assistant([call("hang-1")], "toolUse")]); + + expect(session.prepareStep()).toBe(true); + expect(modelCalls).toBe(0); + + const settled = session.agent.state.messages.filter((m) => m.role === "toolResult") as ToolResultMessage[]; + expect(settled.map((m) => m.toolCallId)).toEqual(["hang-1"]); + expect(settled[0].isError).toBe(true); + expect(persisted().filter((m) => m.role === "toolResult").length).toBe(1); + + expect(await session.step()).toEqual({ done: true }); + expect(modelCalls).toBe(1); + expect(session.agent.state.messages.filter((m) => m.role === "user").length).toBe(1); + assertValidToolPairing(session.agent.state.messages); + }); + + it("keeps retry handling, so a transient provider error asks for another step", async () => { + const failed = assistant([{ type: "text", text: "" }], "error"); + failed.errorMessage = "overloaded"; + await createSession([failed, answer("answer")]); + session.agent.state.messages = [user("go")]; + + // Post-run handling owns the retry, so the step must not report itself finished. + expect(await session.step()).toEqual({ done: false }); + }); });