From 4220bbb63adfdd50eaab93d7c3f900a5441c3f8f Mon Sep 17 00:00:00 2001 From: Merge Sim Date: Mon, 7 Sep 2026 11:42:30 -0700 Subject: [PATCH] fix(native-chat): never preempt an in-flight handoff when reclaiming a dead owner Removing the lease-deadline precondition let the renewer evict a proven-dead native owner the instant its child exited. That is right for an idle owner, but a suspend parks a live owner at handoffStage 'preparing', and the handoff transition requires that exact stage, its operation id, and a non-null ownerProcess. An eviction landing in that window nulls the owner and bumps the fence, so the handoff then fails agent_session_ownership_unknown. The renewer now reclaims only an owner with no handoff in flight, and the eviction transition takes an optional expected handoff stage so the write is a compare-and-swap: a stage that drifted between the scan and the mutation raises agent_session_checkpoint_stale instead of clobbering the handoff. The guard lives beside the other eviction preconditions rather than in the record store, and the store now reuses the transition's input type instead of restating it. Found by review of the preceding commit, which widened a window the deadline had been masking. --- ...ctured-agent-session-lease-renewer.test.ts | 53 +++++++++++++++++++ .../structured-agent-session-lease-renewer.ts | 11 ++-- .../agent-session-lease-transitions.ts | 14 ++++- .../runtime/agent-session-record-store.ts | 13 ++--- 4 files changed, 78 insertions(+), 13 deletions(-) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts index e5b97dea189..f91ea0a1922 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts @@ -6,6 +6,7 @@ import { agentSessionLeaseFixture, agentSessionRecordFixture } from '../../../shared/agent-session-record.test-fixture' +import { stopStoredAgentSessionOwnerForHandoff } from '../../runtime/agent-session-handoff-record-transitions' import { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' import { StructuredAgentSessionLeaseRenewer } from './structured-agent-session-lease-renewer' @@ -264,6 +265,58 @@ describe('structured agent-session lease renewal', () => { expect(store.getRecord('session-renewal')?.lease.lastRenewedAt).toBe(NOW) }) + it('does not preempt a native owner handoff after the child stops', async () => { + const store = await liveStore() + const initialFence = store.getRecord('session-renewal')!.lease.runtimeFence + const probeStarted = Promise.withResolvers() + const deadProbe = Promise.withResolvers<{ outcome: 'pid-absent' }>() + const onError = vi.fn() + const renewer = new StructuredAgentSessionLeaseRenewer({ + store, + probe: async () => { + probeStarted.resolve() + return deadProbe.promise + }, + now: () => NOW + 10_000, + onError + }) + const renewal = renewer.renewNow() + await probeStarted.promise + await store.transitionHandoff('session-renewal', (record) => ({ + ...record, + lease: { + ...record.lease, + handoffStage: 'preparing', + handoffOperationId: 'handoff-1' + } + })) + deadProbe.resolve({ outcome: 'pid-absent' }) + await renewal + + expect(store.getRecord('session-renewal')?.lease).toMatchObject({ + runtimeFence: initialFence, + handoffStage: 'preparing', + handoffOperationId: 'handoff-1', + claimStatus: 'live' + }) + expect(onError).toHaveBeenCalledWith({ + sessionId: 'session-renewal', + error: expect.objectContaining({ message: 'agent_session_checkpoint_stale' }) + }) + const stopped = await stopStoredAgentSessionOwnerForHandoff(store, { + sessionId: 'session-renewal', + expectedFence: initialFence, + operationId: 'handoff-1', + now: NOW + 10_001 + }) + expect(stopped.lease).toMatchObject({ + runtimeFence: initialFence + 1, + handoffStage: 'old-owner-stopped', + handoffOperationId: 'handoff-1', + claimStatus: 'released' + }) + }) + it('routes a proven dead TUI owner into handoff recovery', async () => { const store = await liveStore() await store.transitionHandoff('session-renewal', (record) => ({ diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts index e4ea2bd30f5..949f1c305ac 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts @@ -76,13 +76,18 @@ export class StructuredAgentSessionLeaseRenewer { if (!probe) { continue } - // Positive death evidence settles it; the deadline only guards a live owner that has not - // renewed yet. Matches the dead-TUI-owner branch below, which never waits for expiry. - if (record.lease.runtimeKind === 'native' && isProvenDeadProbe(probe)) { + // Positive death evidence settles an idle owner without waiting for expiry. An active + // handoff owns its stop transition and fence; the renewer must not preempt it. + if ( + record.lease.runtimeKind === 'native' && + record.lease.handoffStage === null && + isProvenDeadProbe(probe) + ) { try { await this.input.store.evictProvenDeadOwner({ sessionId: record.sessionId, expectedFence: record.lease.runtimeFence, + expectedHandoffStage: null, probe, now }) diff --git a/src/main/runtime/agent-session-lease-transitions.ts b/src/main/runtime/agent-session-lease-transitions.ts index 28617817647..d4f41e5e3a1 100644 --- a/src/main/runtime/agent-session-lease-transitions.ts +++ b/src/main/runtime/agent-session-lease-transitions.ts @@ -203,15 +203,25 @@ export function renewAgentSessionLease(args: { } /** Proven eviction — the only other thing besides acquisition that may move the fence. */ -export function evictAgentSessionOwner(args: { +export type EvictAgentSessionOwnerInput = { record: AgentSessionRecord expectedFence: number + /** Compare-and-swap guard: a stage that drifted since the scan must not be clobbered. */ + expectedHandoffStage?: AgentSessionRecord['lease']['handoffStage'] probe: AgentSessionOwnerProbe now: number journalSettlement: 'required' | 'not-required' -}): AgentSessionRecord { +} + +export function evictAgentSessionOwner(args: EvictAgentSessionOwnerInput): AgentSessionRecord { const { record } = args assertFence(record.lease, args.expectedFence) + if ( + args.expectedHandoffStage !== undefined && + record.lease.handoffStage !== args.expectedHandoffStage + ) { + throw new Error('agent_session_checkpoint_stale') + } if (record.lease.settlementRetryRequired) { throw new Error('agent_session_ownership_unknown') } diff --git a/src/main/runtime/agent-session-record-store.ts b/src/main/runtime/agent-session-record-store.ts index 4325410ed81..07f573fc576 100644 --- a/src/main/runtime/agent-session-record-store.ts +++ b/src/main/runtime/agent-session-record-store.ts @@ -13,7 +13,6 @@ import { admitAgentSessionOperationRow, type AgentSessionOperationAdmission } from './agent-session-operation-admission' -import type { AgentSessionOwnerProbe } from '../../shared/agent-session-lease-adjudication' import { classifyObservedAgentSessionSpawnToken } from '../../shared/agent-session-lease-adjudication' import type { AgentSessionProviderHandleLink } from '../../shared/agent-session-provider-handle' import { @@ -28,7 +27,8 @@ import { evictAgentSessionOwner, proveAgentSessionOwner, setAgentSessionJournalCheckpoint, - type AgentSessionProcessIdentityCommit + type AgentSessionProcessIdentityCommit, + type EvictAgentSessionOwnerInput } from './agent-session-lease-transitions' import { settleFailedAgentSessionAcquisition, @@ -248,12 +248,9 @@ export class AgentSessionRecordStore { ) } - async evictProvenDeadOwner(args: { - sessionId: string - expectedFence: number - probe: AgentSessionOwnerProbe - now: number - }): Promise { + async evictProvenDeadOwner( + args: Omit & { sessionId: string } + ): Promise { return this.mutate(args.sessionId, (record) => evictAgentSessionOwner({ ...args, record, journalSettlement: 'required' }) )