🤖 fix: refresh post-compaction context after plan writes

src/node/services/agentSession.postCompactionRefresh.test.ts

@@ -0,0 +1,168 @@
+import { describe, expect, test, mock } from "bun:test";
+import { AgentSession } from "./agentSession";
+import type { Config } from "@/node/config";
+import type { HistoryService } from "./historyService";
+import type { PartialService } from "./partialService";
+import type { AIService } from "./aiService";
+import type { InitStateManager } from "./initStateManager";
+import type { BackgroundProcessManager } from "./backgroundProcessManager";
+
+// NOTE: These tests focus on the event wiring (tool-call-end -> callback).
+// The actual post-compaction state computation is covered elsewhere.
+
+describe("AgentSession post-compaction refresh trigger", () => {
+  test("triggers callback on file_edit_* tool-call-end when experiment enabled", () => {
+    const handlers = new Map<string, (...args: unknown[]) => void>();
+
+    const aiService: AIService = {
+      on(eventName: string | symbol, listener: (...args: unknown[]) => void) {
+        handlers.set(String(eventName), listener);
+        return this;
+      },
+      off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+      stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })),
+    } as unknown as AIService;
+
+    const historyService: HistoryService = {
+      getHistory: mock(() => Promise.resolve({ success: true as const, data: [] })),
+    } as unknown as HistoryService;
+
+    const initStateManager: InitStateManager = {
+      on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+      off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+    } as unknown as InitStateManager;
+
+    const backgroundProcessManager: BackgroundProcessManager = {
+      setMessageQueued: mock(() => undefined),
+      cleanup: mock(() => Promise.resolve()),
+    } as unknown as BackgroundProcessManager;
+
+    const config: Config = { srcDir: "/tmp" } as unknown as Config;
+    const partialService: PartialService = {} as unknown as PartialService;
+
+    const onPostCompactionStateChange = mock(() => undefined);
+
+    const session = new AgentSession({
+      workspaceId: "ws",
+      config,
+      historyService,
+      partialService,
+      aiService,
+      initStateManager,
+      backgroundProcessManager,
+      onPostCompactionStateChange,
+    });
+
+    // Enable the experiment gate (normally set during sendMessage()).
+    (session as unknown as { postCompactionContextEnabled: boolean }).postCompactionContextEnabled =
+      true;
+
+    const toolEnd = handlers.get("tool-call-end");
+    expect(toolEnd).toBeDefined();
+
+    toolEnd!({
+      type: "tool-call-end",
+      workspaceId: "ws",
+      messageId: "m1",
+      toolCallId: "t1b",
+      toolName: "file_edit_replace_lines",
+      result: {},
+      timestamp: Date.now(),
+    });
+
+    toolEnd!({
+      type: "tool-call-end",
+      workspaceId: "ws",
+      messageId: "m1",
+      toolCallId: "t1",
+      toolName: "file_edit_insert",
+      result: {},
+      timestamp: Date.now(),
+    });
+
+    toolEnd!({
+      type: "tool-call-end",
+      workspaceId: "ws",
+      messageId: "m1",
+      toolCallId: "t2",
+      toolName: "bash",
+      result: {},
+      timestamp: Date.now(),
+    });
+
+    expect(onPostCompactionStateChange).toHaveBeenCalledTimes(2);
+
+    session.dispose();
+  });
+
+  test("does not trigger callback when experiment disabled", () => {
+    const handlers = new Map<string, (...args: unknown[]) => void>();
+
+    const aiService: AIService = {
+      on(eventName: string | symbol, listener: (...args: unknown[]) => void) {
+        handlers.set(String(eventName), listener);
+        return this;
+      },
+      off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+      stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })),
+    } as unknown as AIService;
+
+    const initStateManager: InitStateManager = {
+      on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+      off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+    } as unknown as InitStateManager;
+
+    const backgroundProcessManager: BackgroundProcessManager = {
+      setMessageQueued: mock(() => undefined),
+      cleanup: mock(() => Promise.resolve()),
+    } as unknown as BackgroundProcessManager;
+
+    const config: Config = { srcDir: "/tmp" } as unknown as Config;
+    const historyService: HistoryService = {
+      getHistory: mock(() => Promise.resolve({ success: true as const, data: [] })),
+    } as unknown as HistoryService;
+    const partialService: PartialService = {} as unknown as PartialService;
+
+    const onPostCompactionStateChange = mock(() => undefined);
+
+    const session = new AgentSession({
+      workspaceId: "ws",
+      config,
+      historyService,
+      partialService,
+      aiService,
+      initStateManager,
+      backgroundProcessManager,
+      onPostCompactionStateChange,
+    });
+
+    const toolEnd = handlers.get("tool-call-end");
+    expect(toolEnd).toBeDefined();
+
+    toolEnd!({
+      type: "tool-call-end",
+      workspaceId: "ws",
+      messageId: "m1",
+      toolCallId: "t1",
+      toolName: "file_edit_replace_string",
+      result: {},
+      timestamp: Date.now(),
+    });
+
+    expect(onPostCompactionStateChange).toHaveBeenCalledTimes(0);
+
+    session.dispose();
+  });
+});

