From 765e345ca50f2002dc7ad150676abbfcfb274697 Mon Sep 17 00:00:00 2001 From: m4air Date: Tue, 15 Sep 2026 23:48:10 -0700 Subject: [PATCH] fix(chat): release provider children after lost resume holds --- .../structured-hold-retention/README.md | 50 +++ .../structured-hold-retention/reproduce.mjs | 120 ++++++++ .../structured-hold-retention/results.json | 18 ++ ...red-agent-session-hold-resume-race.test.ts | 289 ++++++++++++++++++ .../structured-agent-session-holders.ts | 28 +- .../structured-agent-session-holds.ts | 17 +- .../structured-agent-session-hold.test.ts | 101 ++++++ .../methods/structured-agent-session-hold.ts | 4 +- ...ss-version-agent-session-wire.unit.test.ts | 20 +- 9 files changed, 621 insertions(+), 26 deletions(-) create mode 100644 docs/audits/structured-hold-retention/README.md create mode 100644 docs/audits/structured-hold-retention/reproduce.mjs create mode 100644 docs/audits/structured-hold-retention/results.json create mode 100644 src/main/native-chat/agent-session-wire/structured-agent-session-hold-resume-race.test.ts diff --git a/docs/audits/structured-hold-retention/README.md b/docs/audits/structured-hold-retention/README.md new file mode 100644 index 00000000000..030db63a9c6 --- /dev/null +++ b/docs/audits/structured-hold-retention/README.md @@ -0,0 +1,50 @@ +# Structured session hold lost during resume + +Run from the repository root: + +```sh +ORCA_BACKGROUND_LAUNCH=1 node docs/audits/structured-hold-retention/reproduce.mjs +``` + +The script bundles the actual `StructuredAgentSessionHolds` implementation in memory. It runs it +once without the post-resume holder check and once with the current source. It uses a deferred +provider acquisition, an isolated fake child, and a 5 ms release grace. It launches no application, +provider, or terminal and reads no user profile. + +The last surface releases its hold while acquisition is pending. At that point the session has no +provider child, so `release()` cannot arm the release clock. Before the fix, acquisition completes +with a child, zero holders, and no scheduled eviction. With the fix, successful acquisition checks +for surviving holders and schedules the existing release clock. The recorded child is released once. + +The RPC path registers connection cleanup before awaiting `host.hold()` in +`src/main/runtime/rpc/methods/structured-agent-session-hold.ts`. Runtime socket close calls +`cleanupSubscriptionsForConnection()` in `runtime-rpc/runtime-rpc-lifecycle.ts`. That supplies the +production release-during-acquisition ordering reproduced here. + +## Ownership limits + +- This proves a lifecycle race, not that it caused any particular OOM report. No process RSS was + measured. It applies to structured sessions acquiring a provider child, not ordinary PTY tabs. +- The release clock preserves its 15-second production grace, waits while a turn is active, and + cancels when a new holder arrives. Acquisition failures retain their existing handling. +- Disposal prevents late acquisition or release callbacks from restarting the clock. Host teardown + owns cleanup after disposal. The host's broader pre-attach shutdown admission is outside this fix. +- A restored childless journal is not necessarily abandoned. Startup selects persisted visible + tabs; `host.sessions` supplies `listSessionTabs()`, and childless sessions can retain live TUI + owners. `host.close()` closes that TUI owner before removing the journal. This fix neither evicts + childless history nor infers process exit from transport loss. +- Execution remains on the owning runtime, with no wire or SSH routing changes. + +Targeted regressions live in +`src/main/native-chat/agent-session-wire/structured-agent-session-hold-resume-race.test.ts`. +They cover last-holder loss, active turns, new holders, reconnection, failed acquisition, explicit +close, and disposal. + +Same-ID replacement is fenced at both ownership layers. Holder entries receive a new incarnation +after release and re-add, so an old failed acquisition cannot remove a replacement. The RPC uses +the subscription registry's `releaseIfCurrent()` cleanup, so its failure cannot unregister the +replacement's connection cleanup. Duplicate adds remain one holder. Class tests and real host/RPC +tests cover the old failure arriving before and after replacement success; disconnect still releases +the replacement normally. They also cover the reverse outcome: an old acquisition succeeds and the +replacement refuses a stale fence. Its last-holder rollback starts the same turn-aware release clock +for the acquired child. No RPC fields or published frame shapes change. diff --git a/docs/audits/structured-hold-retention/reproduce.mjs b/docs/audits/structured-hold-retention/reproduce.mjs new file mode 100644 index 00000000000..68658113a52 --- /dev/null +++ b/docs/audits/structured-hold-retention/reproduce.mjs @@ -0,0 +1,120 @@ +import { createHash } from 'node:crypto' +import { readFile } from 'node:fs/promises' +import { createRequire } from 'node:module' +import { fileURLToPath } from 'node:url' +import { build } from 'esbuild' + +if (process.env.ORCA_BACKGROUND_LAUNCH !== '1') { + throw new Error('Run with ORCA_BACKGROUND_LAUNCH=1.') +} + +const root = fileURLToPath(new URL('../../../', import.meta.url)) +const sourcePath = fileURLToPath( + new URL( + '../../../src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts', + import.meta.url + ) +) +const source = await readFile(sourcePath, 'utf8') +const postResumeCheck = + ' // The last surface can disconnect before acquisition makes a child available to release.\n' + + ' if (!this.disposed && !this.holders.isHeld(sessionId)) {\n' + + ' this.clock.arm(sessionId)\n' + + ' }\n' +if (!source.includes(postResumeCheck)) { + throw new Error('Source changed: review the before-fix transform before running this proof.') +} + +async function loadHolds(withPostResumeCheck) { + const result = await build({ + absWorkingDir: root, + entryPoints: [sourcePath], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + logLevel: 'silent', + plugins: [ + { + name: 'compare-post-resume-holder-check', + setup(plugin) { + plugin.onLoad({ filter: /structured-agent-session-holds\.ts$/ }, () => ({ + contents: withPostResumeCheck ? source : source.replace(postResumeCheck, ''), + loader: 'ts' + })) + } + } + ] + }) + const loaded = { exports: {} } + new Function('module', 'exports', 'require', result.outputFiles[0].text)( + loaded, + loaded.exports, + createRequire(import.meta.url) + ) + return loaded.exports.StructuredAgentSessionHolds +} + +async function reproduce(Holds) { + const gate = Promise.withResolvers() + let child = false + let evictions = 0 + const holds = new Holds({ + resume: async () => { + await gate.promise + child = true + }, + hasProviderChild: () => child, + isTurnActive: () => false, + evict: async () => { + evictions += 1 + child = false + }, + graceMs: 5 + }) + try { + const acquiring = holds.hold('restored-session', 'connection:surface') + holds.release('restored-session', 'connection:surface') + gate.resolve() + await acquiring + const releasePendingAfterAcquisition = holds.isReleasePending('restored-session') + await new Promise((resolve) => setTimeout(resolve, 30)) + return { + child, + held: holds.isHeld('restored-session'), + releasePendingAfterAcquisition, + evictions + } + } finally { + holds.dispose() + } +} + +const before = await reproduce(await loadHolds(false)) +const after = await reproduce(await loadHolds(true)) +const passed = + before.child && + !before.held && + !before.releasePendingAfterAcquisition && + before.evictions === 0 && + !after.child && + !after.held && + after.releasePendingAfterAcquisition && + after.evictions === 1 +console.log( + JSON.stringify( + { + source: 'src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts', + sourceSha256: createHash('sha256').update(source).digest('hex'), + comparison: 'same source, before omits only the post-resume holder check', + before, + after, + passed + }, + null, + 2 + ) +) +if (!passed) { + process.exitCode = 1 +} diff --git a/docs/audits/structured-hold-retention/results.json b/docs/audits/structured-hold-retention/results.json new file mode 100644 index 00000000000..6019410f918 --- /dev/null +++ b/docs/audits/structured-hold-retention/results.json @@ -0,0 +1,18 @@ +{ + "source": "src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts", + "sourceSha256": "c95e27518d6cb09f1c97e5ff18bbb9fc5b15c7680d4a99e90cc14353d602230f", + "comparison": "same source, before omits only the post-resume holder check", + "before": { + "child": true, + "held": false, + "releasePendingAfterAcquisition": false, + "evictions": 0 + }, + "after": { + "child": false, + "held": false, + "releasePendingAfterAcquisition": true, + "evictions": 1 + }, + "passed": true +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-hold-resume-race.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-hold-resume-race.test.ts new file mode 100644 index 00000000000..bd3a840d9ad --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-hold-resume-race.test.ts @@ -0,0 +1,289 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { StructuredAgentSessionHolds } from './structured-agent-session-holds' + +const GRACE_MS = 15_000 +const pendingHolds: StructuredAgentSessionHolds[] = [] + +function resumeHarness() { + const resumeGate = Promise.withResolvers() + let child = false + let turnActive = false + const evict = vi.fn(async () => { + child = false + }) + const holds = new StructuredAgentSessionHolds({ + resume: async () => { + await resumeGate.promise + child = true + }, + hasProviderChild: () => child, + isTurnActive: () => turnActive, + evict, + graceMs: GRACE_MS + }) + pendingHolds.push(holds) + return { + holds, + resumeGate, + evict, + hasChild: () => child, + setTurnActive: (value: boolean) => { + turnActive = value + } + } +} + +beforeEach(() => vi.useFakeTimers()) +afterEach(() => { + for (const holds of pendingHolds.splice(0)) { + holds.dispose() + } + vi.useRealTimers() +}) + +describe('a surface leaving while its structured session resumes', () => { + it('releases the acquired child after the last surface disconnects during resume', async () => { + const { holds, resumeGate, evict, hasChild } = resumeHarness() + const hold = holds.hold('session-1', 'connection-1:chat') + + holds.release('session-1', 'connection-1:chat') + expect(holds.isReleasePending('session-1')).toBe(false) + resumeGate.resolve() + await hold + + expect(hasChild()).toBe(true) + expect(holds.isHeld('session-1')).toBe(false) + expect(holds.isReleasePending('session-1')).toBe(true) + await vi.advanceTimersByTimeAsync(GRACE_MS - 1) + expect(evict).not.toHaveBeenCalled() + await vi.advanceTimersByTimeAsync(1) + expect(evict).toHaveBeenCalledExactlyOnceWith('session-1') + expect(hasChild()).toBe(false) + }) + + it('waits for an active turn before releasing the late child', async () => { + const { holds, resumeGate, evict, setTurnActive } = resumeHarness() + const hold = holds.hold('session-1', 'connection-1:chat') + holds.release('session-1', 'connection-1:chat') + setTurnActive(true) + resumeGate.resolve() + await hold + + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(evict).not.toHaveBeenCalled() + expect(holds.isReleasePending('session-1')).toBe(true) + + setTurnActive(false) + await vi.advanceTimersByTimeAsync(GRACE_MS) + expect(evict).toHaveBeenCalledExactlyOnceWith('session-1') + }) + + it.each([false, true])('preserves an arriving holder with resume=%s', async (resume) => { + const { holds, resumeGate, evict } = resumeHarness() + const first = holds.hold('session-1', 'connection-1:chat') + holds.release('session-1', 'connection-1:chat') + const replacement = holds.hold('session-1', 'connection-2:chat', { resume }) + resumeGate.resolve() + await Promise.all([first, replacement]) + + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(holds.isHeld('session-1')).toBe(true) + expect(holds.isReleasePending('session-1')).toBe(false) + expect(evict).not.toHaveBeenCalled() + + holds.release('session-1', 'connection-2:chat') + await vi.advanceTimersByTimeAsync(GRACE_MS) + expect(evict).toHaveBeenCalledExactlyOnceWith('session-1') + }) + + it('cancels the late-child release when a surface reconnects during grace', async () => { + const { holds, resumeGate, evict } = resumeHarness() + const hold = holds.hold('session-1', 'connection-1:chat') + holds.release('session-1', 'connection-1:chat') + resumeGate.resolve() + await hold + expect(holds.isReleasePending('session-1')).toBe(true) + + await holds.hold('session-1', 'connection-2:chat') + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(holds.isReleasePending('session-1')).toBe(false) + expect(evict).not.toHaveBeenCalled() + }) + + it('preserves a failed resume without scheduling eviction', async () => { + const { holds, resumeGate, evict, hasChild } = resumeHarness() + const failure = new Error('provider acquisition failed') + const hold = holds.hold('session-1', 'connection-1:chat') + const rejected = expect(hold).rejects.toBe(failure) + holds.release('session-1', 'connection-1:chat') + resumeGate.reject(failure) + await rejected + + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(hasChild()).toBe(false) + expect(holds.isHeld('session-1')).toBe(false) + expect(holds.isReleasePending('session-1')).toBe(false) + expect(evict).not.toHaveBeenCalled() + }) + + it('leaves late acquisition cleanup to host teardown after disposal', async () => { + const { holds, resumeGate, evict } = resumeHarness() + const hold = holds.hold('session-1', 'connection-1:chat') + holds.release('session-1', 'connection-1:chat') + holds.dispose() + resumeGate.resolve() + await hold + + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(holds.isReleasePending('session-1')).toBe(false) + expect(evict).not.toHaveBeenCalled() + }) + + it('does not restart release timers when a surface leaves after disposal', async () => { + const { holds, resumeGate, evict } = resumeHarness() + const hold = holds.hold('session-1', 'connection-1:chat') + resumeGate.resolve() + await hold + holds.dispose() + holds.release('session-1', 'connection-1:chat') + + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(holds.isHeld('session-1')).toBe(false) + expect(holds.isReleasePending('session-1')).toBe(false) + expect(evict).not.toHaveBeenCalled() + }) + + it('releases a late child acquired after explicit close forgot its holders', async () => { + const { holds, resumeGate, evict } = resumeHarness() + const hold = holds.hold('session-1', 'connection-1:chat') + holds.forget('session-1') + resumeGate.resolve() + await hold + + await vi.advanceTimersByTimeAsync(GRACE_MS) + expect(holds.isHeld('session-1')).toBe(false) + expect(evict).toHaveBeenCalledExactlyOnceWith('session-1') + }) + + it.each([false, true])( + 'keeps a reused holder when old resume fails (replacement finished=%s)', + async (replacementFinished) => { + const firstGate = Promise.withResolvers() + const replacementGate = Promise.withResolvers() + let child = false + const resume = vi + .fn() + .mockImplementationOnce(() => firstGate.promise) + .mockImplementationOnce(async () => { + await replacementGate.promise + child = true + }) + const evict = vi.fn(async () => {}) + const holds = new StructuredAgentSessionHolds({ + resume, + hasProviderChild: () => child, + isTurnActive: () => false, + evict, + graceMs: GRACE_MS + }) + pendingHolds.push(holds) + const first = holds.hold('session-1', 'same-holder') + const rejected = expect(first).rejects.toThrow('old acquisition failed') + holds.release('session-1', 'same-holder') + const replacement = holds.hold('session-1', 'same-holder') + if (replacementFinished) { + replacementGate.resolve() + await replacement + } + + firstGate.reject(new Error('old acquisition failed')) + await rejected + expect(holds.isHeld('session-1')).toBe(true) + replacementGate.resolve() + await replacement + await vi.advanceTimersByTimeAsync(GRACE_MS * 2) + expect(evict).not.toHaveBeenCalled() + + holds.release('session-1', 'same-holder') + await vi.advanceTimersByTimeAsync(GRACE_MS) + expect(evict).toHaveBeenCalledExactlyOnceWith('session-1') + } + ) + + it('removes a failed replacement while the released old hold is still pending', async () => { + const firstGate = Promise.withResolvers() + const resume = vi + .fn() + .mockImplementationOnce(() => firstGate.promise) + .mockRejectedValueOnce(new Error('replacement acquisition failed')) + const holds = new StructuredAgentSessionHolds({ + resume, + hasProviderChild: () => false, + isTurnActive: () => false, + evict: async () => {}, + graceMs: GRACE_MS + }) + pendingHolds.push(holds) + const first = holds.hold('session-1', 'same-holder') + const rejected = expect(first).rejects.toThrow('old acquisition failed') + holds.release('session-1', 'same-holder') + + await expect(holds.hold('session-1', 'same-holder')).rejects.toThrow( + 'replacement acquisition failed' + ) + expect(holds.isHeld('session-1')).toBe(false) + expect(holds.isReleasePending('session-1')).toBe(false) + + firstGate.reject(new Error('old acquisition failed')) + await rejected + expect(holds.isHeld('session-1')).toBe(false) + }) + + it.each(['old-holder', 'different-holder'])( + 'releases the old acquisition after replacement %s fails, once its turn finishes', + async (replacementHolder) => { + const firstGate = Promise.withResolvers() + const replacementGate = Promise.withResolvers() + let child = false + let turnActive = true + const resume = vi + .fn() + .mockImplementationOnce(async () => { + await firstGate.promise + child = true + }) + .mockImplementationOnce(() => replacementGate.promise) + const evict = vi.fn(async () => { + child = false + }) + const holds = new StructuredAgentSessionHolds({ + resume, + hasProviderChild: () => child, + isTurnActive: () => turnActive, + evict, + graceMs: GRACE_MS + }) + pendingHolds.push(holds) + const first = holds.hold('session-1', 'old-holder') + holds.release('session-1', 'old-holder') + const replacement = holds.hold('session-1', replacementHolder) + const rejected = expect(replacement).rejects.toThrow('replacement acquisition failed') + firstGate.resolve() + await first + expect(holds.isReleasePending('session-1')).toBe(false) + + replacementGate.reject(new Error('replacement acquisition failed')) + await rejected + expect(holds.isHeld('session-1')).toBe(false) + expect(holds.isReleasePending('session-1')).toBe(true) + await vi.advanceTimersByTimeAsync(GRACE_MS) + expect(evict).not.toHaveBeenCalled() + expect(child).toBe(true) + + turnActive = false + await vi.advanceTimersByTimeAsync(GRACE_MS) + expect(evict).toHaveBeenCalledExactlyOnceWith('session-1') + expect(child).toBe(false) + } + ) +}) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-holders.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-holders.ts index ed0451efb79..531717168b1 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-holders.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-holders.ts @@ -6,23 +6,35 @@ // still looking at, and a lost one leaks the child forever. A set answers both idempotently, // because it records WHICH surface holds the session, not how many do. +type Holder = { resumeCapable: boolean; incarnation: symbol } + export class StructuredAgentSessionHolders { - private readonly bySession = new Map>() + private readonly bySession = new Map>() /** True when the session gained its FIRST holder — the edge that ends a pending release. */ add(sessionId: string, holderId: string, resumeCapable = true): boolean { const holders = this.bySession.get(sessionId) if (!holders) { - this.bySession.set(sessionId, new Map([[holderId, resumeCapable]])) + this.bySession.set(sessionId, new Map([[holderId, { resumeCapable, incarnation: Symbol() }]])) return true } - holders.set(holderId, (holders.get(holderId) ?? false) || resumeCapable) + const previous = holders.get(holderId) + holders.set(holderId, { + resumeCapable: (previous?.resumeCapable ?? false) || resumeCapable, + incarnation: previous?.incarnation ?? Symbol() + }) return false } /** True when the session lost its LAST holder — the edge that starts one. */ - remove(sessionId: string, holderId: string): boolean { + remove(sessionId: string, holderId: string, expectedIncarnation?: symbol): boolean { const holders = this.bySession.get(sessionId) + if ( + expectedIncarnation !== undefined && + holders?.get(holderId)?.incarnation !== expectedIncarnation + ) { + return false + } if (!holders?.delete(holderId) || holders.size > 0) { return false } @@ -38,12 +50,18 @@ export class StructuredAgentSessionHolders { return this.bySession.get(sessionId)?.has(holderId) ?? false } + incarnation(sessionId: string, holderId: string): symbol | undefined { + return this.bySession.get(sessionId)?.get(holderId)?.incarnation + } + holderIds(sessionId: string): string[] { return [...(this.bySession.get(sessionId)?.keys() ?? [])] } hasResumeCapableHolder(sessionId: string): boolean { - return [...(this.bySession.get(sessionId)?.values() ?? [])].some(Boolean) + return [...(this.bySession.get(sessionId)?.values() ?? [])].some( + (holder) => holder.resumeCapable + ) } /** Drops every holder of one session without evaluating the edge, for a session that is gone. */ diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts index 6afc8abc052..027c2d54772 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-holds.ts @@ -36,6 +36,7 @@ export type StructuredAgentSessionHoldOptions = { export class StructuredAgentSessionHolds { private readonly holders = new StructuredAgentSessionHolders() private readonly clock: StructuredAgentSessionReleaseClock + private disposed = false constructor(private readonly deps: StructuredAgentSessionHoldsDeps) { const clockDeps: StructuredAgentSessionReleaseClockDeps = { @@ -55,6 +56,7 @@ export class StructuredAgentSessionHolds { ): Promise { const alreadyHeld = this.holders.has(sessionId, holderId) this.holders.add(sessionId, holderId, options.resume !== false) + const incarnation = this.holders.incarnation(sessionId, holderId) // Unconditional, not only on the first-holder edge: a second surface arriving during the grace // window must cancel the pending release too. this.clock.cancel(sessionId) @@ -66,19 +68,23 @@ export class StructuredAgentSessionHolds { if (!this.deps.hasProviderChild(sessionId)) { throw new Error('agent_session_ownership_unknown') } + // The last surface can disconnect before acquisition makes a child available to release. + if (!this.disposed && !this.holders.isHeld(sessionId)) { + this.clock.arm(sessionId) + } } catch (error) { - if (!alreadyHeld) { - this.holders.remove(sessionId, holderId) + if (!alreadyHeld && incarnation !== undefined) { + this.release(sessionId, holderId, incarnation) } throw error } } - release(sessionId: string, holderId: string): void { - if (!this.holders.remove(sessionId, holderId)) { + release(sessionId: string, holderId: string, expectedIncarnation?: symbol): void { + if (!this.holders.remove(sessionId, holderId, expectedIncarnation)) { return } - if (this.deps.hasProviderChild(sessionId)) { + if (!this.disposed && this.deps.hasProviderChild(sessionId)) { this.clock.arm(sessionId) } } @@ -102,6 +108,7 @@ export class StructuredAgentSessionHolds { } dispose(): void { + this.disposed = true this.clock.dispose() } } diff --git a/src/main/runtime/rpc/methods/structured-agent-session-hold.test.ts b/src/main/runtime/rpc/methods/structured-agent-session-hold.test.ts index 4e6dfdf45bf..da34108604a 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session-hold.test.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session-hold.test.ts @@ -162,6 +162,107 @@ describe('a client that holds a session', () => { }) describe('a client that disappears without cleanup', () => { + it('releases a late child after its same-ID replacement refuses the stale fence', async () => { + await host.close(SESSION) + await host.restoreReadableSessions() + closeSession.mockClear() + const firstEntered = Promise.withResolvers() + const firstGate = Promise.withResolvers() + const replacementEntered = Promise.withResolvers() + const replacementGate = Promise.withResolvers() + const attach = host.attach.bind(host) + const attachSpy = vi + .spyOn(host, 'attach') + .mockImplementationOnce(async (...args) => { + firstEntered.resolve() + await firstGate.promise + return attach(...args) + }) + .mockImplementationOnce(async (...args) => { + replacementEntered.resolve() + await replacementGate.promise + return attach(...args) + }) + try { + const params = { sessionId: SESSION, holderId: 'same-chat' } + const first = call('agentSession.hold', params) + await firstEntered.promise + const replacement = call('agentSession.hold', params) + await replacementEntered.promise + firstGate.resolve() + expect(await first).toMatchObject({ ok: true }) + expect(host.isHeld(SESSION)).toBe(true) + expect(closeSession).not.toHaveBeenCalled() + + replacementGate.resolve() + expect(await replacement).toMatchObject({ + ok: false, + error: { code: 'agent_session_checkpoint_stale' } + }) + expect(host.isHeld(SESSION)).toBe(false) + await vi.waitFor(() => expect(host.hasSession(SESSION)).toBe(false)) + expect(closeSession).toHaveBeenCalledExactlyOnceWith(SESSION) + } finally { + firstGate.resolve() + replacementGate.resolve() + attachSpy.mockRestore() + } + }) + + it.each([false, true])( + 'keeps replacement hold and cleanup after an old request fails (replacement finished=%s)', + async (replacementFinished) => { + await host.close(SESSION) + await host.restoreReadableSessions() + closeSession.mockClear() + const firstEntered = Promise.withResolvers() + const firstGate = Promise.withResolvers() + const replacementEntered = Promise.withResolvers() + const replacementGate = Promise.withResolvers() + const attach = host.attach.bind(host) + const attachSpy = vi + .spyOn(host, 'attach') + .mockImplementationOnce(async () => { + firstEntered.resolve() + await firstGate.promise + throw new Error('old acquisition failed') + }) + .mockImplementationOnce(async (...args) => { + replacementEntered.resolve() + await replacementGate.promise + return attach(...args) + }) + try { + const params = { sessionId: SESSION, holderId: 'same-chat' } + const first = call('agentSession.hold', params) + await firstEntered.promise + const replacement = call('agentSession.hold', params) + await replacementEntered.promise + if (replacementFinished) { + replacementGate.resolve() + expect(await replacement).toMatchObject({ ok: true }) + } + + firstGate.resolve() + expect(await first).toMatchObject({ ok: false }) + expect(host.isHeld(SESSION)).toBe(true) + replacementGate.resolve() + expect(await replacement).toMatchObject({ ok: true }) + await new Promise((resolve) => setTimeout(resolve, GRACE_MS * 4)) + expect(host.hasSession(SESSION)).toBe(true) + expect(closeSession).not.toHaveBeenCalled() + + runtime.cleanupSubscriptionsForConnection(CONNECTION) + await vi.waitFor(() => expect(host.hasSession(SESSION)).toBe(false)) + expect(closeSession).toHaveBeenCalledExactlyOnceWith(SESSION) + } finally { + firstGate.resolve() + replacementGate.resolve() + attachSpy.mockRestore() + } + } + ) + it('still releases the session when its transport closes', async () => { await call('agentSession.hold', { sessionId: SESSION, holderId: 'chat-1' }) diff --git a/src/main/runtime/rpc/methods/structured-agent-session-hold.ts b/src/main/runtime/rpc/methods/structured-agent-session-hold.ts index 18fb600a949..54753fcb889 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session-hold.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session-hold.ts @@ -36,7 +36,7 @@ export const STRUCTURED_AGENT_SESSION_HOLD_METHODS = [ await ensureStructuredHostInstalled(ctx) const host = requireStructuredHost(ctx) const holderKey = holderKeyFor(ctx, params.holderId) - ctx.runtime.registerSubscriptionCleanup( + const registration = ctx.runtime.registerOwnedSubscriptionCleanup( holdCleanupIdFor(params.sessionId, holderKey), () => host.release(params.sessionId, holderKey), ctx.connectionId @@ -44,7 +44,7 @@ export const STRUCTURED_AGENT_SESSION_HOLD_METHODS = [ try { await host.hold(params.sessionId, holderKey) } catch (error) { - ctx.runtime.cleanupSubscription(holdCleanupIdFor(params.sessionId, holderKey)) + registration.releaseIfCurrent() throw error } return { held: true as const } diff --git a/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts b/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts index ce41d778ce2..e9697d3ac64 100644 --- a/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts +++ b/tests/e2e/cross-version-wire/cross-version-agent-session-wire.unit.test.ts @@ -21,6 +21,7 @@ import type { AgentSessionAttachParams } from '../../../src/main/native-chat/age import { StructuredAgentSessionHost } from '../../../src/main/native-chat/agent-session-wire/structured-agent-session-host' import { setStructuredAgentSessionHost } from '../../../src/main/native-chat/agent-session-wire/structured-agent-session-registry' import { AgentSessionRecordStore } from '../../../src/main/runtime/agent-session-record-store' +import { RuntimeSubscriptionRegistry } from '../../../src/main/runtime/runtime-subscription-registry' import { computeAgentSessionPayloadFingerprint } from '../../../src/shared/agent-session-mutation-envelope' import type { AgentSessionSubscribeEvent } from '../../../src/shared/agent-session-wire' import { @@ -272,7 +273,7 @@ function paramsFor(method: string): unknown { } function runtimeStub(): unknown { - const cleanups = new Map void>() + const subscriptions = new RuntimeSubscriptionRegistry() return { getRuntimeId: () => 'runtime-1', getClientSettings: () => ({ experimentalStructuredNativeChat: true }), @@ -287,19 +288,10 @@ function runtimeStub(): unknown { return resolved }, publishStructuredAgentSessionTab: () => {}, - registerSubscriptionCleanup: (id: string, cleanup: () => void) => cleanups.set(id, cleanup), - cleanupSubscription: (id: string) => { - cleanups.get(id)?.() - cleanups.delete(id) - }, - cleanupSubscriptionsByPrefix: (prefix: string) => { - for (const [id, cleanup] of cleanups) { - if (id.startsWith(prefix)) { - cleanup() - cleanups.delete(id) - } - } - } + registerSubscriptionCleanup: subscriptions.register.bind(subscriptions), + registerOwnedSubscriptionCleanup: subscriptions.registerOwned.bind(subscriptions), + cleanupSubscription: subscriptions.cleanup.bind(subscriptions), + cleanupSubscriptionsByPrefix: subscriptions.cleanupByPrefix.bind(subscriptions) } }