fix(chat): release provider children after lost resume holds

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<void>()
+  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<void>()
+      const replacementGate = Promise.withResolvers<void>()
+      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<void>()
+    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<void>()
+      const replacementGate = Promise.withResolvers<void>()
+      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)
+    }
+  )
+})

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<string, Map<string, boolean>>()
+  private readonly bySession = new Map<string, Map<string, Holder>>()
 
   /** 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. */

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<void> {
     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()
   }
 }

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<void>()
+    const firstGate = Promise.withResolvers<void>()
+    const replacementEntered = Promise.withResolvers<void>()
+    const replacementGate = Promise.withResolvers<void>()
+    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<void>()
+      const firstGate = Promise.withResolvers<void>()
+      const replacementEntered = Promise.withResolvers<void>()
+      const replacementGate = Promise.withResolvers<void>()
+      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' })
 

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 }