fix(desktop): preserve session control action state

(cherry picked from commit 0efb410ed49ed6977352b313a639ff3f9d9375c1)
This commit is contained in:
Jerry Gooch
2026-09-05 16:01:59 -05:00
committed by Teknium
parent dffd8d62c2
commit 4f008c36bb
3 changed files with 271 additions and 14 deletions
+196 -5
View File
@@ -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<unknown>()
const readResponse = deferred<unknown>()
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<unknown>()
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<unknown>()
const newerResponse = deferred<unknown>()
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<unknown>()
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<unknown>()
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 () => {
+62 -8
View File
@@ -140,6 +140,7 @@ const DISPATCH_TYPES = new Set<SessionControlDispatch['type']>(['exec', 'send'])
const ERROR_LIMIT = 240
const versions = new Map<string, number>()
const eventVersions = new Map<string, number>()
export const $sessionControlBySession = atom<Record<string, SessionControlEntry>>({})
@@ -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)
+13 -1
View File
@@ -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