fix(listener): apply websocket permission mode changes during active turns (#1499)

Co-authored-by: Letta Code <noreply@letta.com>

src/tests/websocket/listen-client-concurrency.test.ts

@@ -954,4 +954,122 @@ describe("listen-client multi-worker concurrency", () => {
     expect(runtime.loopStatus).toBe("WAITING_ON_INPUT");
     expect(runtime.queuedMessagesByItemId.size).toBe(0);
   });
+
+  test("mid-turn mode changes apply to same-turn approval classification", async () => {
+    const listener = __listenClientTestUtils.createListenerRuntime();
+    __listenClientTestUtils.setActiveRuntime(listener);
+    const runtime = __listenClientTestUtils.getOrCreateScopedRuntime(
+      listener,
+      "agent-1",
+      "conv-mid",
+    );
+    const socket = new MockSocket();
+
+    let releaseFirstDrain!: () => void;
+    const firstDrainGate = new Promise<void>((resolve) => {
+      releaseFirstDrain = resolve;
+    });
+    let drainCount = 0;
+    drainHandlers.set("conv-mid", async () => {
+      drainCount += 1;
+      if (drainCount === 1) {
+        await firstDrainGate;
+        return {
+          stopReason: "requires_approval",
+          approvals: [
+            {
+              toolCallId: "tc-1",
+              toolName: "Bash",
+              toolArgs: '{"command":"pwd"}',
+            },
+          ],
+          apiDurationMs: 0,
+        };
+      }
+      return {
+        stopReason: "end_turn",
+        approvals: [],
+        apiDurationMs: 0,
+      };
+    });
+
+    let capturedModeAtClassification: string | null = null;
+    (classifyApprovalsMock as any).mockImplementationOnce(
+      async (_approvals: any, opts: any) => {
+        capturedModeAtClassification = opts?.permissionModeState?.mode ?? null;
+        return {
+          autoAllowed: [
+            {
+              approval: {
+                toolCallId: "tc-1",
+                toolName: "Bash",
+                toolArgs: '{"command":"pwd"}',
+              },
+              permission: { decision: "allow" },
+              context: null,
+              parsedArgs: { command: "pwd" },
+            },
+          ],
+          autoDenied: [],
+          needsUserInput: [],
+        };
+      },
+    );
+    (executeApprovalBatchMock as any).mockResolvedValueOnce([
+      {
+        type: "tool",
+        tool_call_id: "tc-1",
+        status: "success",
+        tool_return: "ok",
+      },
+    ]);
+
+    const turnPromise = __listenClientTestUtils.handleIncomingMessage(
+      makeIncomingMessage("agent-1", "conv-mid", "run it"),
+      socket as unknown as WebSocket,
+      runtime,
+    );
+
+    await waitFor(() => sendMessageStreamMock.mock.calls.length >= 1);
+
+    await __listenClientTestUtils.handleChangeDeviceStateInput(listener, {
+      command: {
+        type: "change_device_state",
+        runtime: { agent_id: "agent-1", conversation_id: "conv-mid" },
+        payload: { mode: "bypassPermissions" },
+      },
+      socket: socket as unknown as WebSocket,
+      opts: {},
+      processQueuedTurn: async () => {},
+    });
+
+    releaseFirstDrain();
+
+    await turnPromise;
+
+    expect(capturedModeAtClassification === "bypassPermissions").toBe(true);
+  });
+
+  test("change_device_state does not prune default-state entry mid-turn", async () => {
+    const listener = __listenClientTestUtils.createListenerRuntime();
+    __listenClientTestUtils.setActiveRuntime(listener);
+    const socket = new MockSocket();
+
+    await __listenClientTestUtils.handleChangeDeviceStateInput(listener, {
+      command: {
+        type: "change_device_state",
+        runtime: { agent_id: "agent-1", conversation_id: "default" },
+        payload: { mode: "default" },
+      },
+      socket: socket as unknown as WebSocket,
+      opts: {},
+      processQueuedTurn: async () => {},
+    });
+
+    expect(
+      listener.permissionModeByConversation.has(
+        "agent:agent-1::conversation:default",
+      ),
+    ).toBe(true);
+  });
 });

src/tests/websocket/listener-permission-mode.test.ts

