fix(llm): drain delegated retry recovery

packages/llm/llm-retry/README.md

@@ -2,11 +2,11 @@
 
 Function plugin that applies exact-provider retry policy on the agent loop's closed-step recovery seam. It does not wrap `ctx.llm.stream()`: every adapter call remains one provider attempt, and every retry opens a fresh numbered step.
 
-Each provider adapter owns an optional nested `retryPolicy`, captured when its route registers on `ctx.llm` and carried with each call that reaches that registration's final adapter boundary. An in-flight failure retains that serving policy if the route is later disposed or replaced; a failure before any final adapter is selected has no provider policy and delegates. Omission uses normal mode: two retries for `RATE_LIMIT`, `SERVER`, `TIMEOUT`, and `TRANSPORT`. A normal policy can change its finite budget, eligible codes, and backoff. Always mode asks downstream recovery first, then retries every model-request failure without an attempt limit; success, cancellation, or plugin disposal stops it.
+Each provider adapter owns an optional nested `retryPolicy`, captured when its route registers on `ctx.llm` and carried with each call that reaches that registration's final adapter boundary. An in-flight failure retains that serving policy if the route is later disposed or replaced; a failure before any final adapter is selected has no provider policy and delegates. Omission uses normal mode: two retries for `RATE_LIMIT`, `SERVER`, `TIMEOUT`, and `TRANSPORT`. A normal policy can change its finite budget, eligible codes, and backoff. Always mode asks downstream recovery first, then retries every model-request failure without an attempt limit; success, cancellation, or plugin disposal stops it after active delegated recovery reaches quiescence.
 
 Both modes use bounded exponential backoff with symmetric jitter. A valid `providerRetryAfterMs` at or below `maxDelayMs` replaces local backoff without jitter. An over-cap provider delay makes normal mode delegate, while always mode uses its configured local backoff so it cannot terminate on that instruction.
 
-Before waiting, the plugin appends a non-surface `llm/retry` event with the provider, mode, failure, and scheduled delay. Normal events include the finite maximum; always events omit it, and UIs render `∞`. Cancellation and plugin disposal abort the wait; disposal drains active backoffs, and a callback captured before disposal fails closed.
+Before waiting, the plugin appends a non-surface `llm/retry` event with the provider, mode, failure, and scheduled delay. Normal events include the finite maximum; always events omit it, and UIs render `∞`. Cancellation and plugin disposal abort active backoff, drain active delegated recovery before applying the abort, and make a callback captured before disposal fail closed.
 
 The separately published `./invariant` companion checks that every retry record names the current open turn and latest closed step, matches the failed request's durable provider, has a unique step record and correct provider-policy retry number, and carries a valid mode-specific budget and bounded timer delay. Full jitter may schedule zero milliseconds at its lower boundary.
 
@@ -46,5 +46,5 @@ The reconstructed request preserves the prior prefix and is eligible for provide
 
 - **Agent steps are the only retry boundary** — direct `ctx.llm.stream()` consumers remain single-attempt because a raw stream cannot separate already-emitted chunks durably.
 - **Always mode retries permanent failures** — authentication, quota, invalid-request, protocol, and unrecoverable context errors continue until success, cancellation, or disposal; deployments own provider-specific cost and latency controls.
-- **Recovery policies compose by waterfall order** — always mode accepts a downstream retry before applying its fallback. A later policy that never settles also prevents the fallback from running.
+- **Recovery policies compose by waterfall order** — always mode accepts a downstream retry before applying its fallback. A later policy that ignores cancellation and never settles also prevents fallback, turn quiescence, and plugin disposal from completing.
 - **`llm/retry` records scheduling, not completion** — later step and turn events establish success, exhaustion, or cancellation.

packages/llm/llm-retry/src/index.ts

