diff --git a/packages/server/src/server/agent/providers/acp-agent.test.ts b/packages/server/src/server/agent/providers/acp-agent.test.ts index 07e392793..49adf40a2 100644 --- a/packages/server/src/server/agent/providers/acp-agent.test.ts +++ b/packages/server/src/server/agent/providers/acp-agent.test.ts @@ -2150,6 +2150,14 @@ describe("ACPAgentSession", () => { content: { type: "text", text: "lo" }, } as SessionUpdate, }); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "agent_message_chunk", + messageId: "assistant-2", + content: { type: "text", text: "" }, + } as SessionUpdate, + }); const timeline = events .filter((event) => event.type === "timeline") @@ -2161,7 +2169,6 @@ describe("ACPAgentSession", () => { { type: "assistant_message", text: " How are you?", messageId: "assistant-1" }, { type: "reasoning", text: "Thinking" }, { type: "reasoning", text: " more" }, - { type: "user_message", text: "hel", messageId: "user-1" }, { type: "user_message", text: "hello", messageId: "user-1" }, ]); }); @@ -2528,6 +2535,102 @@ describe("ACPAgentSession", () => { ]); }); + test("startTurn dedupes a provider-owned user echo streamed as text and image chunks", async () => { + const session = createSession(); + const events: AgentStreamEvent[] = []; + const prompt = vi.fn(() => new Promise(() => {})); + + asInternals(session).sessionId = "session-1"; + asInternals(session).connection = { prompt }; + + session.subscribe((event) => { + events.push(event); + }); + + await session.startTurn( + [ + { type: "text", text: "hey" }, + { type: "image", data: "AA==", mimeType: "image/png" }, + ], + { clientMessageId: "msg-client-1" }, + ); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + messageId: "msg-provider-1", + content: { type: "text", text: "hey" }, + } as SessionUpdate, + }); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + messageId: "msg-provider-1", + content: { type: "image", data: "AA==", mimeType: "image/png" }, + } as SessionUpdate, + }); + + expect( + events.filter((event) => event.type === "timeline" && event.item.type === "user_message"), + ).toEqual([ + { + type: "timeline", + provider: session.provider, + item: { + type: "user_message", + text: "hey", + messageId: "msg-client-1", + clientMessageId: "msg-client-1", + }, + turnId: expect.any(String), + }, + ]); + }); + + test("startTurn keeps an image-only provider echo when no canonical user message was emitted", async () => { + const session = createSession(); + const events: AgentStreamEvent[] = []; + const prompt = vi.fn(() => new Promise(() => {})); + + asInternals(session).sessionId = "session-1"; + asInternals(session).connection = { prompt }; + + session.subscribe((event) => { + events.push(event); + }); + + await session.startTurn([{ type: "image", data: "AA==", mimeType: "image/png" }], { + clientMessageId: "msg-client-1", + }); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + content: { type: "image", data: "AA==", mimeType: "image/png" }, + } as SessionUpdate, + }); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "agent_message_chunk", + messageId: "assistant-1", + content: { type: "text", text: "I see it" }, + } as SessionUpdate, + }); + + expect( + events.filter((event) => event.type === "timeline" && event.item.type === "user_message"), + ).toEqual([ + { + type: "timeline", + provider: session.provider, + item: { type: "user_message", text: "[image]" }, + turnId: expect.any(String), + }, + ]); + }); + test("startTurn converts background prompt rejections into turn_failed events", async () => { const session = createSession(); const events: Array<{ type: string; turnId?: string; error?: string }> = []; @@ -2561,6 +2664,49 @@ describe("ACPAgentSession", () => { expect(asInternals(session).activeForegroundTurnId).toBeNull(); }); + test("flushes an image-only provider echo before a rejected turn finishes", async () => { + const session = createSession(); + const events: AgentStreamEvent[] = []; + let rejectPrompt!: (error: Error) => void; + const prompt = vi.fn( + () => + new Promise((_, reject) => { + rejectPrompt = reject; + }), + ); + + asInternals(session).sessionId = "session-1"; + asInternals(session).connection = { prompt }; + session.subscribe((event) => events.push(event)); + + const { turnId } = await session.startTurn([ + { type: "image", data: "AA==", mimeType: "image/png" }, + ]); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + content: { type: "image", data: "AA==", mimeType: "image/png" }, + } as SessionUpdate, + }); + + rejectPrompt(new Error("prompt failed")); + await Promise.resolve(); + await Promise.resolve(); + + expect( + events.filter((event) => event.type === "timeline" || event.type === "turn_failed"), + ).toEqual([ + { + type: "timeline", + provider: session.provider, + item: { type: "user_message", text: "[image]" }, + turnId, + }, + expect.objectContaining({ type: "turn_failed", turnId, error: "prompt failed" }), + ]); + }); + test("startTurn preserves JSON-RPC error details from a real ACP prompt response", async () => { const session = createSession(); const clientToAgent = new TransformStream(); @@ -2663,6 +2809,32 @@ describe("ACPAgentSession close() tree-kill", () => { vi.restoreAllMocks(); }); + test("close() flushes a buffered provider-owned user message before unsubscribing", async () => { + const session = createSession(); + const events: AgentStreamEvent[] = []; + session.subscribe((event) => events.push(event)); + asInternals(session).sessionId = "session-1"; + + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + content: { type: "image", data: "AA==", mimeType: "image/png" }, + } as SessionUpdate, + }); + expect(events).toEqual([]); + + await session.close(); + + expect(events).toEqual([ + { + type: "timeline", + provider: session.provider, + item: { type: "user_message", text: "[image]" }, + }, + ]); + }); + test("close() terminates the main child process via the process tree", async () => { const terminator = new FakeTerminator(); const session = createSession(terminator.terminate); @@ -2984,6 +3156,68 @@ describe("ACP session/load invariant — cwd and mcpServers always passed", () = ]); }); + test("coalesces an ID-less text and image user message during loadSession replay", async () => { + let session!: ACPAgentSession; + const loadSession = async () => { + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "hey" }, + } as SessionUpdate, + }); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "user_message_chunk", + content: { type: "image", data: "AA==", mimeType: "image/png" }, + } as SessionUpdate, + }); + await session.sessionUpdate({ + sessionId: "session-1", + update: { + sessionUpdate: "agent_message_chunk", + messageId: "assistant-replay-1", + content: { type: "text", text: "Hello" }, + } as SessionUpdate, + }); + return { + sessionId: "session-1", + modes: null, + models: null, + configOptions: [], + }; + }; + ({ session } = makeTestSession({ + capabilities: { loadSession: true }, + handle: { sessionId: "session-1", provider: "test-acp" }, + loadSession, + })); + + await session.initializeResumedSession(); + + const history: AgentStreamEvent[] = []; + for await (const event of session.streamHistory()) { + history.push(event); + } + expect(history).toEqual([ + { + type: "timeline", + provider: session.provider, + item: { type: "user_message", text: "hey[image]" }, + }, + { + type: "timeline", + provider: session.provider, + item: { + type: "assistant_message", + text: "Hello", + messageId: "assistant-replay-1", + }, + }, + ]); + }); + test("assigns stable fallback IDs to ID-less assistant messages during loadSession replay", async () => { let session!: ACPAgentSession; const loadSession = async () => { diff --git a/packages/server/src/server/agent/providers/acp-agent.ts b/packages/server/src/server/agent/providers/acp-agent.ts index 033e21238..399ed8942 100644 --- a/packages/server/src/server/agent/providers/acp-agent.ts +++ b/packages/server/src/server/agent/providers/acp-agent.ts @@ -467,14 +467,9 @@ interface PendingPermission { turnId: string | null; } -interface MessageAssemblyState { +interface PendingUserMessage { text: string; -} - -interface SubmittedUserMessageEcho { - messageId: string; - text: string; - turnId: string; + messageId?: string; } export type SessionStateResponse = NewSessionResponse | LoadSessionResponse | ResumeSessionResponse; @@ -1302,9 +1297,8 @@ export class ACPAgentSession implements AgentSession, ACPClient { private readonly launchEnv?: Record; private readonly subscribers = new Set<(event: AgentStreamEvent) => void>(); private readonly pendingPermissions = new Map(); - private readonly messageAssemblies = new Map(); - private readonly submittedUserMessageIds = new Set(); - private activeSubmittedUserMessage: SubmittedUserMessageEcho | null = null; + private pendingUserMessage: PendingUserMessage | null = null; + private submittedUserMessageTurnId: string | null = null; private readonly toolCalls = new Map(); private readonly terminalEntries = new Map(); private readonly persistedHistory: AgentTimelineItem[] = []; @@ -1428,6 +1422,7 @@ export class ACPAgentSession implements AgentSession, ACPClient { mcpServers: this.acpMcpServers(), }), ); + this.deliverTranslatedEvents(this.flushPendingUserMessage()); this.replayingHistory = false; this.historyPending = this.persistedHistory.length > 0; this.applySessionState(response); @@ -1493,11 +1488,12 @@ export class ACPAgentSession implements AgentSession, ACPClient { throw new Error("A foreground turn is already active"); } + this.deliverTranslatedEvents(this.flushPendingUserMessage()); const turnId = randomUUID(); const messageId = options?.clientMessageId ?? randomUUID(); this.activeForegroundTurnId = turnId; this.fallbackAssistantMessageId = null; - this.activeSubmittedUserMessage = null; + this.submittedUserMessageTurnId = null; this.emitBootstrapThreadEvent(); this.pushEvent({ type: "turn_started", provider: this.provider, turnId }); this.emitSubmittedUserMessage(prompt, messageId, turnId, options?.clientMessageId); @@ -2043,6 +2039,7 @@ export class ACPAgentSession implements AgentSession, ACPClient { } this.closed = true; + this.deliverTranslatedEvents(this.flushPendingUserMessage()); this.settleCommandsReady(); for (const pending of this.pendingPermissions.values()) { @@ -2141,6 +2138,10 @@ export class ACPAgentSession implements AgentSession, ACPClient { }, "provider.acp.parsed_event", ); + this.deliverTranslatedEvents(events); + } + + private deliverTranslatedEvents(events: AgentStreamEvent[]): void { if (this.replayingHistory) { for (const event of events) { if (event.type === "timeline") { @@ -2467,45 +2468,43 @@ export class ACPAgentSession implements AgentSession, ACPClient { } private translateSessionUpdate(update: SessionUpdate): AgentStreamEvent[] { + if (update.sessionUpdate === "user_message_chunk") { + return this.handleUserMessageChunk(update); + } + + const pendingUserEvents = this.flushPendingUserMessage(); switch (update.sessionUpdate) { - case "user_message_chunk": { - this.fallbackAssistantMessageId = null; - const item = this.createMessageTimelineItem("user_message", update); - if (!item) { - return []; - } - if (item.type !== "user_message") { - return [this.wrapTimeline(item)]; - } - if (this.isSubmittedUserMessageEcho(item)) { - return []; - } - return [this.wrapTimeline(item)]; - } case "agent_message_chunk": { const item = this.createMessageTimelineItem("assistant_message", update); - return item ? [this.wrapTimeline(item)] : []; + return item ? [...pendingUserEvents, this.wrapTimeline(item)] : pendingUserEvents; } case "agent_thought_chunk": { this.fallbackAssistantMessageId = null; const item = this.createMessageTimelineItem("reasoning", update); - return item ? [this.wrapTimeline(item)] : []; + return item ? [...pendingUserEvents, this.wrapTimeline(item)] : pendingUserEvents; } case "tool_call": this.fallbackAssistantMessageId = null; - return this.handleToolCallUpdate(update.toolCallId, update, undefined); + return [ + ...pendingUserEvents, + ...this.handleToolCallUpdate(update.toolCallId, update, undefined), + ]; case "tool_call_update": - return this.handleToolCallUpdate( - update.toolCallId, - update, - this.toolCalls.get(update.toolCallId), - ); + return [ + ...pendingUserEvents, + ...this.handleToolCallUpdate( + update.toolCallId, + update, + this.toolCalls.get(update.toolCallId), + ), + ]; case "plan": this.fallbackAssistantMessageId = null; - return [this.wrapTimeline(mapPlanToTimeline(update))]; + return [...pendingUserEvents, this.wrapTimeline(mapPlanToTimeline(update))]; case "current_mode_update": this.handleCurrentModeUpdate(update); return [ + ...pendingUserEvents, { type: "mode_changed", provider: this.provider, @@ -2514,13 +2513,13 @@ export class ACPAgentSession implements AgentSession, ACPClient { }, ]; case "config_option_update": - return this.handleConfigOptionUpdate(update); + return [...pendingUserEvents, ...this.handleConfigOptionUpdate(update)]; case "session_info_update": this.handleSessionInfoUpdate(update); - return []; + return pendingUserEvents; case "usage_update": this.handleUsageUpdate(update); - return []; + return pendingUserEvents; case "available_commands_update": this.cachedCommands = update.availableCommands.map((command) => ({ name: command.name, @@ -2529,12 +2528,60 @@ export class ACPAgentSession implements AgentSession, ACPClient { kind: "command", })); this.settleCommandsReady(); - return []; + return pendingUserEvents; default: - return []; + return pendingUserEvents; } } + private handleUserMessageChunk( + update: Extract, + ): AgentStreamEvent[] { + this.fallbackAssistantMessageId = null; + if ( + this.activeForegroundTurnId && + this.submittedUserMessageTurnId === this.activeForegroundTurnId + ) { + return []; + } + + const chunkText = contentBlockToText(update.content); + if (!chunkText) { + return []; + } + + const messageId = update.messageId ?? undefined; + const pending = this.pendingUserMessage; + const startsNewMessage = Boolean( + pending?.messageId && messageId && pending.messageId !== messageId, + ); + const events = startsNewMessage ? this.flushPendingUserMessage() : []; + this.pendingUserMessage ??= { + text: "", + ...(messageId ? { messageId } : {}), + }; + if (!this.pendingUserMessage.messageId && messageId) { + this.pendingUserMessage.messageId = messageId; + } + this.pendingUserMessage.text += chunkText; + return events; + } + + private flushPendingUserMessage(): AgentStreamEvent[] { + const pending = this.pendingUserMessage; + if (!pending) { + return []; + } + this.pendingUserMessage = null; + return [ + this.wrapTimeline({ + type: "user_message", + text: pending.text, + ...(pending.messageId ? { messageId: pending.messageId } : {}), + }), + ]; + } + private handleToolCallUpdate( toolCallId: string, update: ToolCall | ToolCallUpdate, @@ -2549,13 +2596,12 @@ export class ACPAgentSession implements AgentSession, ACPClient { } private createMessageTimelineItem( - type: "user_message" | "assistant_message" | "reasoning", + type: "assistant_message" | "reasoning", update: Extract< SessionUpdate, - { sessionUpdate: "user_message_chunk" | "agent_message_chunk" | "agent_thought_chunk" } + { sessionUpdate: "agent_message_chunk" | "agent_thought_chunk" } >, ): - | { type: "user_message"; text: string; messageId?: string } | { type: "assistant_message"; text: string; messageId: string } | { type: "reasoning"; text: string } | null { @@ -2563,14 +2609,6 @@ export class ACPAgentSession implements AgentSession, ACPClient { if (!chunkText) { return null; } - const key = this.messageAssemblyKey(type, update.messageId); - const state = this.messageAssemblies.get(key) ?? { text: "" }; - state.text += chunkText; - this.messageAssemblies.set(key, state); - - if (type === "user_message") { - return { type: "user_message", text: state.text, messageId: update.messageId ?? undefined }; - } if (type === "assistant_message") { return { type: "assistant_message", @@ -2590,15 +2628,6 @@ export class ACPAgentSession implements AgentSession, ACPClient { return this.fallbackAssistantMessageId; } - private messageAssemblyKey( - type: "user_message" | "assistant_message" | "reasoning", - messageId: string | null | undefined, - ): string { - const fallbackId = - type === "user_message" ? (this.activeForegroundTurnId ?? "default") : "default"; - return `${type}:${messageId ?? fallbackId}`; - } - private handleCurrentModeUpdate(update: CurrentModeUpdate): void { this.currentMode = this.transformModeId(update.currentModeId); } @@ -2717,8 +2746,7 @@ export class ACPAgentSession implements AgentSession, ACPClient { if (text.trim().length === 0) { return; } - this.submittedUserMessageIds.add(messageId); - this.activeSubmittedUserMessage = { messageId, text, turnId }; + this.submittedUserMessageTurnId = turnId; this.pushEvent({ type: "timeline", provider: this.provider, @@ -2749,29 +2777,15 @@ export class ACPAgentSession implements AgentSession, ACPClient { private finishTurn( event: Extract, ): void { + this.deliverTranslatedEvents(this.flushPendingUserMessage()); this.activeForegroundTurnId = null; this.fallbackAssistantMessageId = null; - if (this.activeSubmittedUserMessage?.turnId === event.turnId) { - this.activeSubmittedUserMessage = null; + if (this.submittedUserMessageTurnId === event.turnId) { + this.submittedUserMessageTurnId = null; } this.pushEvent(event); } - private isSubmittedUserMessageEcho( - item: Extract, - ): boolean { - const active = this.activeSubmittedUserMessage; - if (!active || active.turnId !== this.activeForegroundTurnId) { - return false; - } - if (item.messageId) { - if (this.submittedUserMessageIds.has(item.messageId)) { - return true; - } - } - return active.text.startsWith(item.text); - } - private emitBootstrapThreadEvent(): void { if (!this.bootstrapThreadEventPending || !this.sessionId) { return;