diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index 89f5c2c15c..bc43a7858b 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -25,6 +25,7 @@ import { it, vi, } from "vitest"; +import { POSTHOG_NOTIFICATIONS } from "../acp-extensions"; import { getSessionJsonlPath } from "../adapters/claude/session/jsonl-hydration"; import type { PermissionMode } from "../execution-mode"; import type { PostHogAPIClient } from "../posthog-api"; @@ -1433,6 +1434,126 @@ describe("AgentServer HTTP Mode", () => { expect(body.error).toBe("No active session for this run"); }, 20000); + it("continues a cloud task after a manual compact command", async () => { + const s = createServer(); + await s.start(); + const broadcastEvent = vi.fn(); + let serverInternals!: { + session: { clientConnection: { prompt: typeof prompt } }; + broadcastEvent: typeof broadcastEvent; + handleAcpTransportMessage(message: unknown): void; + }; + const prompt = vi.fn(async (_params: { prompt: ContentBlock[] }) => { + serverInternals.handleAcpTransportMessage({ + jsonrpc: "2.0", + method: POSTHOG_NOTIFICATIONS.TURN_COMPLETE, + params: { sessionId: "session-1", stopReason: "end_turn" }, + }); + return { stopReason: "end_turn" }; + }); + serverInternals = s as unknown as typeof serverInternals; + serverInternals.session.clientConnection.prompt = prompt; + serverInternals.broadcastEvent = broadcastEvent; + + const token = createToken(); + const response = await fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "compact-and-continue", + method: "user_message", + params: { + content: + "/compact Continue with the task using the question tool and plan.", + }, + }), + }); + + expect(response.status).toBe(200); + const body = (await response.json()) as { + result?: { stopReason?: string }; + }; + expect(body.result?.stopReason).toBe("end_turn"); + expect(prompt).toHaveBeenCalledTimes(2); + expect(prompt.mock.calls[0]?.[0].prompt).toEqual([ + { + type: "text", + text: "/compact Continue with the task using the question tool and plan.", + }, + ]); + expect(prompt.mock.calls[1]?.[0].prompt).toEqual([ + { + type: "text", + text: expect.stringContaining("Continue working on the task"), + _meta: { ui: { hidden: true } }, + }, + ]); + const turnCompleteEvents = broadcastEvent.mock.calls.filter( + ([event]) => + (event as { notification?: { method?: string } }).notification + ?.method === POSTHOG_NOTIFICATIONS.TURN_COMPLETE, + ); + expect(turnCompleteEvents).toHaveLength(1); + }, 20000); + + it("retries only the continuation after compact follow-up failure", async () => { + const s = createServer(); + await s.start(); + const prompt = vi + .fn(async (_params: { prompt: ContentBlock[] }) => ({ + stopReason: "end_turn", + })) + .mockResolvedValueOnce({ stopReason: "end_turn" }) + .mockRejectedValueOnce(new Error("sdk connection lost")); + const serverInternals = s as unknown as { + session: { clientConnection: { prompt: typeof prompt } }; + }; + serverInternals.session.clientConnection.prompt = prompt; + + const token = createToken(); + const send = async () => + fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "compact-retry", + method: "user_message", + params: { + content: "/compact Continue the task.", + messageId: "compact-retry", + }, + }), + }); + + const first = await send(); + expect(first.status).toBe(200); + expect(prompt).toHaveBeenCalledTimes(2); + + const retry = await send(); + expect(retry.status).toBe(200); + expect(prompt).toHaveBeenCalledTimes(3); + expect(prompt.mock.calls[0]?.[0].prompt[0]).toMatchObject({ + type: "text", + text: "/compact Continue the task.", + }); + expect(prompt.mock.calls[1]?.[0].prompt[0]).toMatchObject({ + type: "text", + text: expect.stringContaining("Continue working on the task"), + }); + expect(prompt.mock.calls[2]?.[0].prompt[0]).toMatchObject({ + type: "text", + text: expect.stringContaining("Continue working on the task"), + }); + }, 20000); + it("rewrites a bundled local skill slash command before sending the prompt", async () => { const skillDefinition = [ "---", diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 1c14af21ae..5210b7540f 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -289,6 +289,10 @@ function hiddenTextBlock(text: string): ContentBlock { } as ContentBlock; } +function isManualCompactPrompt(prompt: ContentBlock[]): boolean { + return /^\/compact(?:\s|$)/.test(promptBlocksToText(prompt).trimStart()); +} + interface LocalSkillPromptContext { /** Set when the message is a bare `/skill` invocation the adapter should strip. */ skillName?: string; @@ -339,6 +343,7 @@ export class AgentServer { private rtkSavingsAttempted = false; private questionRelayedToSlack = false; private adapterEmittedTurnComplete = false; + private suppressAdapterTurnComplete = false; private runUsage = new RunUsageAccumulator(); private detectedPrUrl: string | null = null; // Reset per session. `evaluatedPrUrls` dedupes per URL; `prAttributionChain` serializes @@ -364,6 +369,7 @@ export class AgentServer { private initializationPromise: Promise | null = null; private pendingEvents: Record[] = []; private deliveredMessageIds = new Set(); + private pendingCompactContinuationMessageIds = new Set(); private pendingPermissions = new Map< string, { @@ -909,18 +915,28 @@ export class AgentServer { typeof params.messageId === "string" && params.messageId ? params.messageId : undefined; + let retryCompactContinuation = false; if (messageId) { if (this.deliveredMessageIds.has(messageId)) { - this.logger.info("Duplicate user_message delivery ignored", { - messageId, - }); - return { stopReason: "duplicate_delivery", duplicate: true }; + if (this.pendingCompactContinuationMessageIds.has(messageId)) { + retryCompactContinuation = true; + this.logger.info("Retrying pending compact continuation", { + messageId, + }); + } else { + this.logger.info("Duplicate user_message delivery ignored", { + messageId, + }); + return { stopReason: "duplicate_delivery", duplicate: true }; + } + } else { + this.deliveredMessageIds.add(messageId); } - this.deliveredMessageIds.add(messageId); if (this.deliveredMessageIds.size > 500) { const oldest = this.deliveredMessageIds.values().next().value; if (oldest !== undefined) { this.deliveredMessageIds.delete(oldest); + this.pendingCompactContinuationMessageIds.delete(oldest); } } } @@ -951,17 +967,53 @@ export class AgentServer { : {}), }; + const manualCompactPrompt = isManualCompactPrompt(prompt); + const acpSessionId = this.session.acpSessionId; + const continueAfterCompaction = (): Promise => + this.promptWithUpstreamRetry({ + sessionId: acpSessionId, + prompt: [ + hiddenTextBlock( + "Compaction is complete. Continue working on the task from the compacted context, following the user's instructions from the /compact command.", + ), + ], + }); + + let compactCommandCompleted = retryCompactContinuation; let result: PromptResponse; + this.suppressAdapterTurnComplete = + manualCompactPrompt || retryCompactContinuation; try { - result = await this.session.clientConnection.prompt({ - sessionId: this.session.acpSessionId, - prompt, - ...(Object.keys(promptMeta).length > 0 - ? { _meta: promptMeta } - : {}), - }); + if (retryCompactContinuation) { + result = await continueAfterCompaction(); + if (messageId) { + this.pendingCompactContinuationMessageIds.delete(messageId); + } + } else { + result = await this.session.clientConnection.prompt({ + sessionId: this.session.acpSessionId, + prompt, + ...(Object.keys(promptMeta).length > 0 + ? { _meta: promptMeta } + : {}), + }); + + if (result.stopReason === "end_turn" && manualCompactPrompt) { + compactCommandCompleted = true; + if (messageId) { + this.pendingCompactContinuationMessageIds.add(messageId); + } + // `/compact` is an SDK-local command, so without a follow-up the + // cloud run reports completion before the model resumes the task. + this.recordTurnUsage(result.usage); + result = await continueAfterCompaction(); + if (messageId) { + this.pendingCompactContinuationMessageIds.delete(messageId); + } + } + } } catch (error) { - if (messageId) { + if (messageId && !compactCommandCompleted) { this.deliveredMessageIds.delete(messageId); } await this.session.logWriter.flushAll(); @@ -974,6 +1026,8 @@ export class AgentServer { throw error; } return { stopReason: "error_recoverable" }; + } finally { + this.suppressAdapterTurnComplete = false; } this.logger.debug("User message completed", { @@ -1299,16 +1353,8 @@ export class AgentServer { // Tap both streams to broadcast all ACP messages via SSE (mimics local transport) this.adapterEmittedTurnComplete = false; - const onAcpMessage = (message: unknown) => { - if (isTurnCompleteNotification(message)) { - this.adapterEmittedTurnComplete = true; - } - this.broadcastEvent({ - type: "notification", - timestamp: new Date().toISOString(), - notification: message, - }); - }; + const onAcpMessage = (message: unknown) => + this.handleAcpTransportMessage(message); const tappedReadable = createTappedReadableStream( acpConnection.clientStreams.readable as ReadableStream, @@ -4045,6 +4091,20 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} }); } + private handleAcpTransportMessage(message: unknown): void { + if (isTurnCompleteNotification(message)) { + if (this.suppressAdapterTurnComplete) { + return; + } + this.adapterEmittedTurnComplete = true; + } + this.broadcastEvent({ + type: "notification", + timestamp: new Date().toISOString(), + notification: message, + }); + } + private broadcastTurnComplete(stopReason: string): void { if (!this.session) return; if (this.adapterEmittedTurnComplete) {