From 7c4904898ddc977760d9d38bb89a01918450a2ce Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:38:44 -0700 Subject: [PATCH] fix(agent-session): harden persisted ownership invariants --- .../agent-session-lease-transitions.ts | 8 ++ .../agent-session-record-store-file.ts | 38 +++++- .../agent-session-record-store.test.ts | 123 +++++++++++++++++- .../runtime/agent-session-record-store.ts | 18 ++- .../agent-session-reservation-admission.ts | 12 +- .../agent-session-store-transaction-queue.ts | 31 +++++ src/shared/agent-session-record.ts | 46 ++++++- 7 files changed, 256 insertions(+), 20 deletions(-) diff --git a/src/main/runtime/agent-session-lease-transitions.ts b/src/main/runtime/agent-session-lease-transitions.ts index 37584f4b4f4..0541039942a 100644 --- a/src/main/runtime/agent-session-lease-transitions.ts +++ b/src/main/runtime/agent-session-lease-transitions.ts @@ -6,6 +6,7 @@ * observed process identity, and a proved provider handle — in that order, at one fence. */ +import { isDeepStrictEqual } from 'node:util' import { adjudicateAgentSessionRestart, evaluateAgentSessionAcquisition, @@ -46,6 +47,13 @@ function assertFence(lease: AgentSessionLease, fence: number): void { } } +export function agentSessionReconciliationTargetMatches( + current: AgentSessionRecord, + probed: AgentSessionRecord +): boolean { + return isDeepStrictEqual(current, probed) +} + /** * Compare-and-swap reservation. Writes the intent at fence + 1 before any process exists, so the * loser of a concurrent swap is refused and never spawns. diff --git a/src/main/runtime/agent-session-record-store-file.ts b/src/main/runtime/agent-session-record-store-file.ts index ae6eb2749b8..0be4f1652df 100644 --- a/src/main/runtime/agent-session-record-store-file.ts +++ b/src/main/runtime/agent-session-record-store-file.ts @@ -85,7 +85,24 @@ function parseState(raw: string, hostId: string): { state: AgentSessionStoreStat operations?: unknown retiredClaimKeys?: unknown } - if (!Number.isSafeInteger(file.schemaVersion) || typeof file.hostId !== 'string') { + if ( + !Number.isSafeInteger(file.schemaVersion) || + (file.schemaVersion as number) < 0 || + typeof file.hostId !== 'string' + ) { + return null + } + if ( + (file.schemaVersion as number) <= AGENT_SESSION_STORE_SCHEMA_VERSION && + (typeof file.records !== 'object' || + file.records === null || + Array.isArray(file.records) || + typeof file.operations !== 'object' || + file.operations === null || + Array.isArray(file.operations) || + ((file.schemaVersion as number) === AGENT_SESSION_STORE_SCHEMA_VERSION && + !Array.isArray(file.retiredClaimKeys))) + ) { return null } const state = emptyState(hostId) @@ -128,6 +145,7 @@ export async function loadAgentSessionStore( filePath: string, hostId: string ): Promise { + let unusableStoreFound = false for (const [candidate, recoveredFromBackup] of [ [filePath, false], [backupPath(filePath), true] @@ -135,11 +153,15 @@ export async function loadAgentSessionStore( let raw: string try { raw = await readFile(candidate, 'utf-8') - } catch { + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') { + unusableStoreFound = true + } continue } const parsed = parseState(raw, hostId) if (!parsed) { + unusableStoreFound = true continue } return { @@ -148,11 +170,14 @@ export async function loadAgentSessionStore( recoveredFromBackup } } + if (unusableStoreFound) { + throw new Error('agent_session_store_corrupt') + } return { state: emptyState(hostId), readOnly: false, recoveredFromBackup: false } } function serializeState(state: AgentSessionStoreState): string { - const records: Record = {} + const records: Record = Object.create(null) for (const [sessionId, value] of state.unreadableRecords) { records[sessionId] = value } @@ -180,8 +205,11 @@ export async function saveAgentSessionStore( try { await rm(backupPath(filePath), { force: true }) await rename(filePath, backupPath(filePath)) - } catch { - // First write, or the primary is already gone; the backup stays as it was. + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') { + throw error + } + // First write, or recovery already found the primary missing. } } try { diff --git a/src/main/runtime/agent-session-record-store.test.ts b/src/main/runtime/agent-session-record-store.test.ts index 01076a1d7ba..00f7fdcb5cc 100644 --- a/src/main/runtime/agent-session-record-store.test.ts +++ b/src/main/runtime/agent-session-record-store.test.ts @@ -239,6 +239,27 @@ describe('acquisition path', () => { 'agent_session_conflict' ) }) + + it.each([ + [ + 'provider', + { provider: 'codex', accountHome: { variable: 'CODEX_HOME', path: '/home/dev/.codex' } } + ], + ['account', { accountHome: { variable: 'CLAUDE_CONFIG_DIR', path: '/home/dev/.claude-other' } }] + ] as const)("refuses to change a session's pinned %s", async (_name, overrides) => { + const store = await open() + await establishOwner(store) + await expect( + store.reserveOwner( + reserveRequest({ + ...overrides, + expectedFence: 1, + probe: { outcome: 'pid-absent' }, + operation: { callerKey: 'client-1', operationId: operationId(), fingerprint: 'fp-2' } + }) + ) + ).rejects.toThrow('agent_session_conflict') + }) }) describe('concurrent claims', () => { @@ -458,6 +479,66 @@ describe('restart reconciliation', () => { expect(record?.lease.provenHandleLinkId).toBe('link-1') }) + it('never applies a stale restart probe to a replacement owner', async () => { + const writer = await open() + await establishOwner(writer) + const reconciler = await open() + let releaseProbe!: (probe: AgentSessionOwnerProbe) => void + let markProbeStarted!: () => void + const probeStarted = new Promise((resolve) => (markProbeStarted = resolve)) + const probeResult = new Promise((resolve) => (releaseProbe = resolve)) + const reconciliation = reconciler.reconcileOnRestart({ + probe: async () => { + markProbeStarted() + return probeResult + }, + now: NOW + 1_000 + }) + + await probeStarted + await writer.evictProvenDeadOwner({ + sessionId: 'session-alpha', + expectedFence: 1, + probe: { outcome: 'pid-absent' }, + now: NOW + 100 + }) + const replacement = await writer.reserveOwner( + reserveRequest({ + expectedFence: 2, + probe: UNUSED, + spawnToken: 'spawn-b', + operation: { + callerKey: 'client-1', + operationId: operationId(NOW + 200), + fingerprint: 'fp-2' + }, + now: NOW + 200 + }) + ) + await writer.commitProcessIdentity({ + sessionId: 'session-alpha', + fence: replacement.record.lease.runtimeFence, + process: processIdentity({ pid: 5252, spawnToken: 'spawn-b' }), + now: NOW + 300 + }) + await writer.proveOwner({ + sessionId: 'session-alpha', + fence: replacement.record.lease.runtimeFence, + link: handleLink({ linkId: 'link-2', origin: 'resumed', mintedAtFence: 3 }), + now: NOW + 300 + }) + + releaseProbe({ outcome: 'pid-absent' }) + expect(await reconciliation).toEqual(new Map()) + const persisted = JSON.parse(await readFile(agentSessionStorePath(directory), 'utf-8')) + expect(persisted.records['session-alpha'].lease).toMatchObject({ + runtimeFence: 3, + claimStatus: 'live', + ownerProcess: { pid: 5252, spawnToken: 'spawn-b' }, + unreconciled: false + }) + }) + it('keeps the fence monotonic across a restart and never reuses a retired fence', async () => { const first = await open() await establishOwner(first) @@ -572,7 +653,8 @@ describe('host and workspace isolation', () => { it.each([ ['WSL', WSL], ['SSH', SSH], - ['another workspace', FOLDER] + ['another workspace', FOLDER], + ['another workspace kind', { ...NATIVE, workspaceKind: 'folder' }] ] as const)('refuses to move one session id to %s', async (_name, location) => { const store = await open() await establishOwner(store) @@ -590,7 +672,7 @@ describe('host and workspace isolation', () => { it('keeps native, WSL, and SSH sessions in separate scopes', async () => { const store = await open() - await establishOwner(store, { sessionId: 'session-native' }) + await establishOwner(store, { sessionId: '__proto__' }) await establishOwner(store, { sessionId: 'session-wsl', location: WSL, spawnToken: 'spawn-b' }) await establishOwner(store, { sessionId: 'session-ssh', location: SSH, spawnToken: 'spawn-c' }) await establishOwner(store, { @@ -599,11 +681,11 @@ describe('host and workspace isolation', () => { spawnToken: 'spawn-d' }) - expect(store.listByScope(NATIVE).map((record) => record.sessionId)).toEqual(['session-native']) + expect(store.listByScope(NATIVE).map((record) => record.sessionId)).toEqual(['__proto__']) expect(store.listByScope(WSL).map((record) => record.sessionId)).toEqual(['session-wsl']) expect(store.listByScope(SSH).map((record) => record.sessionId)).toEqual(['session-ssh']) expect(store.listByScope(FOLDER).map((record) => record.sessionId)).toEqual(['session-folder']) - expect(store.listRecords()).toHaveLength(4) + expect((await open()).getRecord('__proto__')).not.toBeNull() }) it('preserves the workspace kind so a folder workspace is never read back as a worktree', async () => { @@ -708,6 +790,26 @@ describe('orphans, claim keys, checkpoints, and unreadable rows', () => { expect(persisted.records['session-alpha'].lease.runtimeFence).toBe('not-a-number') }) + it.each([ + [ + 'invalid checkpoint', + (record: AgentSessionRecord) => + Object.assign(record.lease, { journalCheckpoint: { epoch: 'bad', sequence: 1 } }) + ], + [ + 'missing live proof', + (record: AgentSessionRecord) => Object.assign(record.lease, { provenHandleLinkId: null }) + ] + ])('quarantines a record with %s', async (_name, corrupt) => { + const first = await open() + await establishOwner(first) + const filePath = agentSessionStorePath(directory) + const raw = JSON.parse(await readFile(filePath, 'utf-8')) + corrupt(raw.records['session-alpha']) + await writeFile(filePath, JSON.stringify(raw)) + expect((await open()).isSessionUnreadable('session-alpha')).toBe(true) + }) + it('recovers the previous committed state when the primary file is corrupt', async () => { const first = await open() await establishOwner(first) @@ -729,6 +831,19 @@ describe('orphans, claim keys, checkpoints, and unreadable rows', () => { expect(backup.records['session-alpha'].lease.runtimeFence).toBe(1) }) + it.each([ + ['corrupt', ['{ truncated']], + ['missing required collections', ['{"schemaVersion":1,"hostId":"local"}']], + ['corrupt in both committed copies', ['{ truncated', '{ also truncated']] + ])('fails closed when the store is %s', async (_name, copies) => { + const filePath = agentSessionStorePath(directory) + await writeFile(filePath, copies[0]) + if (copies[1]) { + await writeFile(`${filePath}.bak`, copies[1]) + } + await expect(open()).rejects.toThrow('agent_session_store_corrupt') + }) + it('refuses to write a store written by a newer schema', async () => { const filePath = agentSessionStorePath(directory) await writeFile( diff --git a/src/main/runtime/agent-session-record-store.ts b/src/main/runtime/agent-session-record-store.ts index 622f8504502..569d3f5402a 100644 --- a/src/main/runtime/agent-session-record-store.ts +++ b/src/main/runtime/agent-session-record-store.ts @@ -28,6 +28,7 @@ import { } from '../../shared/agent-session-record' import { applyAgentSessionRestartAdjudication, + agentSessionReconciliationTargetMatches, commitAgentSessionProcessIdentity, evictAgentSessionOwner, proveAgentSessionOwner, @@ -247,18 +248,25 @@ export class AgentSessionRecordStore { now: number }): Promise> { const pending = this.listRecords().filter((record) => record.lease.unreconciled) - const probes = new Map() + const probes = new Map() for (const record of pending) { - probes.set(record.sessionId, await args.probe(record)) + probes.set(record.sessionId, { record, probe: await args.probe(record) }) } return this.transact(() => { const reconciled = new Map() - for (const [sessionId, probe] of probes) { + for (const [sessionId, probed] of probes) { const record = this.state.records.get(sessionId) - if (!record?.lease.unreconciled) { + if ( + !record?.lease.unreconciled || + !agentSessionReconciliationTargetMatches(record, probed.record) + ) { continue } - const next = applyAgentSessionRestartAdjudication({ record, probe, now: args.now }) + const next = applyAgentSessionRestartAdjudication({ + record, + probe: probed.probe, + now: args.now + }) this.state.records.set(sessionId, next) reconciled.set(sessionId, next) } diff --git a/src/main/runtime/agent-session-reservation-admission.ts b/src/main/runtime/agent-session-reservation-admission.ts index 4e9c967c84f..7f5c60b2a7a 100644 --- a/src/main/runtime/agent-session-reservation-admission.ts +++ b/src/main/runtime/agent-session-reservation-admission.ts @@ -15,7 +15,7 @@ import { import type { AgentSessionOwnerProbe } from '../../shared/agent-session-lease-adjudication' import { AGENT_SESSION_RECORD_SCHEMA_VERSION, - agentSessionScopeKey, + agentSessionExecutionLocationsEqual, type AgentSessionAccountHome, type AgentSessionExecutionLocation, type AgentSessionRecord @@ -111,9 +111,13 @@ export function applyAgentSessionReservation( } return { record: createAgentSessionRecord(request, reservation), disposition: 'created' } } - if (agentSessionScopeKey(existing.location) !== agentSessionScopeKey(request.location)) { - // Why: one session id may not name two execution locations; native, WSL, and SSH copies of a - // workspace are different sessions with different provider handles. + if ( + !agentSessionExecutionLocationsEqual(existing.location, request.location) || + existing.provider !== request.provider || + existing.accountHome.variable !== request.accountHome.variable || + existing.accountHome.path !== request.accountHome.path + ) { + // Why: location, provider, and account are the session identity; changing one is a fork. throw new Error('agent_session_conflict') } if (request.expectedFence === null) { diff --git a/src/main/runtime/agent-session-store-transaction-queue.ts b/src/main/runtime/agent-session-store-transaction-queue.ts index b5be1973eb7..f85ca92c13f 100644 --- a/src/main/runtime/agent-session-store-transaction-queue.ts +++ b/src/main/runtime/agent-session-store-transaction-queue.ts @@ -1,3 +1,5 @@ +import type { AgentSessionOperationRow } from '../../shared/agent-session-operation-ledger' +import type { AgentSessionRecord } from '../../shared/agent-session-record' import { AGENT_SESSION_STORE_SCHEMA_VERSION, agentSessionStoreRevision, @@ -16,6 +18,32 @@ function markLoadedLeasesUnreconciled(state: AgentSessionStoreState): void { } } +function mapEntriesMatch(left: ReadonlyMap, right: ReadonlyMap): boolean { + if (left.size !== right.size) { + return false + } + for (const [key, value] of left) { + if (right.get(key) !== value) { + return false + } + } + return true +} + +function agentSessionStoreStateChanged( + state: AgentSessionStoreState, + records: ReadonlyMap, + operations: ReadonlyMap, + retiredClaimKeys: AgentSessionStoreState['retiredClaimKeys'] +): boolean { + return ( + !mapEntriesMatch(state.records, records) || + !mapEntriesMatch(state.operations, operations) || + state.retiredClaimKeys.length !== retiredClaimKeys.length || + state.retiredClaimKeys.some((entry, index) => entry !== retiredClaimKeys[index]) + ) +} + export class AgentSessionStoreTransactionQueue { private queue: Promise = Promise.resolve() private diskRecoveredFromBackup: boolean @@ -43,6 +71,9 @@ export class AgentSessionStoreTransactionQueue { const retiredClaimKeys = [...this.state.retiredClaimKeys] try { const result = apply() + if (!agentSessionStoreStateChanged(this.state, records, operations, retiredClaimKeys)) { + return result + } await saveAgentSessionStore(this.filePath, this.state, { recoveredFromBackup: this.diskRecoveredFromBackup }) diff --git a/src/shared/agent-session-record.ts b/src/shared/agent-session-record.ts index 4bdb52d7e79..761cde5ecb1 100644 --- a/src/shared/agent-session-record.ts +++ b/src/shared/agent-session-record.ts @@ -188,6 +188,34 @@ function isAgentSessionAccountHome(value: unknown): value is AgentSessionAccount ) } +function isAgentSessionJournalCheckpoint(value: unknown): value is AgentSessionJournalCheckpoint { + if (typeof value !== 'object' || value === null) { + return false + } + const checkpoint = value as Partial + return ( + Number.isSafeInteger(checkpoint.epoch) && + (checkpoint.epoch as number) >= 0 && + Number.isSafeInteger(checkpoint.sequence) && + (checkpoint.sequence as number) >= 0 + ) +} + +function isAgentSessionDeathEvidence(value: unknown): value is AgentSessionDeathEvidence { + if (typeof value !== 'object' || value === null) { + return false + } + const evidence = value as Partial + return ( + (evidence.kind === 'exit-observed' || + evidence.kind === 'pid-absent' || + evidence.kind === 'identity-mismatch') && + isBoundedString(evidence.detail, MAX_ID_LENGTH) && + Number.isSafeInteger(evidence.observedAt) && + (evidence.observedAt as number) >= 0 + ) +} + function isAgentSessionLease(value: unknown): value is AgentSessionLease { if (typeof value !== 'object' || value === null) { return false @@ -212,12 +240,15 @@ function isAgentSessionLease(value: unknown): value is AgentSessionLease { Number.isSafeInteger(lease.lastRenewedAt) && (lease.handoffOperationId === null || isBoundedString(lease.handoffOperationId, MAX_ID_LENGTH)) && + (lease.journalCheckpoint === null || + isAgentSessionJournalCheckpoint(lease.journalCheckpoint)) && isBoundedString(lease.claimKeyId, MAX_ID_LENGTH) && (lease.claimStatus === 'reserved' || lease.claimStatus === 'live' || lease.claimStatus === 'conflicted' || lease.claimStatus === 'released') && - typeof lease.unreconciled === 'boolean' + typeof lease.unreconciled === 'boolean' && + (lease.deathEvidence === null || isAgentSessionDeathEvidence(lease.deathEvidence)) ) } @@ -226,7 +257,7 @@ export function isAgentSessionRecord(value: unknown): value is AgentSessionRecor return false } const record = value as Partial - return ( + const shapeValid = record.schemaVersion === AGENT_SESSION_RECORD_SCHEMA_VERSION && isAgentSessionId(record.sessionId) && isAgentSessionExecutionLocation(record.location) && @@ -237,5 +268,16 @@ export function isAgentSessionRecord(value: unknown): value is AgentSessionRecor record.lease.sessionId === record.sessionId && Number.isSafeInteger(record.createdAt) && Number.isSafeInteger(record.updatedAt) + if (!shapeValid) { + return false + } + const validated = record as AgentSessionRecord + const head = validated.providerHandleChain.at(-1) + return ( + validated.providerHandleChain.every((link) => link.handle.provider === validated.provider) && + (validated.lease.claimStatus !== 'live' || + (validated.lease.ownerProcess !== null && + head?.linkId === validated.lease.provenHandleLinkId && + head.mintedAtFence === validated.lease.runtimeFence)) ) }