🤖 fix: refresh post-compaction context after plan writes (#1119)

🤖 Fix post-compaction context sidebar staleness after plan writes.

- Backend now emits refreshed workspace metadata (including
`postCompaction` state) after relevant tool completions (`file_edit_*`,
`propose_plan`) when the experiment is enabled.
- Refresh is debounced in `WorkspaceService` to coalesce rapid tool
sequences.

Validation:
- `make static-check`
- Added unit tests for the tool-call-end trigger and the debounce
behavior.

---

<details>
<summary>📋 Implementation Plan</summary>

# Fix: Post-compaction context UI not updating when plan is written

## Summary
When the **Post-Compaction Context** experiment is enabled, the
right-sidebar section (Costs tab) should show a “Plan file” item as soon
as the agent writes the plan file. Currently it **doesn’t update until
the user toggles the experiment off/on**.

## Repro (current behavior)
1. Enable **Settings → Experiments → Post-Compaction Context**.
2. In a workspace, have the agent write a plan (uses `file_edit_*` tools
to create/update the plan file).
3. Observe: the right-sidebar **Post-Compaction Context** section does
**not** appear / does not show “Plan file”.
4. Toggle the experiment off then on.
5. Observe: the section updates and now shows the “Plan file” item.

## Root cause (why this happens)
**Frontend data source is stale.**

- UI path:
- `CostsTab.tsx` renders `<PostCompactionSection … />` only when
`postCompactionEnabled`.
- `usePostCompactionState(workspaceId)` reads
`workspaceMetadata.get(workspaceId)?.postCompaction` from
`WorkspaceContext`.
- If `planPath` is `null`, the section won’t render (it early-returns
when there’s no plan and no tracked files).

- Where does `workspaceMetadata[*].postCompaction` come from?
- It is only included when `WorkspaceContext.loadWorkspaceMetadata()`
calls `api.workspace.list({ includePostCompaction: true })`.
- That list call is triggered on mount and when the user toggles the
experiment (see `ExperimentsSection.tsx` calling
`refreshWorkspaceMetadata()`).

- Critically:
- Writing the plan uses `file_edit_*` tools → creates history entries
and writes the plan file on disk.
- **No code path automatically refreshes workspace metadata (or
postCompaction state) after a file-edit tool finishes.**
- Backend only emits metadata enriched with postCompaction in one place
today: **after compaction completes** (`onCompactionComplete` callback
wired in `WorkspaceService.getOrCreateSession`).

So the UI is correct given its inputs; it’s just never told to
recompute/reload post-compaction state when plan/file edits happen.

## Recommended fix (Approach A): backend-driven metadata refresh on tool
completion
**Goal:** after a `file_edit_*` tool finishes (and after
`propose_plan`), emit an updated workspace metadata event that includes
fresh `postCompaction` state.

This keeps the UI model simple: `WorkspaceContext` already subscribes to
metadata events and will update `workspaceMetadata`, which will flow
into `usePostCompactionState` and re-render the right sidebar.

### Implementation plan

#### 1) Add a new “post-compaction state may have changed” callback to
`AgentSession`
- Extend `AgentSession` constructor args to include something like:
  - `onPostCompactionStateChange?: () => void`
- Store it on the session.

#### 2) Trigger that callback on relevant tool completions
In `AgentSession.attachAiListeners()`:
- In the `"tool-call-end"` handler, after `emitChatEvent(payload)`:
  - If tool name is one of:
    - `file_edit_insert`
    - `file_edit_replace_string`
- (optionally) `propose_plan` (since it can create/validate plan state)
  - Then call `this.onPostCompactionStateChange?.()`.

**Optional gating:** only do this if the session has seen
`options?.experiments?.postCompactionContext === true` recently.
- Store `this.postCompactionExperimentEnabled` from the latest
`sendMessage` options.
- Gate callback invocation behind it.

#### 3) Implement the callback in `WorkspaceService` (debounced + safe)
In `WorkspaceService.getOrCreateSession()` pass:
- `onPostCompactionStateChange: () =>
schedulePostCompactionMetadataRefresh(workspaceId)`

Implement `schedulePostCompactionMetadataRefresh` in `WorkspaceService`:
- Coalesce bursts (multiple file edits) with a short debounce (e.g.
100–250ms).
- On fire:
  - `const metadata = await this.getInfo(workspaceId)`
- If present, compute `const postCompaction = await
this.getPostCompactionState(workspaceId)`
- `this.sessions.get(workspaceId)?.emitMetadata({ ...metadata,
postCompaction })`
- Wrap in try/catch; treat runtime-unreachable as “don’t crash; skip
emitting postCompaction”.

#### 4) Ensure metadata updates don’t accidentally drop postCompaction
- Today, some metadata events may not include `postCompaction`.
- With this approach, we explicitly re-emit enriched metadata after file
edits.
- (Optional follow-up) When experiment is enabled, also enrich *other*
metadata emits (create/title update/etc.) to avoid losing the field.

### Validation
- Manual:
  - Enable experiment.
  - Write plan.
- Confirm Post-Compaction Context section appears immediately (no
experiment toggle needed).
  - Edit plan again; confirm it remains visible.
- Edit a normal file (exec mode) and ensure tracked files list updates
as expected.
- Regression:
- With experiment disabled, ensure no extra metadata spam and section
stays hidden.

### Tests
- Add a unit test around the new trigger behavior:
- Simulate an `AgentSession` receiving a `tool-call-end` event for
`file_edit_*`.
- Assert `onPostCompactionStateChange` callback is invoked (and
debounced once for bursts).
- Add a `WorkspaceService` unit test for debounced refresh:
  - Stub `getInfo()` and `getPostCompactionState()`.
- Call scheduler multiple times; assert a single
`session.emitMetadata()` with enriched metadata.

### Net LoC estimate (product code)
- **~80–160 LoC**
  - AgentSession: add callback field + invoke on tool-call-end.
  - WorkspaceService: debounce + enriched metadata emit.

## Alternative fix (Approach B): frontend-driven fetch of
post-compaction state
Instead of relying on `workspaceMetadata.postCompaction`, update
`usePostCompactionState` to call `api.workspace.getPostCompactionState({
workspaceId })`:
- On mount / workspace change.
- On tool-call-end events (would require subscribing to WorkspaceStore
or a custom event).

### Pros
- Uses existing API (`getPostCompactionState`) without expanding backend
event behavior.

### Cons
- Needs a reliable “tool-call-end happened” signal on the frontend.
- More moving parts in UI (event wiring, polling, or store integration).

### Net LoC estimate (product code)
- **~120–220 LoC**
  - Hook changes + new event/listener plumbing.

## Decision
Proceed with **Approach A (backend-driven metadata refresh)**:
- It matches existing architecture (metadata subscription already
exists).
- It fixes the immediate bug (plan write doesn’t update sidebar) without
adding UI polling.
- It can be debounced centrally and is straightforward to reason about.

</details>

---
_Generated with `mux`_

---------

Signed-off-by: Thomas Kosiewski <tk@coder.com>

Author: ThomasK33

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

@@ -91,6 +91,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 {
@@ -102,6 +104,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 }> =
     [];
@@ -128,6 +131,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");
@@ -140,6 +148,7 @@ export class AgentSession {
       initStateManager,
       backgroundProcessManager,
       onCompactionComplete,
+      onPostCompactionStateChange,
     } = options;
 
     assert(typeof workspaceId === "string", "workspaceId must be a string");
@@ -154,6 +163,7 @@ export class AgentSession {
     this.initStateManager = initStateManager;
     this.backgroundProcessManager = backgroundProcessManager;
     this.onCompactionComplete = onCompactionComplete;
+    this.onPostCompactionStateChange = onPostCompactionStateChange;
 
     this.compactionHandler = new CompactionHandler({
       workspaceId: this.workspaceId,
@@ -546,6 +556,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));
     }
@@ -611,6 +625,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);
+  });
+});