src/node/services/agentSession.ts

@@ -90,6 +90,8 @@ interface AgentSessionOptions {
   backgroundProcessManager: BackgroundProcessManager;
   /** Called after compaction completes to trigger metadata refresh */
   onCompactionComplete?: () => void;
+  /** Called when post-compaction context state may have changed (plan/file edits) */
+  onPostCompactionStateChange?: () => void;
 }
 
 export class AgentSession {
@@ -101,6 +103,7 @@ export class AgentSession {
   private readonly initStateManager: InitStateManager;
   private readonly backgroundProcessManager: BackgroundProcessManager;
   private readonly onCompactionComplete?: () => void;
+  private readonly onPostCompactionStateChange?: () => void;
   private readonly emitter = new EventEmitter();
   private readonly aiListeners: Array<{ event: string; handler: (...args: unknown[]) => void }> =
     [];
@@ -127,6 +130,11 @@ export class AgentSession {
    * Used to enable the cooldown-based attachment injection.
    */
   private compactionOccurred = false;
+  /**
+   * Cache the last-known experiment state so we don't spam metadata refresh
+   * when post-compaction context is disabled.
+   */
+  private postCompactionContextEnabled = false;
 
   constructor(options: AgentSessionOptions) {
     assert(options, "AgentSession requires options");
@@ -139,6 +147,7 @@ export class AgentSession {
       initStateManager,
       backgroundProcessManager,
       onCompactionComplete,
+      onPostCompactionStateChange,
     } = options;
 
     assert(typeof workspaceId === "string", "workspaceId must be a string");
@@ -153,6 +162,7 @@ export class AgentSession {
     this.initStateManager = initStateManager;
     this.backgroundProcessManager = backgroundProcessManager;
     this.onCompactionComplete = onCompactionComplete;
+    this.onPostCompactionStateChange = onPostCompactionStateChange;
 
     this.compactionHandler = new CompactionHandler({
       workspaceId: this.workspaceId,
@@ -536,6 +546,10 @@ export class AgentSession {
     }
 
     const historyResult = await this.historyService.getHistory(this.workspaceId);
+    // Cache whether post-compaction context is enabled for this session.
+    // Used to decide whether tool-call-end should trigger metadata refresh.
+    this.postCompactionContextEnabled = Boolean(options?.experiments?.postCompactionContext);
+
     if (!historyResult.success) {
       return Err(createUnknownSendMessageError(historyResult.error));
     }
@@ -601,6 +615,18 @@ export class AgentSession {
     forward("tool-call-delta", (payload) => this.emitChatEvent(payload));
     forward("tool-call-end", (payload) => {
       this.emitChatEvent(payload);
+
+      // If post-compaction context is enabled, certain tools can change what should
+      // be displayed/injected (plan writes, tracked file diffs). Trigger a metadata
+      // refresh so the right sidebar updates without requiring an experiment toggle.
+      if (
+        this.postCompactionContextEnabled &&
+        payload.type === "tool-call-end" &&
+        (payload.toolName === "propose_plan" || payload.toolName.startsWith("file_edit_"))
+      ) {
+        this.onPostCompactionStateChange?.();
+      }
+
       // Tool call completed: auto-send queued messages
       this.sendQueuedMessages();
     });

src/node/services/workspaceService.test.ts

@@ -6,6 +6,7 @@ import type { PartialService } from "./partialService";
 import type { AIService } from "./aiService";
 import type { InitStateManager } from "./initStateManager";
 import type { ExtensionMetadataService } from "./ExtensionMetadataService";
+import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
 import type { BackgroundProcessManager } from "./backgroundProcessManager";
 
 // Helper to access private renamingWorkspaces set
@@ -14,6 +15,8 @@ function addToRenamingWorkspaces(service: WorkspaceService, workspaceId: string)
   (service as any).renamingWorkspaces.add(workspaceId);
 }
 
+// NOTE: This test file uses bun:test mocks (not Jest).
+
 describe("WorkspaceService rename lock", () => {
   let workspaceService: WorkspaceService;
   let mockAIService: AIService;
@@ -116,3 +119,116 @@ describe("WorkspaceService rename lock", () => {
     }
   });
 });
+
+describe("WorkspaceService post-compaction metadata refresh", () => {
+  let workspaceService: WorkspaceService;
+
+  beforeEach(() => {
+    const aiService: AIService = {
+      isStreaming: mock(() => false),
+      getWorkspaceMetadata: mock(() =>
+        Promise.resolve({ success: false as const, error: "not found" })
+      ),
+      on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+      off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) {
+        return this;
+      },
+    } as unknown as AIService;
+
+    const mockHistoryService: Partial<HistoryService> = {
+      getHistory: mock(() => Promise.resolve({ success: true as const, data: [] })),
+      appendToHistory: mock(() => Promise.resolve({ success: true as const, data: undefined })),
+    };
+
+    const mockConfig: Partial<Config> = {
+      srcDir: "/tmp/test",
+      getSessionDir: mock(() => "/tmp/test/sessions"),
+      generateStableId: mock(() => "test-id"),
+      findWorkspace: mock(() => null),
+    };
+
+    const mockPartialService: Partial<PartialService> = {
+      commitToHistory: mock(() => Promise.resolve({ success: true as const, data: undefined })),
+    };
+
+    const mockInitStateManager: Partial<InitStateManager> = {};
+    const mockExtensionMetadataService: Partial<ExtensionMetadataService> = {};
+    const mockBackgroundProcessManager: Partial<BackgroundProcessManager> = {
+      cleanup: mock(() => Promise.resolve()),
+    };
+
+    workspaceService = new WorkspaceService(
+      mockConfig as Config,
+      mockHistoryService as HistoryService,
+      mockPartialService as PartialService,
+      aiService,
+      mockInitStateManager as InitStateManager,
+      mockExtensionMetadataService as ExtensionMetadataService,
+      mockBackgroundProcessManager as BackgroundProcessManager
+    );
+  });
+
+  test("debounces multiple refresh requests into a single metadata emit", async () => {
+    const workspaceId = "ws-post-compaction";
+
+    const emitMetadata = mock(() => undefined);
+
+    interface WorkspaceServiceTestAccess {
+      sessions: Map<string, { emitMetadata: (metadata: unknown) => void }>;
+      getInfo: (workspaceId: string) => Promise<FrontendWorkspaceMetadata | null>;
+      getPostCompactionState: (workspaceId: string) => Promise<{
+        planPath: string | null;
+        trackedFilePaths: string[];
+        excludedItems: string[];
+      }>;
+      schedulePostCompactionMetadataRefresh: (workspaceId: string) => void;
+    }
+
+    const svc = workspaceService as unknown as WorkspaceServiceTestAccess;
+    svc.sessions.set(workspaceId, { emitMetadata });
+
+    const fakeMetadata: FrontendWorkspaceMetadata = {
+      id: workspaceId,
+      name: "ws",
+      projectName: "proj",
+      projectPath: "/tmp/proj",
+      namedWorkspacePath: "/tmp/proj/ws",
+      runtimeConfig: { type: "local", srcBaseDir: "/tmp" },
+    };
+
+    const getInfoMock: WorkspaceServiceTestAccess["getInfo"] = mock(() =>
+      Promise.resolve(fakeMetadata)
+    );
+
+    const postCompactionState = {
+      planPath: "~/.mux/plans/cmux/plan.md",
+      trackedFilePaths: ["/tmp/proj/file.ts"],
+      excludedItems: [],
+    };
+
+    const getPostCompactionStateMock: WorkspaceServiceTestAccess["getPostCompactionState"] = mock(
+      () => Promise.resolve(postCompactionState)
+    );
+
+    svc.getInfo = getInfoMock;
+    svc.getPostCompactionState = getPostCompactionStateMock;
+
+    svc.schedulePostCompactionMetadataRefresh(workspaceId);
+    svc.schedulePostCompactionMetadataRefresh(workspaceId);
+    svc.schedulePostCompactionMetadataRefresh(workspaceId);
+
+    // Debounce is short, but use a safe buffer.
+    await new Promise((resolve) => setTimeout(resolve, 150));
+
+    expect(getInfoMock).toHaveBeenCalledTimes(1);
+    expect(getPostCompactionStateMock).toHaveBeenCalledTimes(1);
+    expect(emitMetadata).toHaveBeenCalledTimes(1);
+
+    const enriched = (emitMetadata as ReturnType<typeof mock>).mock.calls[0][0] as {
+      postCompaction?: { planPath: string | null };
+    };
+    expect(enriched.postCompaction?.planPath).toBe(postCompactionState.planPath);
+  });
+});