fix(agent): continue cloud tasks after compaction (#3463)

Author: tatoalo

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 = [
         "---",

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<void> | null = null;
   private pendingEvents: Record<string, unknown>[] = [];
   private deliveredMessageIds = new Set<string>();
+  private pendingCompactContinuationMessageIds = new Set<string>();
   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<PromptResponse> =>
+          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<Uint8Array>,
@@ -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) {