From 82b6df9377a6af544b21b96b91394da29860d627 Mon Sep 17 00:00:00 2001 From: Charles Packer Date: Mon, 23 Mar 2026 18:42:03 -0700 Subject: [PATCH] fix(listener): apply websocket permission mode changes during active turns (#1499) Co-authored-by: Letta Code --- .../listen-client-concurrency.test.ts | 118 +++++++++++++++++ .../listener-permission-mode.test.ts | 93 ++++++++++++++ src/websocket/listener/client.ts | 20 ++- src/websocket/listener/permissionMode.ts | 119 ++++++++++++++---- src/websocket/listener/send.ts | 4 +- src/websocket/listener/turn.ts | 40 +++--- 6 files changed, 333 insertions(+), 61 deletions(-) create mode 100644 src/tests/websocket/listener-permission-mode.test.ts diff --git a/src/tests/websocket/listen-client-concurrency.test.ts b/src/tests/websocket/listen-client-concurrency.test.ts index 28b1b56..4e117d2 100644 --- a/src/tests/websocket/listen-client-concurrency.test.ts +++ b/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((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); + }); }); diff --git a/src/tests/websocket/listener-permission-mode.test.ts b/src/tests/websocket/listener-permission-mode.test.ts new file mode 100644 index 0000000..0e7b626 --- /dev/null +++ b/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(); + }); +}); diff --git a/src/websocket/listener/client.ts b/src/websocket/listener/client.ts index f905fa5..fb7c557 100644 --- a/src/websocket/listener/client.ts +++ b/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); diff --git a/src/websocket/listener/permissionMode.ts b/src/websocket/listener/permissionMode.ts index 03462e8..bd21174 100644 --- a/src/websocket/listener/permissionMode.ts +++ b/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 { 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. diff --git a/src/websocket/listener/send.ts b/src/websocket/listener/send.ts index af0bee8..e6f9394 100644 --- a/src/websocket/listener/send.ts +++ b/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, diff --git a/src/websocket/listener/turn.ts b/src/websocket/listener/turn.ts index 8ceb225..9dc08db 100644 --- a/src/websocket/listener/turn.ts +++ b/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( - runtime.listener, - normalizedAgentId, - conversationId, - ), - }; + // 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,