From 4f008c36bb953dcf19a4da75508083a5c2ddccfb Mon Sep 17 00:00:00 2001 From: Jerry Gooch Date: Sat, 5 Sep 2026 16:01:59 -0500 Subject: [PATCH] fix(desktop): preserve session control action state (cherry picked from commit 0efb410ed49ed6977352b313a639ff3f9d9375c1) --- .../desktop/src/store/session-control.test.ts | 201 +++++++++++++++++- apps/desktop/src/store/session-control.ts | 70 +++++- tests/tui_gateway/test_session_control.py | 14 +- 3 files changed, 271 insertions(+), 14 deletions(-) diff --git a/apps/desktop/src/store/session-control.test.ts b/apps/desktop/src/store/session-control.test.ts index c4cc9f1f72..49055975d8 100644 --- a/apps/desktop/src/store/session-control.test.ts +++ b/apps/desktop/src/store/session-control.test.ts @@ -205,6 +205,58 @@ describe('session-control store', () => { expect($sessionControlBySession.get().s1!.capability).toBe('supported') }) + it('rebind clears busy state and transport errors without discarding every cached snapshot', async () => { + applySessionControlSnapshot('supported', FULL_SNAPSHOT) + const actionResponse = deferred() + const readResponse = deferred() + useGateway( + vi.fn((_method, params) => (params.session_id === 'supported' ? actionResponse.promise : readResponse.promise)) + ) + + const action = runSessionControlAction('supported', 'goal.pause') + const read = refreshSessionControl('unknown') + $sessionControlBySession.set({ + ...$sessionControlBySession.get(), + unsupported: { + capability: 'unsupported', + error: 'stale transport error', + loading: true, + pendingAction: 'goal.pause', + snapshot: { ...FULL_SNAPSHOT, revision: 'unsupported-snapshot' } + } + }) + + resetSessionControlAfterGatewayRebind() + + expect($sessionControlBySession.get()).toMatchObject({ + supported: { + capability: 'supported', + error: null, + loading: false, + pendingAction: null, + snapshot: { revision: FULL_SNAPSHOT.revision } + }, + unknown: { capability: 'unknown', error: null, loading: false, pendingAction: null, snapshot: null }, + unsupported: { + capability: 'unknown', + error: null, + loading: false, + pendingAction: null, + snapshot: { revision: 'unsupported-snapshot' } + } + }) + + actionResponse.reject(new Error('stale action response')) + readResponse.resolve({ control: { ...FULL_SNAPSHOT, revision: 'stale read response' } }) + await expect(action).rejects.toThrow('stale action response') + await read + + expect($sessionControlBySession.get()).toMatchObject({ + supported: { loading: false, pendingAction: null, snapshot: { revision: FULL_SNAPSHOT.revision } }, + unknown: { loading: false, pendingAction: null, snapshot: null } + }) + }) + it('retains the last good snapshot and reports a bounded ordinary read error', async () => { applySessionControlSnapshot('s1', FULL_SNAPSHOT) const message = 'x'.repeat(500) @@ -247,6 +299,84 @@ describe('session-control store', () => { expect($sessionControlBySession.get().s1!.snapshot!.revision).toBe(FULL_SNAPSHOT.revision) }) + it('downgrades a current action method-not-found once, preserves its snapshot, and still rejects', async () => { + applySessionControlSnapshot('s1', FULL_SNAPSHOT) + useGateway( + vi.fn(async () => { + throw new JsonRpcGatewayError('method not found', { code: -32601 }) + }) + ) + + await expect(runSessionControlAction('s1', 'goal.pause')).rejects.toThrow('method not found') + + expect($sessionControlBySession.get().s1).toMatchObject({ + capability: 'unsupported', + error: null, + loading: false, + pendingAction: null, + snapshot: { revision: FULL_SNAPSHOT.revision } + }) + expect(refreshLegacyGoal).toHaveBeenCalledTimes(1) + expect(refreshLegacyGoal).toHaveBeenCalledWith('s1') + }) + + it('does not let a method-not-found action response downgrade a newer event snapshot', async () => { + applySessionControlSnapshot('s1', FULL_SNAPSHOT) + const response = deferred() + useGateway(vi.fn(() => response.promise)) + + const action = runSessionControlAction('s1', 'goal.pause') + applySessionControlUpdate('s1', { ...FULL_SNAPSHOT, revision: 'event-newer' }) + response.reject(new JsonRpcGatewayError('method not found', { code: -32601 })) + + await expect(action).rejects.toThrow('method not found') + expect($sessionControlBySession.get().s1).toMatchObject({ + capability: 'supported', + error: 'method not found', + loading: false, + pendingAction: null, + snapshot: { revision: 'event-newer' } + }) + expect(refreshLegacyGoal).not.toHaveBeenCalled() + }) + + it('does not let a stale method-not-found response disturb a newer action', async () => { + applySessionControlSnapshot('s1', FULL_SNAPSHOT) + const olderResponse = deferred() + const newerResponse = deferred() + useGateway( + vi + .fn() + .mockImplementationOnce(() => olderResponse.promise) + .mockImplementationOnce(() => newerResponse.promise) + ) + + const olderAction = runSessionControlAction('s1', 'goal.pause') + const newerAction = runSessionControlAction('s1', 'goal.resume') + olderResponse.reject(new JsonRpcGatewayError('method not found', { code: -32601 })) + + await expect(olderAction).rejects.toThrow('method not found') + expect($sessionControlBySession.get().s1).toMatchObject({ + capability: 'supported', + error: null, + loading: true, + pendingAction: 'goal.resume', + snapshot: { revision: FULL_SNAPSHOT.revision } + }) + expect(refreshLegacyGoal).not.toHaveBeenCalled() + + newerResponse.resolve({ + control: { ...FULL_SNAPSHOT, revision: 'newer-action' }, + dispatch: { display: null, message: null, notice: null, output: 'resumed', type: 'exec' } + }) + await expect(newerAction).resolves.toMatchObject({ output: 'resumed' }) + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: false, + pendingAction: null, + snapshot: { revision: 'newer-action' } + }) + }) + it('makes an event-first session supported and leaves malformed events as a no-op', () => { applySessionControlUpdate('s1', FULL_SNAPSHOT) const first = $sessionControlBySession.get().s1 @@ -265,7 +395,7 @@ describe('session-control store', () => { const action = runSessionControlAction('s1', 'subgoal.add', { text: 'verify hydration' }) expect($sessionControlBySession.get().s1).toMatchObject({ capability: 'unknown', - loading: false, + loading: true, pendingAction: 'subgoal.add' }) expect($sessionControlBySession.get().s2).toBeUndefined() @@ -302,6 +432,33 @@ describe('session-control store', () => { expect($sessionControlBySession.get().s1!.snapshot!.revision).toBe(FULL_SNAPSHOT.revision) }) + it('keeps a pending action current through an event and publishes its later failure', async () => { + applySessionControlSnapshot('s1', FULL_SNAPSHOT) + const response = deferred() + useGateway(vi.fn(() => response.promise)) + + const action = runSessionControlAction('s1', 'goal.pause') + expect($sessionControlBySession.get().s1).toMatchObject({ loading: true, pendingAction: 'goal.pause' }) + + applySessionControlUpdate('s1', { ...FULL_SNAPSHOT, revision: 'event-newer' }) + expect($sessionControlBySession.get().s1).toMatchObject({ + error: null, + loading: true, + pendingAction: 'goal.pause', + snapshot: { revision: 'event-newer' } + }) + + response.reject(new Error('backend unavailable')) + await expect(action).rejects.toThrow('backend unavailable') + + expect($sessionControlBySession.get().s1).toMatchObject({ + error: 'backend unavailable', + loading: false, + pendingAction: null, + snapshot: { revision: 'event-newer' } + }) + }) + it('rejects malformed action data without promoting an unknown capability', async () => { useGateway( vi.fn(async () => ({ @@ -323,11 +480,23 @@ describe('session-control store', () => { useGateway(vi.fn(() => slow.promise)) const read = refreshSessionControl('s1') + expect($sessionControlBySession.get().s1).toMatchObject({ loading: true, pendingAction: null }) + applySessionControlUpdate('s1', { ...FULL_SNAPSHOT, revision: 'event-newer' }) + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: false, + pendingAction: null, + snapshot: { revision: 'event-newer' } + }) + slow.resolve({ control: { ...FULL_SNAPSHOT, revision: 'read-stale' } }) await read - expect($sessionControlBySession.get().s1!.snapshot!.revision).toBe('event-newer') + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: false, + pendingAction: null, + snapshot: { revision: 'event-newer' } + }) }) it('does not let a late read overwrite a newer action response', async () => { @@ -391,10 +560,20 @@ describe('session-control store', () => { const refresh = refreshSupportedSessionControlAfterTurn('s1') applySessionControlUpdate('s1', { ...FULL_SNAPSHOT, revision: 'event-newer' }) + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: false, + pendingAction: null, + snapshot: { revision: 'event-newer' } + }) + slow.resolve({ control: { ...FULL_SNAPSHOT, revision: 'post-turn-stale' } }) await refresh - expect($sessionControlBySession.get().s1!.snapshot!.revision).toBe('event-newer') + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: false, + pendingAction: null, + snapshot: { revision: 'event-newer' } + }) }) it('skips post-turn refreshes until the session is known to support the control RPC', async () => { @@ -418,19 +597,31 @@ describe('session-control store', () => { expect(unsupportedRequest).not.toHaveBeenCalled() }) - it('returns a valid stale action dispatch without letting it overwrite a newer event', async () => { + it('applies an authoritative action response after an intermediate event', async () => { const slow = deferred() useGateway(vi.fn(() => slow.promise)) const action = runSessionControlAction('s1', 'goal.pause') + expect($sessionControlBySession.get().s1).toMatchObject({ loading: true, pendingAction: 'goal.pause' }) + applySessionControlUpdate('s1', { ...FULL_SNAPSHOT, revision: 'event-newer' }) + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: true, + pendingAction: 'goal.pause', + snapshot: { revision: 'event-newer' } + }) + slow.resolve({ control: { ...FULL_SNAPSHOT, revision: 'action-stale' }, dispatch: { display: null, message: null, notice: null, output: 'paused', type: 'exec' } }) await expect(action).resolves.toMatchObject({ output: 'paused', type: 'exec' }) - expect($sessionControlBySession.get().s1!.snapshot!.revision).toBe('event-newer') + expect($sessionControlBySession.get().s1).toMatchObject({ + loading: false, + pendingAction: null, + snapshot: { revision: 'action-stale' } + }) }) it('does not schedule a timer or poller while hydrating or dispatching control state', async () => { diff --git a/apps/desktop/src/store/session-control.ts b/apps/desktop/src/store/session-control.ts index 9e47430090..044bd2fb1c 100644 --- a/apps/desktop/src/store/session-control.ts +++ b/apps/desktop/src/store/session-control.ts @@ -140,6 +140,7 @@ const DISPATCH_TYPES = new Set(['exec', 'send']) const ERROR_LIMIT = 240 const versions = new Map() +const eventVersions = new Map() export const $sessionControlBySession = atom>({}) @@ -502,6 +503,17 @@ function advanceVersion(sessionId: string): number { return next } +function currentEventVersion(sessionId: string): number { + return eventVersions.get(sessionId) ?? 0 +} + +function advanceEventVersion(sessionId: string): number { + const next = currentEventVersion(sessionId) + 1 + eventVersions.set(sessionId, next) + + return next +} + function isCurrent(sessionId: string, token: number): boolean { return currentVersion(sessionId) === token } @@ -556,7 +568,35 @@ export function applySessionControlSnapshot(sessionId: string, rawSnapshot: unkn /** Applies an event update; invalid updates are deliberately a claimed no-op. */ export function applySessionControlUpdate(sessionId: string, rawSnapshot: unknown): SessionControlEntry | undefined { - return applySessionControlSnapshot(sessionId, rawSnapshot) + if (!sessionId) { + return undefined + } + + const snapshot = parseSessionControlSnapshot(rawSnapshot) + + if (!snapshot) { + return undefined + } + + const current = $sessionControlBySession.get()[sessionId] ?? emptyEntry() + const actionIsPending = current.pendingAction !== null + + advanceEventVersion(sessionId) + + if (!actionIsPending) { + advanceVersion(sessionId) + } + + const nextSnapshot = current.snapshot?.revision === snapshot.revision ? current.snapshot : snapshot + + return publishEntry(sessionId, { + ...current, + capability: 'supported', + error: null, + loading: actionIsPending ? current.loading : false, + pendingAction: actionIsPending ? current.pendingAction : null, + snapshot: nextSnapshot + }) } export function clearSessionControl(sessionId: string): void { @@ -584,12 +624,16 @@ export function resetSessionControlAfterGatewayRebind(): void { for (const [sessionId, entry] of Object.entries(entries)) { advanceVersion(sessionId) - if (entry.capability === 'unsupported') { - nextEntries[sessionId] = { ...entry, capability: 'unknown', error: null, loading: false, pendingAction: null } - changed = true - } else { - nextEntries[sessionId] = entry + const nextEntry: SessionControlEntry = { + ...entry, + capability: entry.capability === 'unsupported' ? 'unknown' : entry.capability, + error: null, + loading: false, + pendingAction: null } + + nextEntries[sessionId] = nextEntry + changed ||= !sameEntry(entry, nextEntry) } if (changed) { @@ -605,6 +649,7 @@ export function resetSessionControlForTests(): void { $sessionControlBySession.set({}) versions.clear() + eventVersions.clear() } function beginRead(sessionId: string, background: boolean): number { @@ -627,7 +672,7 @@ function beginAction(sessionId: string, action: SessionControlAction): number { publishEntry(sessionId, { ...current, error: null, - loading: false, + loading: true, pendingAction: action }) @@ -794,6 +839,7 @@ export async function runSessionControlAction( throw new Error('Session control gateway is unavailable') } + const eventVersion = currentEventVersion(sessionId) const token = beginAction(sessionId, action) try { @@ -818,7 +864,15 @@ export async function runSessionControlAction( return dispatch } catch (error) { - if (isSessionGoneForBackgroundPolling(error)) { + if (isMethodNotFound(error)) { + const transitioned = eventVersion === currentEventVersion(sessionId) ? markUnsupported(sessionId, token) : false + + if (transitioned) { + await refreshSessionGoal(sessionId) + } else { + publishFailure(sessionId, token, error, true) + } + } else if (isSessionGoneForBackgroundPolling(error)) { finishGoneRequest(sessionId, token, true) } else { publishFailure(sessionId, token, error, true) diff --git a/tests/tui_gateway/test_session_control.py b/tests/tui_gateway/test_session_control.py index 549bc7fcf6..2b05cd28e5 100644 --- a/tests/tui_gateway/test_session_control.py +++ b/tests/tui_gateway/test_session_control.py @@ -454,14 +454,26 @@ class TestErrorsAndEvents: "message": "Continue toward the goal.", "display": "/goal resume", } + dispatch_result = { + **expected, + "model_context": {"private_trace": "never expose this nested value"}, + "private_reasoning": "never expose this model-facing field", + } + emitted = [] + monkeypatch.setattr(server, "_emit", lambda event, event_sid, payload=None: emitted.append((event, event_sid, payload))) monkeypatch.setitem( server._methods, "command.dispatch", - lambda rid, params: {"jsonrpc": "2.0", "id": rid, "result": expected}, + lambda rid, params: {"jsonrpc": "2.0", "id": rid, "result": dispatch_result}, ) response = _call(server, "session.control", session_id=sid, action="goal.resume") assert response["result"]["dispatch"] == expected + assert "model_context" not in response["result"]["dispatch"] + assert "private_reasoning" not in response["result"]["dispatch"] + assert emitted == [("session.control.update", sid, {"control": response["result"]["control"]})] + assert "model_context" not in emitted[0][2] + assert "private_reasoning" not in emitted[0][2] def test_emit_failure_must_not_turn_mutation_into_rpc_error(self, server, session, monkeypatch): """Event delivery is best-effort: a failing _emit must not swallow an