@@ -0,0 +1,93 @@
+import { describe, expect, test } from "bun:test";
+import { __listenClientTestUtils } from "../../websocket/listen-client";
+import {
+  getConversationPermissionModeState,
+  getOrCreateConversationPermissionModeStateRef,
+  getPermissionModeScopeKey,
+  pruneConversationPermissionModeStateIfDefault,
+} from "../../websocket/listener/permissionMode";
+
+describe("listener permission mode helpers", () => {
+  test("getOrCreate ref preserves identity across legacy default-key migration", () => {
+    const listener = __listenClientTestUtils.createListenerRuntime();
+    const legacyKey = getPermissionModeScopeKey(null, "default");
+
+    const legacyState = {
+      mode: "acceptEdits" as const,
+      planFilePath: null,
+      modeBeforePlan: null,
+    };
+    listener.permissionModeByConversation.set(legacyKey, legacyState);
+
+    const canonicalRef = getOrCreateConversationPermissionModeStateRef(
+      listener,
+      "agent-123",
+      "default",
+    );
+
+    expect(canonicalRef).toBe(legacyState);
+    expect(listener.permissionModeByConversation.has(legacyKey)).toBe(false);
+    expect(
+      listener.permissionModeByConversation.get(
+        getPermissionModeScopeKey("agent-123", "default"),
+      ),
+    ).toBe(legacyState);
+  });
+
+  test("read getter returns default snapshot without materializing map entry", () => {
+    const listener = __listenClientTestUtils.createListenerRuntime();
+    const scopeKey = getPermissionModeScopeKey("agent-xyz", "conv-1");
+
+    const state = getConversationPermissionModeState(
+      listener,
+      "agent-xyz",
+      "conv-1",
+    );
+
+    expect(state.mode).toBeDefined();
+    expect(listener.permissionModeByConversation.has(scopeKey)).toBe(false);
+  });
+
+  test("prune removes only default-equivalent canonical entries", () => {
+    const listener = __listenClientTestUtils.createListenerRuntime();
+    const ref = getOrCreateConversationPermissionModeStateRef(
+      listener,
+      "agent-1",
+      "conv-prune",
+    );
+
+    const prunedDefault = pruneConversationPermissionModeStateIfDefault(
+      listener,
+      "agent-1",
+      "conv-prune",
+    );
+    expect(prunedDefault).toBe(true);
+    expect(
+      listener.permissionModeByConversation.has(
+        getPermissionModeScopeKey("agent-1", "conv-prune"),
+      ),
+    ).toBe(false);
+
+    const ref2 = getOrCreateConversationPermissionModeStateRef(
+      listener,
+      "agent-1",
+      "conv-prune",
+    );
+    ref2.mode = "bypassPermissions";
+
+    const prunedNonDefault = pruneConversationPermissionModeStateIfDefault(
+      listener,
+      "agent-1",
+      "conv-prune",
+    );
+    expect(prunedNonDefault).toBe(false);
+    expect(
+      listener.permissionModeByConversation.get(
+        getPermissionModeScopeKey("agent-1", "conv-prune"),
+      ),
+    ).toBe(ref2);
+
+    // keep typechecker happy about intentionally unused ref
+    expect(ref).toBeDefined();
+  });
+});

src/websocket/listener/client.ts

