diff --git a/packages/server/src/server/agent/providers/acp-agent.ts b/packages/server/src/server/agent/providers/acp-agent.ts index 4944cb9e6..214b2f613 100644 --- a/packages/server/src/server/agent/providers/acp-agent.ts +++ b/packages/server/src/server/agent/providers/acp-agent.ts @@ -86,6 +86,7 @@ import { type ProviderRuntimeSettings, } from "../provider-launch-config.js"; import { renderPromptAttachmentAsText } from "../prompt-attachments.js"; +import { appendOrReplaceGrowingAssistantMessage, runProviderTurn } from "./provider-runner.js"; import { findExecutable } from "../../../utils/executable.js"; import { spawnProcess } from "../../../utils/spawn.js"; @@ -857,84 +858,20 @@ export class ACPAgentSession implements AgentSession, ACPClient { } async run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise { - const timeline: AgentTimelineItem[] = []; - let finalText = ""; - let usage: AgentUsage | undefined; - let turnId: string | null = null; - let settled = false; - let resolveCompletion!: () => void; - let rejectCompletion!: (error: Error) => void; - const buffered: AgentStreamEvent[] = []; - - const completion = new Promise((resolve, reject) => { - resolveCompletion = resolve; - rejectCompletion = reject; + const result = await runProviderTurn({ + prompt, + runOptions: options, + startTurn: (p, o) => this.startTurn(p, o), + subscribe: (callback) => this.subscribe(callback), + getSessionId: () => this.sessionId ?? "", + reduceFinalText: appendOrReplaceGrowingAssistantMessage, }); - const processEvent = (event: AgentStreamEvent) => { - if (settled) { - return; - } - if (turnId && "turnId" in event && event.turnId && event.turnId !== turnId) { - return; - } - if (event.type === "timeline") { - timeline.push(event.item); - if (event.item.type === "assistant_message") { - finalText = event.item.text.startsWith(finalText) - ? event.item.text - : `${finalText}${event.item.text}`; - } - return; - } - if (event.type === "turn_completed") { - usage = event.usage; - settled = true; - resolveCompletion(); - return; - } - if (event.type === "turn_failed") { - settled = true; - rejectCompletion(new Error(event.error)); - return; - } - if (event.type === "turn_canceled") { - settled = true; - resolveCompletion(); - } - }; - - const unsubscribe = this.subscribe((event) => { - if (!turnId) { - buffered.push(event); - return; - } - processEvent(event); - }); - - try { - const started = await this.startTurn(prompt, options); - turnId = started.turnId; - for (const event of buffered) { - processEvent(event); - } - if (!settled) { - await completion; - } - } finally { - unsubscribe(); - } - if (!this.sessionId) { throw new Error("ACP session did not expose a session id"); } - return { - sessionId: this.sessionId, - finalText, - usage, - timeline, - }; + return result; } async startTurn( diff --git a/packages/server/src/server/agent/providers/claude-agent.ts b/packages/server/src/server/agent/providers/claude-agent.ts index 8776881bd..6dbb0059d 100644 --- a/packages/server/src/server/agent/providers/claude-agent.ts +++ b/packages/server/src/server/agent/providers/claude-agent.ts @@ -42,6 +42,7 @@ import { formatProviderDiagnosticError, toDiagnosticErrorMessage, } from "./diagnostic-utils.js"; +import { appendOrReplaceGrowingAssistantMessage, runProviderTurn } from "./provider-runner.js"; import { renderPromptAttachmentAsText } from "../prompt-attachments.js"; import type { @@ -1543,77 +1544,14 @@ class ClaudeAgentSession implements AgentSession { } async run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise { - const timeline: AgentTimelineItem[] = []; - let finalText = ""; - let usage: AgentUsage | undefined; - let turnId: string | null = null; - const bufferedEvents: AgentStreamEvent[] = []; - let settled = false; - let resolveCompletion!: () => void; - let rejectCompletion!: (error: Error) => void; - - const processEvent = (event: AgentStreamEvent) => { - if (settled) { - return; - } - const eventTurnId = (event as { turnId?: string }).turnId; - if (turnId && eventTurnId && eventTurnId !== turnId) { - return; - } - if (event.type === "timeline") { - timeline.push(event.item); - if (event.item.type === "assistant_message") { - if (!finalText) { - finalText = event.item.text; - } else if (event.item.text.startsWith(finalText)) { - finalText = event.item.text; - } else { - finalText += event.item.text; - } - } - return; - } - if (event.type === "turn_completed") { - usage = event.usage; - settled = true; - resolveCompletion(); - return; - } - if (event.type === "turn_failed") { - settled = true; - rejectCompletion(new Error(event.error)); - return; - } - if (event.type === "turn_canceled") { - settled = true; - resolveCompletion(); - } - }; - - const completion = new Promise((resolve, reject) => { - resolveCompletion = resolve; - rejectCompletion = reject; + const result = await runProviderTurn({ + prompt, + runOptions: options, + startTurn: (p, o) => this.startTurn(p, o), + subscribe: (callback) => this.subscribe(callback), + getSessionId: () => this.claudeSessionId ?? "", + reduceFinalText: appendOrReplaceGrowingAssistantMessage, }); - const unsubscribe = this.subscribe((event) => { - if (!turnId) { - bufferedEvents.push(event); - return; - } - processEvent(event); - }); - - try { - const result = await this.startTurn(prompt, options); - turnId = result.turnId; - for (const event of bufferedEvents) { - processEvent(event); - } - if (!settled) { - await completion; - } - } finally { - unsubscribe(); - } this.cachedRuntimeInfo = { provider: "claude", @@ -1626,12 +1564,7 @@ class ClaudeAgentSession implements AgentSession { throw new Error("Session ID not set after run completed"); } - return { - sessionId: this.claudeSessionId, - finalText, - usage, - timeline, - }; + return result; } async startTurn( diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index ace2484d4..a4fc1352c 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.ts @@ -60,6 +60,7 @@ import { resolveBinaryVersion, toDiagnosticErrorMessage, } from "./diagnostic-utils.js"; +import { runProviderTurn } from "./provider-runner.js"; import type { WorkspaceGitService } from "../../workspace-git-service.js"; const DEFAULT_TIMEOUT_MS = 14 * 24 * 60 * 60 * 1000; @@ -2962,81 +2963,22 @@ class CodexAppServerAgentSession implements AgentSession { } async run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise { - const timeline: AgentTimelineItem[] = []; - let finalText = ""; - let usage: AgentUsage | undefined; - let turnId: string | null = null; - const bufferedEvents: AgentStreamEvent[] = []; - let settled = false; - let resolveCompletion!: () => void; - let rejectCompletion!: (error: Error) => void; - - const processEvent = (event: AgentStreamEvent) => { - if (settled) { - return; - } - const eventTurnId = (event as { turnId?: string }).turnId; - if (turnId && eventTurnId && eventTurnId !== turnId) { - return; - } - if (event.type === "timeline") { - timeline.push(event.item); - if (event.item.type === "assistant_message") { - finalText = event.item.text; - } else if (event.item.type === "tool_call" && event.item.detail.type === "plan") { - finalText = event.item.detail.text; + return runProviderTurn({ + prompt, + runOptions: options, + startTurn: (p, o) => this.startTurn(p, o), + subscribe: (callback) => this.subscribe(callback), + getSessionId: async () => (await this.getRuntimeInfo()).sessionId ?? "", + reduceFinalText: ({ current, item }) => { + if (item.type === "assistant_message") { + return item.text; } - return; - } - if (event.type === "turn_completed") { - usage = event.usage; - settled = true; - resolveCompletion(); - return; - } - if (event.type === "turn_failed") { - settled = true; - rejectCompletion(new Error(event.error)); - return; - } - if (event.type === "turn_canceled") { - settled = true; - resolveCompletion(); - } - }; - - const completion = new Promise((resolve, reject) => { - resolveCompletion = resolve; - rejectCompletion = reject; + if (item.type === "tool_call" && item.detail.type === "plan") { + return item.detail.text; + } + return current; + }, }); - const unsubscribe = this.subscribe((event) => { - if (!turnId) { - bufferedEvents.push(event); - return; - } - processEvent(event); - }); - - try { - const result = await this.startTurn(prompt, options); - turnId = result.turnId; - for (const event of bufferedEvents) { - processEvent(event); - } - if (!settled) { - await completion; - } - } finally { - unsubscribe(); - } - - const info = await this.getRuntimeInfo(); - return { - sessionId: info.sessionId ?? "", - finalText, - usage, - timeline, - }; } async startTurn( diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index d6050ba89..6c84936b1 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -57,6 +57,7 @@ import { resolveBinaryVersion, toDiagnosticErrorMessage, } from "./diagnostic-utils.js"; +import { runProviderTurn } from "./provider-runner.js"; import { renderPromptAttachmentAsText } from "../prompt-attachments.js"; const OPENCODE_CAPABILITIES: AgentCapabilityFlags = { @@ -2271,78 +2272,13 @@ class OpenCodeAgentSession implements AgentSession { } async run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise { - const timeline: AgentTimelineItem[] = []; - let finalText = ""; - let usage: AgentUsage | undefined; - let turnId: string | null = null; - const bufferedEvents: AgentStreamEvent[] = []; - let settled = false; - let resolveCompletion!: () => void; - let rejectCompletion!: (error: Error) => void; - - const processEvent = (event: AgentStreamEvent) => { - if (settled) { - return; - } - const eventTurnId = (event as { turnId?: string }).turnId; - if (turnId && eventTurnId && eventTurnId !== turnId) { - return; - } - if (event.type === "timeline") { - timeline.push(event.item); - if (event.item.type === "assistant_message") { - finalText = event.item.text; - } - return; - } - if (event.type === "turn_completed") { - usage = event.usage; - settled = true; - resolveCompletion(); - return; - } - if (event.type === "turn_failed") { - settled = true; - rejectCompletion(new Error(event.error)); - return; - } - if (event.type === "turn_canceled") { - settled = true; - resolveCompletion(); - } - }; - - const completion = new Promise((resolve, reject) => { - resolveCompletion = resolve; - rejectCompletion = reject; + return runProviderTurn({ + prompt, + runOptions: options, + startTurn: (p, o) => this.startTurn(p, o), + subscribe: (callback) => this.subscribe(callback), + getSessionId: () => this.sessionId, }); - const unsubscribe = this.subscribe((event) => { - if (!turnId) { - bufferedEvents.push(event); - return; - } - processEvent(event); - }); - - try { - const result = await this.startTurn(prompt, options); - turnId = result.turnId; - for (const event of bufferedEvents) { - processEvent(event); - } - if (!settled) { - await completion; - } - } finally { - unsubscribe(); - } - - return { - sessionId: this.sessionId, - finalText, - usage, - timeline, - }; } async interrupt(): Promise { diff --git a/packages/server/src/server/agent/providers/provider-runner.test.ts b/packages/server/src/server/agent/providers/provider-runner.test.ts new file mode 100644 index 000000000..a2b646507 --- /dev/null +++ b/packages/server/src/server/agent/providers/provider-runner.test.ts @@ -0,0 +1,154 @@ +import { describe, expect, test } from "vitest"; +import type { AgentPromptInput, AgentRunOptions, AgentStreamEvent } from "../agent-sdk-types.js"; +import { appendOrReplaceGrowingAssistantMessage, runProviderTurn } from "./provider-runner.js"; + +class FakeTurnRunner { + readonly events: AgentStreamEvent[] = []; + private subscriber: ((event: AgentStreamEvent) => void) | null = null; + + constructor( + private readonly turnId: string, + private readonly sessionId: string, + private readonly onStart?: () => void, + ) {} + + async startTurn( + _prompt: AgentPromptInput, + _options?: AgentRunOptions, + ): Promise<{ turnId: string }> { + this.onStart?.(); + return { turnId: this.turnId }; + } + + subscribe(callback: (event: AgentStreamEvent) => void): () => void { + this.subscriber = callback; + return () => { + this.subscriber = null; + }; + } + + getSessionId(): string { + return this.sessionId; + } + + emit(event: AgentStreamEvent): void { + this.events.push(event); + this.subscriber?.(event); + } +} + +describe("runProviderTurn", () => { + test("buffers events emitted before startTurn returns", async () => { + const runner = new FakeTurnRunner("turn-1", "session-1", () => { + runner.emit({ + type: "timeline", + provider: "codex", + turnId: "turn-1", + item: { type: "assistant_message", text: "hello" }, + }); + runner.emit({ + type: "turn_completed", + provider: "codex", + turnId: "turn-1", + usage: { inputTokens: 1 }, + }); + }); + + const result = await runProviderTurn({ + prompt: "hi", + startTurn: (prompt, options) => runner.startTurn(prompt, options), + subscribe: (callback) => runner.subscribe(callback), + getSessionId: () => runner.getSessionId(), + }); + + expect(result).toEqual({ + sessionId: "session-1", + finalText: "hello", + usage: { inputTokens: 1 }, + timeline: [{ type: "assistant_message", text: "hello" }], + }); + }); + + test("ignores events from other turns after the turn id is known", async () => { + const runner = new FakeTurnRunner("turn-1", "session-1"); + const resultPromise = runProviderTurn({ + prompt: "hi", + startTurn: (prompt, options) => runner.startTurn(prompt, options), + subscribe: (callback) => runner.subscribe(callback), + getSessionId: () => runner.getSessionId(), + }); + + runner.emit({ + type: "timeline", + provider: "opencode", + turnId: "other-turn", + item: { type: "assistant_message", text: "wrong" }, + }); + runner.emit({ + type: "timeline", + provider: "opencode", + turnId: "turn-1", + item: { type: "assistant_message", text: "right" }, + }); + runner.emit({ type: "turn_completed", provider: "opencode", turnId: "turn-1" }); + + await expect(resultPromise).resolves.toMatchObject({ + finalText: "right", + timeline: [{ type: "assistant_message", text: "right" }], + }); + }); + + test("rejects when the turn fails", async () => { + const runner = new FakeTurnRunner("turn-1", "session-1"); + const resultPromise = runProviderTurn({ + prompt: "hi", + startTurn: (prompt, options) => runner.startTurn(prompt, options), + subscribe: (callback) => runner.subscribe(callback), + getSessionId: () => runner.getSessionId(), + }); + + runner.emit({ + type: "turn_failed", + provider: "claude", + turnId: "turn-1", + error: "provider failed", + }); + + await expect(resultPromise).rejects.toThrow("provider failed"); + }); + + test("supports growing assistant message reducers", async () => { + const runner = new FakeTurnRunner("turn-1", "session-1"); + const resultPromise = runProviderTurn({ + prompt: "hi", + startTurn: (prompt, options) => runner.startTurn(prompt, options), + subscribe: (callback) => runner.subscribe(callback), + getSessionId: () => runner.getSessionId(), + reduceFinalText: appendOrReplaceGrowingAssistantMessage, + }); + + runner.emit({ + type: "timeline", + provider: "acp", + turnId: "turn-1", + item: { type: "assistant_message", text: "hello" }, + }); + runner.emit({ + type: "timeline", + provider: "acp", + turnId: "turn-1", + item: { type: "assistant_message", text: "hello world" }, + }); + runner.emit({ + type: "timeline", + provider: "acp", + turnId: "turn-1", + item: { type: "assistant_message", text: "!" }, + }); + runner.emit({ type: "turn_completed", provider: "acp", turnId: "turn-1" }); + + await expect(resultPromise).resolves.toMatchObject({ + finalText: "hello world!", + }); + }); +}); diff --git a/packages/server/src/server/agent/providers/provider-runner.ts b/packages/server/src/server/agent/providers/provider-runner.ts new file mode 100644 index 000000000..d13d171b6 --- /dev/null +++ b/packages/server/src/server/agent/providers/provider-runner.ts @@ -0,0 +1,128 @@ +import type { + AgentPromptInput, + AgentRunOptions, + AgentRunResult, + AgentStreamEvent, + AgentTimelineItem, +} from "../agent-sdk-types.js"; + +export type ProviderFinalTextReducer = (params: { + current: string; + item: AgentTimelineItem; +}) => string; + +export interface ProviderTurnRunner { + startTurn(prompt: AgentPromptInput, options?: AgentRunOptions): Promise<{ turnId: string }>; + subscribe(callback: (event: AgentStreamEvent) => void): () => void; + getSessionId(): string | Promise; +} + +export interface RunProviderTurnOptions extends ProviderTurnRunner { + prompt: AgentPromptInput; + runOptions?: AgentRunOptions; + reduceFinalText?: ProviderFinalTextReducer; +} + +export async function runProviderTurn({ + prompt, + runOptions, + startTurn, + subscribe, + getSessionId, + reduceFinalText = replaceFinalTextWithAssistantMessage, +}: RunProviderTurnOptions): Promise { + const timeline: AgentTimelineItem[] = []; + let finalText = ""; + let usage: AgentRunResult["usage"]; + let turnId: string | null = null; + const bufferedEvents: AgentStreamEvent[] = []; + let settled = false; + let resolveCompletion!: () => void; + let rejectCompletion!: (error: Error) => void; + + const processEvent = (event: AgentStreamEvent) => { + if (settled) { + return; + } + const eventTurnId = "turnId" in event ? event.turnId : undefined; + if (turnId && eventTurnId && eventTurnId !== turnId) { + return; + } + if (event.type === "timeline") { + timeline.push(event.item); + finalText = reduceFinalText({ current: finalText, item: event.item }); + return; + } + if (event.type === "turn_completed") { + usage = event.usage; + settled = true; + resolveCompletion(); + return; + } + if (event.type === "turn_failed") { + settled = true; + rejectCompletion(new Error(event.error)); + return; + } + if (event.type === "turn_canceled") { + settled = true; + resolveCompletion(); + } + }; + + const completion = new Promise((resolve, reject) => { + resolveCompletion = resolve; + rejectCompletion = reject; + }); + const unsubscribe = subscribe((event) => { + if (!turnId) { + bufferedEvents.push(event); + return; + } + processEvent(event); + }); + + try { + const result = await startTurn(prompt, runOptions); + turnId = result.turnId; + for (const event of bufferedEvents) { + processEvent(event); + } + await completion; + } finally { + unsubscribe(); + } + + return { + sessionId: await getSessionId(), + finalText, + usage, + timeline, + }; +} + +export function replaceFinalTextWithAssistantMessage({ + current, + item, +}: { + current: string; + item: AgentTimelineItem; +}): string { + return item.type === "assistant_message" ? item.text : current; +} + +export function appendOrReplaceGrowingAssistantMessage({ + current, + item, +}: { + current: string; + item: AgentTimelineItem; +}): string { + if (item.type !== "assistant_message") { + return current; + } + if (!current) { + return item.text; + } + return item.text.startsWith(current) ? item.text : `${current}${item.text}`; +}