@@ -63,32 +63,15 @@ export interface RetryInternals {
 type DownstreamOutcome =
   | { readonly type: 'decision'; readonly decision: RequestErrorDecision }
   | { readonly type: 'error'; readonly error: unknown }
-  | { readonly type: 'aborted' }
 
-function downstreamUntilAbort(
+async function settleDownstream(
   next: () => Promise<RequestErrorDecision>,
-  signal: AbortSignal,
 ): Promise<DownstreamOutcome> {
-  if (signal.aborted) return Promise.resolve({ type: 'aborted' })
-  return new Promise((resolve) => {
-    const finish = (outcome: DownstreamOutcome): void => {
-      signal.removeEventListener('abort', onAbort)
-      resolve(outcome)
-    }
-    const onAbort = (): void => { finish({ type: 'aborted' }) }
-    signal.addEventListener('abort', onAbort, { once: true })
-    let downstream: Promise<RequestErrorDecision>
-    try {
-      downstream = next()
-    } catch (error: unknown) {
-      finish({ type: 'error', error })
-      return
-    }
-    void downstream.then(
-      (decision) => { finish({ type: 'decision', decision }) },
-      (error: unknown) => { finish({ type: 'error', error }) },
-    )
-  })
+  try {
+    return { type: 'decision', decision: await next() }
+  } catch (error: unknown) {
+    return { type: 'error', error }
+  }
 }
 
 function localDelay(config: ResolvedRetryPolicy, retry: number, random: () => number): number {
@@ -125,6 +108,12 @@ export function apply(ctx: Context, config: Config = {}, internals: RetryInterna
   const lifetime = new AbortController()
   const active = new Set<Promise<RequestErrorDecision>>()
 
+  function track(operation: Promise<RequestErrorDecision>): Promise<RequestErrorDecision> {
+    const tracked = operation.finally(() => active.delete(tracked))
+    active.add(tracked)
+    return tracked
+  }
+
   async function backoff(
     agent: Agent,
     turn: number,
@@ -163,7 +152,7 @@ export function apply(ctx: Context, config: Config = {}, internals: RetryInterna
     return { action: 'retry' }
   }
 
-  const disposeListener = ctx.on('agent/request-error', async (
+  async function recover(
     agent: Agent,
     turn: number,
     step: number,
@@ -173,11 +162,7 @@ export function apply(ctx: Context, config: Config = {}, internals: RetryInterna
     policy: ResolvedRetryPolicy | undefined,
     signal: AbortSignal,
     next: () => Promise<RequestErrorDecision>,
-  ) => {
-    // A waterfall may have captured this callback before its registration was
-    // removed. Lifetime cancellation must prevent that stale callback from
-    // entering a downstream policy after disposal.
-    if (lifetime.signal.aborted) return Promise.resolve<RequestErrorDecision>({ action: 'fail' })
+  ): Promise<RequestErrorDecision> {
     if (policy === undefined) return next()
     // The call-local policy belongs to the registration that served this
     // failure. Recover only the durable provider identity from the header;
@@ -188,11 +173,12 @@ export function apply(ctx: Context, config: Config = {}, internals: RetryInterna
       throw new Error(`llm-retry: no request provider for closed turn ${turn}/step ${step}`)
     }
     if (policy.mode === 'always') {
-      const downstream = await downstreamUntilAbort(
-        next,
-        AbortSignal.any([signal, lifetime.signal]),
-      )
-      if (downstream.type === 'aborted') return { action: 'fail' }
+      if (signal.aborted || lifetime.signal.aborted) return { action: 'fail' }
+      const fusedSignal = AbortSignal.any([signal, lifetime.signal])
+      // The loop and plugin lifetime stay open until delegated recovery settles.
+      // An abort then wins before the decision or fallback can mutate later state.
+      const downstream = await settleDownstream(next)
+      if (fusedSignal.aborted) return { action: 'fail' }
       if (downstream.type === 'error') {
         ctx.logger.warn(
           `llm-retry: provider "${provider}" always policy ignored a downstream recovery failure: %o`,
@@ -232,15 +218,30 @@ export function apply(ctx: Context, config: Config = {}, internals: RetryInterna
       delayMs = localDelay(policy, retry, random)
     }
 
-    const tracked = backoff(agent, turn, step, failure, provider, policy, retry, delayMs, signal)
-      .finally(() => active.delete(tracked))
-    active.add(tracked)
-    return tracked
+    return backoff(agent, turn, step, failure, provider, policy, retry, delayMs, signal)
+  }
+
+  const disposeListener = ctx.on('agent/request-error', (
+    agent: Agent,
+    turn: number,
+    step: number,
+    error: RequestError,
+    failure: LlmFailure,
+    priorFailures: readonly LlmFailure[],
+    policy: ResolvedRetryPolicy | undefined,
+    signal: AbortSignal,
+    next: () => Promise<RequestErrorDecision>,
+  ) => {
+    // A waterfall may have captured this callback before its registration was
+    // removed. Lifetime cancellation must prevent that stale callback from
+    // entering a downstream policy after disposal.
+    if (lifetime.signal.aborted) return Promise.resolve<RequestErrorDecision>({ action: 'fail' })
+    return track(recover(agent, turn, step, error, failure, priorFailures, policy, signal, next))
   })
 
   ctx.effect(() => async () => {
     disposeListener()
     lifetime.abort(new Error('llm-retry plugin disposed'))
     await Promise.allSettled([...active])
-  }, 'llm-retry: abort and drain backoffs')
+  }, 'llm-retry: abort and drain active recovery')
 }

packages/llm/llm-retry/tests/retry.spec.ts

@@ -706,61 +706,80 @@ describe('provider-routed retry policy', () => {
     expect(vi.getTimerCount()).toBe(0)
   })
 
-  it('does not make plugin disposal wait for a delegated recovery policy', async () => {
+  it('drains delegated recovery before completing plugin disposal', async () => {
     const adapter = new ScriptedAdapter([new LlmError('bad key', 'AUTH')])
     const mounted = await harness(adapter, { mock: alwaysConfig() })
     context = mounted.ctx
-    const downstream = Promise.withResolvers<RequestErrorDecision>()
+    const release = Promise.withResolvers<undefined>()
     const entered = Promise.withResolvers<undefined>()
-    context.on('agent/request-error', () => {
+    const order: string[] = []
+    context.on('agent/request-error', async () => {
       entered.resolve(undefined)
-      return downstream.promise
+      await release.promise
+      order.push('downstream')
+      return { action: 'retry' }
     })
     const agent = context.agentLoop.create(SessionId('retry-delegated-disposal'), {
       provider: 'mock',
       model: 'mock',
     })
-    const idle = waitForIdle(context, agent)
+    const idle = waitForIdle(context, agent).then(() => { order.push('idle') })
     agent.followup([{ type: 'text', text: 'go' }])
     await entered.promise
 
-    const disposing = mounted.retryFiber.dispose()
+    const disposing = mounted.retryFiber.dispose().then(() => { order.push('disposed') })
     let timer: ReturnType<typeof setTimeout> | undefined
     const outcome = await Promise.race([
       disposing.then(() => 'disposed' as const),
       new Promise<'blocked'>((resolve) => { timer = setTimeout(() => { resolve('blocked') }, 100) }),
     ])
     if (timer !== undefined) clearTimeout(timer)
-    downstream.resolve({ action: 'fail' })
+    expect(outcome).toBe('blocked')
+
+    release.resolve(undefined)
     await disposing
     await idle
 
-    expect(outcome).toBe('disposed')
+    expect(order[0]).toBe('downstream')
+    expect(order).toEqual(expect.arrayContaining(['disposed', 'idle']))
     expect(adapter.requests).toHaveLength(1)
+    expect(agent.session.events.some(event => event.type === 'llm/retry')).toBe(false)
   })
 
-  it('lets turn cancellation interrupt a delegated recovery policy', async () => {
+  it('drains delegated recovery before turn cancellation reaches idle', async () => {
     const adapter = new ScriptedAdapter([new LlmError('bad key', 'AUTH')])
     const mounted = await harness(adapter, { mock: alwaysConfig() })
     context = mounted.ctx
     const downstream = Promise.withResolvers<RequestErrorDecision>()
     const entered = Promise.withResolvers<undefined>()
-    context.on('agent/request-error', () => {
+    const order: string[] = []
+    context.on('agent/request-error', async () => {
       entered.resolve(undefined)
-      return downstream.promise
+      const decision = await downstream.promise
+      order.push('downstream')
+      return decision
     })
     const agent = context.agentLoop.create(SessionId('retry-delegated-cancel'), {
       provider: 'mock',
       model: 'mock',
     })
-    const idle = waitForIdle(context, agent)
+    const idle = waitForIdle(context, agent).then(() => { order.push('idle') })
     agent.followup([{ type: 'text', text: 'go' }])
     await entered.promise
 
     agent.cancel({ kind: 'user' })
-    await idle
-    downstream.resolve({ action: 'fail' })
+    let timer: ReturnType<typeof setTimeout> | undefined
+    const outcome = await Promise.race([
+      idle.then(() => 'idle' as const),
+      new Promise<'blocked'>((resolve) => { timer = setTimeout(() => { resolve('blocked') }, 100) }),
+    ])
+    if (timer !== undefined) clearTimeout(timer)
+    expect(outcome).toBe('blocked')
 
+    downstream.resolve({ action: 'retry' })
+    await idle
+
+    expect(order).toEqual(['downstream', 'idle'])
     expect(adapter.requests).toHaveLength(1)
     expect(agent.session.events.at(-1)).toMatchObject({
       type: 'turn/end',
@@ -773,8 +792,10 @@ describe('provider-routed retry policy', () => {
     const mounted = await harness(adapter, { mock: alwaysConfig() })
     context = mounted.ctx
     const downstream = Promise.withResolvers<RequestErrorDecision>()
+    const entered = Promise.withResolvers<undefined>()
     context.on('agent/request-error', (agent) => {
       agent.cancel({ kind: 'user' })
+      entered.resolve(undefined)
       return downstream.promise
     })
     const agent = context.agentLoop.create(SessionId('retry-delegated-sync-cancel'), {
@@ -784,8 +805,17 @@ describe('provider-routed retry policy', () => {
     const idle = waitForIdle(context, agent)
 
     agent.followup([{ type: 'text', text: 'go' }])
+    await entered.promise
+    let timer: ReturnType<typeof setTimeout> | undefined
+    const outcome = await Promise.race([
+      idle.then(() => 'idle' as const),
+      new Promise<'blocked'>((resolve) => { timer = setTimeout(() => { resolve('blocked') }, 100) }),
+    ])
+    if (timer !== undefined) clearTimeout(timer)
+    expect(outcome).toBe('blocked')
+
+    downstream.resolve({ action: 'retry' })
     await idle
-    downstream.resolve({ action: 'fail' })
 
     expect(adapter.requests).toHaveLength(1)
     expect(agent.session.events.at(-1)).toMatchObject({