@@ -57,9 +57,9 @@ import {
   stashRecoveredApprovalInterrupts,
 } from "./interrupts";
 import {
-  getConversationPermissionModeState,
+  getOrCreateConversationPermissionModeStateRef,
   loadPersistedPermissionModeMap,
-  setConversationPermissionModeState,
+  persistPermissionModeMapForRuntime,
 } from "./permissionMode";
 import {
   isListFoldersCommand,
@@ -146,32 +146,30 @@ function handleModeChange(
   try {
     const agentId = scope?.agent_id ?? null;
     const conversationId = scope?.conversation_id ?? "default";
-    const current = getConversationPermissionModeState(
+    const current = getOrCreateConversationPermissionModeStateRef(
       runtime,
       agentId,
       conversationId,
     );
 
-    const next = { ...current };
-
     // Track previous mode so ExitPlanMode can restore it
     if (msg.mode === "plan" && current.mode !== "plan") {
-      next.modeBeforePlan = current.mode;
+      current.modeBeforePlan = current.mode;
     }
-    next.mode = msg.mode;
+    current.mode = msg.mode;
 
     // Generate plan file path when entering plan mode
     if (msg.mode === "plan" && !current.planFilePath) {
-      next.planFilePath = generatePlanFilePath();
+      current.planFilePath = generatePlanFilePath();
     }
 
     // Clear plan-related state when leaving plan mode
     if (msg.mode !== "plan") {
-      next.planFilePath = null;
-      next.modeBeforePlan = null;
+      current.planFilePath = null;
+      current.modeBeforePlan = null;
     }
 
-    setConversationPermissionModeState(runtime, agentId, conversationId, next);
+    persistPermissionModeMapForRuntime(runtime);
 
     emitDeviceStatusUpdate(socket, runtime, scope);
 

src/websocket/listener/permissionMode.ts

@@ -31,11 +31,35 @@ export function getPermissionModeScopeKey(
   return `conversation:${normalizedConversationId}`;
 }
 
+function createDefaultPermissionModeState(): ConversationPermissionModeState {
+  return {
+    mode: globalPermissionMode.getMode(),
+    planFilePath: null,
+    modeBeforePlan: null,
+  };
+}
+
+function isPrunableDefaultState(
+  state: ConversationPermissionModeState,
+): boolean {
+  return (
+    state.mode === globalPermissionMode.getMode() &&
+    state.planFilePath === null &&
+    state.modeBeforePlan === null
+  );
+}
+
+/**
+ * Read-only state lookup for a conversation scope.
+ *
+ * This helper is intended for read paths (status rendering, serialization).
+ * It does not materialize new map entries for missing scopes.
+ */
 export function getConversationPermissionModeState(
   runtime: ListenerRuntime,
   agentId?: string | null,
   conversationId?: string | null,
-): ConversationPermissionModeState {
+): Readonly<ConversationPermissionModeState> {
   const scopeKey = getPermissionModeScopeKey(agentId, conversationId);
   const normalizedConversationId = normalizeConversationId(conversationId);
 
@@ -52,42 +76,78 @@ export function getConversationPermissionModeState(
     const legacyDefault =
       runtime.permissionModeByConversation.get(legacyDefaultKey);
     if (legacyDefault) {
-      if (normalizeCwdAgentId(agentId)) {
-        runtime.permissionModeByConversation.set(scopeKey, {
-          ...legacyDefault,
-        });
+      const normalizedAgentId = normalizeCwdAgentId(agentId);
+      if (normalizedAgentId) {
+        runtime.permissionModeByConversation.set(scopeKey, legacyDefault);
         runtime.permissionModeByConversation.delete(legacyDefaultKey);
       }
       return legacyDefault;
     }
   }
 
-  return {
-    mode: globalPermissionMode.getMode(),
-    planFilePath: null,
-    modeBeforePlan: null,
-  };
+  return createDefaultPermissionModeState();
 }
 
-export function setConversationPermissionModeState(
+/**
+ * Returns the canonical mutable state object for a conversation scope.
+ *
+ * This helper materializes missing entries and guarantees stable identity
+ * during a turn so concurrent mode updates (websocket + tool mutations)
+ * apply to the same object reference.
+ */
+export function getOrCreateConversationPermissionModeStateRef(
   runtime: ListenerRuntime,
-  agentId: string | null,
-  conversationId: string,
-  state: ConversationPermissionModeState,
-): void {
+  agentId?: string | null,
+  conversationId?: string | null,
+): ConversationPermissionModeState {
   const scopeKey = getPermissionModeScopeKey(agentId, conversationId);
-  // Only store if different from the global default to keep the map lean.
-  if (
-    state.mode === globalPermissionMode.getMode() &&
-    state.planFilePath === null &&
-    state.modeBeforePlan === null
-  ) {
-    runtime.permissionModeByConversation.delete(scopeKey);
-  } else {
-    runtime.permissionModeByConversation.set(scopeKey, { ...state });
+  const normalizedConversationId = normalizeConversationId(conversationId);
+
+  const direct = runtime.permissionModeByConversation.get(scopeKey);
+  if (direct) {
+    return direct;
   }
 
-  persistPermissionModeMap(runtime.permissionModeByConversation);
+  if (normalizedConversationId === "default") {
+    const legacyDefaultKey = getPermissionModeScopeKey(null, "default");
+    const legacyDefault =
+      runtime.permissionModeByConversation.get(legacyDefaultKey);
+    if (legacyDefault) {
+      const normalizedAgentId = normalizeCwdAgentId(agentId);
+      if (normalizedAgentId) {
+        runtime.permissionModeByConversation.set(scopeKey, legacyDefault);
+        runtime.permissionModeByConversation.delete(legacyDefaultKey);
+      }
+      return legacyDefault;
+    }
+  }
+
+  const created = createDefaultPermissionModeState();
+  runtime.permissionModeByConversation.set(scopeKey, created);
+  return created;
+}
+
+/**
+ * Remove a canonical state entry when it is equivalent to the default state.
+ *
+ * This should be called at turn finalization boundaries, not on each mode
+ * update, to avoid breaking object identity for in-flight turns.
+ */
+export function pruneConversationPermissionModeStateIfDefault(
+  runtime: ListenerRuntime,
+  agentId?: string | null,
+  conversationId?: string | null,
+): boolean {
+  const scopeKey = getPermissionModeScopeKey(agentId, conversationId);
+  const state = runtime.permissionModeByConversation.get(scopeKey);
+  if (!state) {
+    return false;
+  }
+  if (!isPrunableDefaultState(state)) {
+    return false;
+  }
+  runtime.permissionModeByConversation.delete(scopeKey);
+  return true;
 }
 
 /**
@@ -124,6 +184,15 @@ export function loadPersistedPermissionModeMap(): Map<
   }
 }
 
+/**
+ * Persist permission mode map to remote-settings.json.
+ */
+export function persistPermissionModeMapForRuntime(
+  runtime: ListenerRuntime,
+): void {
+  persistPermissionModeMap(runtime.permissionModeByConversation);
+}
+
 /**
  * Serialize the permission mode map and persist to remote-settings.json.
  * Strips planFilePath (ephemeral). Converts "plan" mode to modeBeforePlan.

src/websocket/listener/send.ts

@@ -41,7 +41,7 @@ import {
   emitToolExecutionFinishedEvents,
   emitToolExecutionStartedEvents,
 } from "./interrupts";
-import { getConversationPermissionModeState } from "./permissionMode";
+import { getOrCreateConversationPermissionModeStateRef } from "./permissionMode";
 import {
   emitDequeuedUserMessage,
   emitRetryDelta,
@@ -157,7 +157,7 @@ export async function resolveStaleApprovals(
         requireArgsForAutoApprove: true,
         missingNameReason: "Tool call incomplete - missing name",
         workingDirectory: recoveryWorkingDirectory,
-        permissionModeState: getConversationPermissionModeState(
+        permissionModeState: getOrCreateConversationPermissionModeStateRef(
           runtime.listener,
           runtime.agentId,
           runtime.conversationId,

src/websocket/listener/turn.ts

@@ -44,8 +44,9 @@ import {
   populateInterruptQueue,
 } from "./interrupts";
 import {
-  getConversationPermissionModeState,
-  setConversationPermissionModeState,
+  getOrCreateConversationPermissionModeStateRef,
+  persistPermissionModeMapForRuntime,
+  pruneConversationPermissionModeStateIfDefault,
 } from "./permissionMode";
 import {
   emitCanonicalMessageDelta,
@@ -98,16 +99,14 @@ export async function handleIncomingMessage(
     conversationId,
   );
 
-  // Build a mutable permission mode state object for this turn, seeded from the
-  // persistent ListenerRuntime map. Tool implementations (EnterPlanMode, ExitPlanMode)
-  // mutate it in place; we sync the final value back to the map after the turn.
-  const turnPermissionModeState = {
-    ...getConversationPermissionModeState(
+  // Get the canonical mutable permission mode state ref for this turn.
+  // Websocket mode changes and tool implementations (EnterPlanMode/ExitPlanMode)
+  // all mutate this same object in place.
+  const turnPermissionModeState = getOrCreateConversationPermissionModeStateRef(
     runtime.listener,
     normalizedAgentId,
     conversationId,
-    ),
-  };
+  );
 
   const msgRunIds: string[] = [];
   let postStopApprovalRecoveryRetries = 0;
@@ -772,22 +771,17 @@ export async function handleIncomingMessage(
       console.error("[Listen] Error handling message:", error);
     }
   } finally {
-    // Sync any permission mode changes made by tools (EnterPlanMode/ExitPlanMode)
-    // back to the persistent ListenerRuntime map so the state survives eviction.
-    setConversationPermissionModeState(
+    // Prune lean defaults only at turn-finalization boundaries (never during
+    // mid-turn mode changes), then persist the canonical map.
+    pruneConversationPermissionModeStateIfDefault(
       runtime.listener,
       normalizedAgentId,
       conversationId,
-      turnPermissionModeState,
     );
+    persistPermissionModeMapForRuntime(runtime.listener);
 
-    // Emit a corrected device status now that the permission mode is synced.
-    // The emitRuntimeStateUpdates() calls earlier in the turn read from the map
-    // before setConversationPermissionModeState() ran, so they emitted a stale
-    // current_permission_mode. This final emission sends the correct value,
-    // ensuring the web UI (and desktop) always reflect mode changes from
-    // EnterPlanMode/ExitPlanMode and that mid-turn web permission changes
-    // are not reverted by a stale emission at turn end.
+    // Emit device status after persistence/pruning so UI reflects the final
+    // canonical state for this scope.
     emitDeviceStatusIfOpen(runtime, {
       agent_id: agentId || null,
       conversation_id: conversationId,