fix(llm-retry): validate retries within their request step

2026-07-30 17:52:03 +08:00
parent 4e5266daa4
commit a43c6da1b9
6 changed files with 92 additions and 125 deletions
--- a/packages/llm/llm-retry/src/history.ts
+++ b/packages/llm/llm-retry/src/history.ts
@@ -1,28 +1,29 @@
-/** Durable request-route lookup for one closed model step. @module @deepseek-ai/dsh-llm-retry/history */
+/** Durable request-route lookup for one open model step. @module @deepseek-ai/dsh-llm-retry/history */

 import type { SessionEvent } from '@deepseek-ai/dsh-session'

 /**
- * Find the provider in force when one step closed, excluding later recovery mutations.
+ * Find the provider in force for one currently open step.
 * Request headers remain effective across turn boundaries until a newer full
 * snapshot changes them; every provider change requires a newer full snapshot.
- * @param events - session events containing the closed step.
+ * @param events - session events ending inside the open step.
 * @param turn - turn that owns the failed step.
 * @param step - failed step whose provider is required.
- * @returns the provider from the request header in force at that step boundary.
+ * @returns the provider from the request header in force for the step.
 */
-export function providerForClosedStep(
+export function providerForOpenStep(
  events: readonly SessionEvent[],
  turn: number,
  step: number,
 ): string | undefined {
-  const stepEndIndex = events.findLastIndex(event =>
-    event.type === 'step/end'
+  const stepStartIndex = events.findLastIndex(event =>
+    event.type === 'step/start'
    && event.data.turn === turn
    && event.data.step === step,
  )
-  if (stepEndIndex < 0) return undefined
-  for (let index = stepEndIndex; index >= 0; index -= 1) {
+  if (stepStartIndex < 0 || events.slice(stepStartIndex + 1).some(event =>
+    event.type === 'step/end' || event.type === 'turn/end')) return undefined
+  for (let index = events.length - 1; index >= 0; index -= 1) {
    // The loop bounds prove this indexed read exists.
    // oxlint-disable-next-line typescript/no-non-null-assertion
    const event = events[index]!
--- a/packages/llm/llm-retry/src/invariant.ts
+++ b/packages/llm/llm-retry/src/invariant.ts
@@ -5,7 +5,7 @@ import type { Session, SessionEvent } from '@deepseek-ai/dsh-session'
 import type { LlmFailure } from '@deepseek-ai/dsh-llm'
 import { MAX_TIMER_DELAY_MS } from '@deepseek-ai/dsh-timeout'
 import type { InvariantFailure, InvariantInstaller } from '@deepseek-ai/dsh-invariants'
-import { providerForClosedStep } from './history.ts'
+import { providerForOpenStep } from './history.ts'
 import type {} from './index.ts'

 const PACKAGE_NAME = '@deepseek-ai/dsh-llm-retry'
@@ -41,7 +41,7 @@ function validateFailure(value: unknown, fail: InvariantFailure): asserts value
  }
 }

-/** Validate one retry record against the open turn and most recently closed step. */
+/** Validate one retry record against the currently open request step. */
 function validateRetry(
  history: readonly SessionEvent[],
  event: SessionEvent<'llm/retry'>,
@@ -78,51 +78,34 @@ function validateRetry(
    fail(`llm/retry delayMs must be a finite number within 0..${MAX_TIMER_DELAY_MS}`)
  }

-  const currentTurnEvents: SessionEvent[] = []
-  let openTurn: number | undefined
-  for (const prior of history.slice().reverse()) {
-    if (prior.type === 'turn/end') fail('llm/retry must be appended inside an open turn')
-    if (prior.type === 'turn/start') {
-      openTurn = prior.data.turn
-      break
-    }
-    currentTurnEvents.push(prior)
+  const turnBoundary = history.findLast(prior =>
+    prior.type === 'turn/start' || prior.type === 'turn/end')
+  if (turnBoundary?.type !== 'turn/start') {
+    fail('llm/retry must be appended inside an open turn')
  }
-  if (openTurn === undefined) fail('llm/retry must be appended inside an open turn')
-  if (turn !== openTurn) {
-    fail(`llm/retry names turn ${turn}, but the open turn is ${openTurn}`)
+  if (turn !== turnBoundary.data.turn) {
+    fail(`llm/retry names turn ${turn}, but the open turn is ${turnBoundary.data.turn}`)
  }

-  let closedStep: number | undefined
-  for (const prior of currentTurnEvents) {
-    if (prior.type === 'step/start') {
-      fail(`llm/retry must follow step/end, but step ${prior.data.step} is still open`)
-    }
-    if (prior.type === 'step/end') {
-      closedStep = prior.data.step
-      break
-    }
+  const stepBoundary = history.findLast(prior =>
+    prior.type === 'step/start' || prior.type === 'step/end')
+  if (stepBoundary?.type !== 'step/start') {
+    fail('llm/retry must be appended inside an open step')
  }
-  if (closedStep === undefined || step !== closedStep) {
-    fail(`llm/retry names step ${step}, but the latest closed step is ${String(closedStep)}`)
+  if (step !== stepBoundary.data.step || turn !== stepBoundary.data.turn) {
+    fail(`llm/retry names turn ${turn}/step ${step}, but the open step is ${stepBoundary.data.turn}/${stepBoundary.data.step}`)
  }
-  const routedProvider = providerForClosedStep(history, turn, step)
+  const routedProvider = providerForOpenStep(history, turn, step)
  if (routedProvider !== provider) {
    fail(`llm/retry provider ${provider} does not match the failed request provider ${String(routedProvider)}`)
  }

-  const chainStart = history.findLastIndex(
-    prior => prior.type === 'turn/start' && prior.data.turn === turn,
-  )
-  const chain = history.slice(Math.max(chainStart, 0))
-  const lastSuccess = chain.findLastIndex(prior => prior.type === 'assistant/message')
-  const chainRetries = chain.slice(lastSuccess + 1)
-    .filter((prior): prior is SessionEvent<'llm/retry'> => prior.type === 'llm/retry')
-  if (chainRetries.some(prior => prior.data.turn === turn && prior.data.step === step)) {
-    fail(`llm/retry duplicates the retry record for turn ${turn}/step ${step}`)
-  }
-  const priorPolicyRetry = chainRetries.findLast(prior =>
-    prior.data.provider === provider && prior.data.policyKey === policyKey)
+  const priorPolicyRetry = history.findLast((prior): prior is SessionEvent<'llm/retry'> =>
+    prior.type === 'llm/retry'
+    && prior.data.turn === turn
+    && prior.data.step === step
+    && prior.data.provider === provider
+    && prior.data.policyKey === policyKey)
  const expectedRetry = (priorPolicyRetry?.data.retry ?? 0) + 1
  if (retry !== expectedRetry) {
    fail(`llm/retry retry ${retry} must equal provider policy retry ${expectedRetry}`)
--- a/packages/llm/llm-retry/tests/invariant.spec.ts
+++ b/packages/llm/llm-retry/tests/invariant.spec.ts
@@ -1,11 +1,11 @@
 import { describe, expect, it } from 'vitest'
 import { Context } from 'cordis'
 import SessionStore, { SessionId, type Session } from '@deepseek-ai/dsh-session'
-import { createUserMessage, ProviderRequestId , createMessage } from '@deepseek-ai/dsh-llm'
+import { createUserMessage, ProviderRequestId } from '@deepseek-ai/dsh-llm'
 import { MAX_TIMER_DELAY_MS } from '@deepseek-ai/dsh-timeout'
 import InvariantService from '@deepseek-ai/dsh-invariants'
 import * as RetryInvariant from '@deepseek-ai/dsh-llm-retry/invariant'
-import { providerForClosedStep } from '../src/history.ts'
+import { providerForOpenStep } from '../src/history.ts'

 async function setup(): Promise<Context> {
  const ctx = new Context()
@@ -15,7 +15,7 @@ async function setup(): Promise<Context> {
  return ctx
 }

-function closeStep(ctx: Context, id: string, turn = 1, step = 1) {
+function openStep(ctx: Context, id: string, turn = 1, step = 1) {
  const session = ctx.sessions.create(SessionId(id))
  session.append('turn/start', { turn })
  session.append('step/start', { turn, step })
@@ -23,7 +23,6 @@ function closeStep(ctx: Context, id: string, turn = 1, step = 1) {
    header: { config: { provider: 'mock', model: 'mock' } },
    reason: 'initial',
  })
-  session.append('step/end', { turn, step })
  return session
 }

@@ -34,7 +33,6 @@ function appendRetryTurn(session: Session, turn: number) {
    header: { config: { provider: 'mock', model: 'mock' } },
    reason: 'initial',
  })
-  session.append('step/end', { turn, step: 1 })
  session.append('llm/retry', { turn, step: 1, ...normal })
 }

@@ -58,28 +56,24 @@ const always = {
 }

 describe('llm-retry invariants', () => {
-  it('has no provider without the requested closed step or a route marker', () => {
-    expect(providerForClosedStep([], 1, 1)).toBeUndefined()
-    expect(providerForClosedStep([{
-      type: 'step/end',
+  it('has no provider without the requested open step or a route marker', () => {
+    expect(providerForOpenStep([], 1, 1)).toBeUndefined()
+    expect(providerForOpenStep([{
+      type: 'step/start',
      data: { turn: 1, step: 1 },
    }] as never, 1, 1)).toBeUndefined()
  })

-  it('accepts bounded and unbounded records after successive closed steps', async () => {
+  it('accepts successive bounded and unbounded records inside their open steps', async () => {
    const ctx = await setup()
-    const session = closeStep(ctx, 'retry-invariant-valid')
+    const session = openStep(ctx, 'retry-invariant-valid')

    expect(() => {
      session.append('llm/retry', { turn: 1, step: 1, ...normal })
-      session.append('turn/end', { turn: 1, reason: { kind: 'error', error: failure } })
-      session.append('turn/start', { turn: 2 })
-      session.append('step/start', { turn: 2, step: 1 })
-      session.append('step/end', { turn: 2, step: 1 })
      session.append('llm/retry', {
-        turn: 2, step: 1, ...normal, retry: 2, delayMs: 0,
+        turn: 1, step: 1, ...normal, retry: 2, delayMs: 0,
      })
-      const unbounded = closeStep(ctx, 'retry-invariant-always')
+      const unbounded = openStep(ctx, 'retry-invariant-always')
      unbounded.append('llm/retry', { turn: 1, step: 1, ...always })
    }).not.toThrow()
    expect(() => { ctx.emit('tools/change') }).not.toThrow()
@@ -87,7 +81,7 @@ describe('llm-retry invariants', () => {

  it('validates the complete durable failure payload', async () => {
    const ctx = await setup()
-    const complete = closeStep(ctx, 'retry-invariant-complete-failure')
+    const complete = openStep(ctx, 'retry-invariant-complete-failure')
    expect(() => {
      complete.append('llm/retry', {
        turn: 1,
@@ -126,7 +120,7 @@ describe('llm-retry invariants', () => {
      ['request-id-empty', { message: 'failed', code: 'RATE_LIMIT', requestId: '' }, /failure\.requestId/],
    ]
    for (const [name, invalidFailure, message] of invalidFailures) {
-      const session = closeStep(ctx, `retry-invariant-failure-${name}`)
+      const session = openStep(ctx, `retry-invariant-failure-${name}`)
      expect(() => {
        session.append('llm/retry', {
          turn: 1, step: 1, ...always, failure: invalidFailure,
@@ -150,43 +144,43 @@ describe('llm-retry invariants', () => {
    ['delay-type', { ...normal, delayMs: '1' }, /delayMs/],
  ])('rejects invalid retry data: %s', async (name, data, message) => {
    const ctx = await setup()
-    const session = closeStep(ctx, `retry-invariant-${name}`)
+    const session = openStep(ctx, `retry-invariant-${name}`)
    expect(() => {
      session.append('llm/retry', { turn: 1, step: 1, ...data } as never)
    }).toThrow(message)
  })

-  it('rejects records outside the latest closed step of an open turn', async () => {
+  it('rejects records outside the currently open turn and step', async () => {
    const ctx = await setup()
    const absent = ctx.sessions.create(SessionId('retry-invariant-no-turn'))
    expect(() => {
      absent.append('llm/retry', { turn: 1, step: 1, ...normal })
    }).toThrow(/inside an open turn/)

-    const wrongTurn = closeStep(ctx, 'retry-invariant-wrong-turn')
+    const wrongTurn = openStep(ctx, 'retry-invariant-wrong-turn')
    expect(() => {
      wrongTurn.append('llm/retry', { turn: 2, step: 1, ...normal })
    }).toThrow(/open turn is 1/)

-    const openStep = ctx.sessions.create(SessionId('retry-invariant-open-step'))
-    openStep.append('turn/start', { turn: 1 })
-    openStep.append('step/start', { turn: 1, step: 1 })
+    const closedStep = openStep(ctx, 'retry-invariant-closed-step')
+    closedStep.append('step/end', { turn: 1, step: 1 })
    expect(() => {
-      openStep.append('llm/retry', { turn: 1, step: 1, ...normal })
-    }).toThrow(/step 1 is still open/)
+      closedStep.append('llm/retry', { turn: 1, step: 1, ...normal })
+    }).toThrow(/inside an open step/)

    const noStep = ctx.sessions.create(SessionId('retry-invariant-no-step'))
    noStep.append('turn/start', { turn: 1 })
    expect(() => {
      noStep.append('llm/retry', { turn: 1, step: 1, ...normal })
-    }).toThrow(/latest closed step is undefined/)
+    }).toThrow(/inside an open step/)

-    const wrongStep = closeStep(ctx, 'retry-invariant-wrong-step')
+    const wrongStep = openStep(ctx, 'retry-invariant-wrong-step')
    expect(() => {
      wrongStep.append('llm/retry', { turn: 1, step: 2, ...normal })
-    }).toThrow(/latest closed step is 1/)
+    }).toThrow(/open step is 1\/1/)

-    const closedTurn = closeStep(ctx, 'retry-invariant-closed-turn')
+    const closedTurn = openStep(ctx, 'retry-invariant-closed-turn')
+    closedTurn.append('step/end', { turn: 1, step: 1 })
    closedTurn.append('turn/end', {
      turn: 1,
      reason: { kind: 'aborted', reason: { kind: 'user' } },
@@ -196,52 +190,31 @@ describe('llm-retry invariants', () => {
    }).toThrow(/inside an open turn/)
  })

-  it('rejects a second retry record for the same step', async () => {
+  it('accepts successive retries in one step and rejects skipped numbering', async () => {
    const ctx = await setup()
-    const session = closeStep(ctx, 'retry-invariant-duplicate')
+    const session = openStep(ctx, 'retry-invariant-number-sequence')
    session.append('llm/retry', { turn: 1, step: 1, ...normal })
+    session.append('llm/retry', { turn: 1, step: 1, ...normal, retry: 2 })

    expect(() => {
-      session.append('llm/retry', { turn: 1, step: 1, ...normal, retry: 2 })
-    }).toThrow(/duplicates the retry record/)
+      session.append('llm/retry', { turn: 1, step: 1, ...always, retry: 2 })
+    }).toThrow(/must equal provider policy retry 1/)
  })

-  it('binds retry numbering to the provider policy and resets it after success', async () => {
+  it('binds retry numbering to the provider policy and resets it for a new step', async () => {
    const ctx = await setup()
-    const mismatch = closeStep(ctx, 'retry-invariant-numbering')
+    const mismatch = openStep(ctx, 'retry-invariant-numbering')
    mismatch.append('llm/retry', { turn: 1, step: 1, ...normal })
-    mismatch.append('turn/end', { turn: 1, reason: { kind: 'error', error: failure } })
-    mismatch.append('turn/start', { turn: 2 })
-    mismatch.append('step/start', { turn: 2, step: 1 })
-    mismatch.append('step/end', { turn: 2, step: 1 })
    expect(() => {
-      mismatch.append('llm/retry', { turn: 2, step: 1, ...normal, retry: 1 })
+      mismatch.append('llm/retry', { turn: 1, step: 1, ...normal, retry: 1 })
    }).toThrow(/must equal provider policy retry 2/)

-    const reset = closeStep(ctx, 'retry-invariant-reset')
+    const reset = openStep(ctx, 'retry-invariant-reset')
    reset.append('llm/retry', { turn: 1, step: 1, ...normal })
-    reset.append('turn/end', { turn: 1, reason: { kind: 'error', error: failure } })
-    reset.append('turn/start', { turn: 2 })
-    reset.append('step/start', { turn: 2, step: 1 })
-    reset.append('assistant/message', {
-      turn: 2,
-      step: 1,
-      message: createMessage({
-        role: 'assistant',
-        content: [{ type: 'text', text: 'success' }],
-        source: {
-          kind: 'model',
-          ...{ provider: 'mock', model: 'mock' },
-        },
-      }),
-    }, { surfaceOp: 'append' })
-    reset.append('step/end', { turn: 2, step: 1 })
-    reset.append('turn/end', { turn: 2, reason: { kind: 'completed' } })
-    reset.append('turn/start', { turn: 3 })
-    reset.append('step/start', { turn: 3, step: 1 })
-    reset.append('step/end', { turn: 3, step: 1 })
+    reset.append('step/end', { turn: 1, step: 1 })
+    reset.append('step/start', { turn: 1, step: 2 })
    expect(() => {
-      reset.append('llm/retry', { turn: 3, step: 1, ...normal })
+      reset.append('llm/retry', { turn: 1, step: 2, ...normal })
    }).not.toThrow()
  })

@@ -277,7 +250,7 @@ describe('llm-retry invariants', () => {

  it('rejects a provider that does not match the failed request route', async () => {
    const ctx = await setup()
-    const session = closeStep(ctx, 'retry-invariant-provider')
+    const session = openStep(ctx, 'retry-invariant-provider')
    expect(() => {
      session.append('llm/retry', { turn: 1, step: 1, ...always, provider: 'other' })
    }).toThrow(/does not match the failed request provider mock/)
@@ -287,7 +260,7 @@ describe('llm-retry invariants', () => {
    const ctx = new Context()
    await ctx.plugin(SessionStore)
    const session = ctx.sessions.create(SessionId('retry-invariant-late'))
-    session.append('step/end', { turn: 1, step: 1 })
+    session.append('step/start', { turn: 1, step: 1 })
    session.append('llm/retry', { turn: 1, step: 1, ...normal })
    await ctx.plugin(InvariantService)
    await expect(ctx.plugin(RetryInvariant)).rejects.toThrow(/inside an open turn/)
--- a/packages/llm/llm-retry/tests/persistence.spec.ts
+++ b/packages/llm/llm-retry/tests/persistence.spec.ts
@@ -38,7 +38,6 @@ describe.each(['jsonl', 'sqlite'] as const)('%s retry-event persistence', (kind)
        header: { config: { provider: 'mock', model: 'mock' } },
        reason: 'initial',
      })
-      session.append('step/end', { turn: 1, step: 1 })
      const event = session.append('llm/retry', {
        turn: 1,
        step: 1,
@@ -49,6 +48,7 @@ describe.each(['jsonl', 'sqlite'] as const)('%s retry-event persistence', (kind)
        delayMs: 750,
        failure: { message: 'provider busy', code: 'RATE_LIMIT', status: 429 },
      })
+      session.append('step/end', { turn: 1, step: 1 })
      session.append('turn/end', {
        turn: 1,
        reason: {
--- a/packages/llm/llm-retry/tests/retry.spec.ts
+++ b/packages/llm/llm-retry/tests/retry.spec.ts
@@ -277,14 +277,21 @@ describe('provider-routed retry policy', () => {
    await vi.advanceTimersByTimeAsync(500)
    await idle

+    const retryEvent = agent.session.events.find(event => event.type === 'llm/retry')
    const failedChunks = agent.session.events.filter(event =>
-      event.type === 'assistant/chunk' && event.data.turn === 1 && event.data.step === 1,
+      event.type === 'assistant/chunk'
+      && retryEvent !== undefined
+      && event.seq < retryEvent.seq,
    )
-    expect(failedChunks).toHaveLength(6)
-    expect(agent.session.events.filter(event => event.type === 'assistant/message').map(event => ({
+    expect(failedChunks).toHaveLength(7)
+    const assistantMessages = agent.session.events.filter(event => event.type === 'assistant/message')
+    expect(assistantMessages.map(event => ({
      turn: event.data.turn,
      step: event.data.step,
    }))).toEqual([{ turn: 1, step: 1 }])
+    expect(failedChunks.every(event =>
+      !assistantMessages[0]?.sourceEventSeqs?.includes(event.seq),
+    )).toBe(true)
    expect(agent.session.events.some(event => event.type === 'tool/call')).toBe(false)
    expect(toolExecutions).toBe(0)
    expect(agent.session.deriveMessages().at(-1)).toMatchObject({
@@ -325,7 +332,7 @@ describe('provider-routed retry policy', () => {
    expect(agent.session.events.filter(event => event.type === 'llm/retry')).toHaveLength(2)
    expect(agent.session.events.at(-1)).toMatchObject({
      type: 'turn/end',
-      data: { reason: { kind: 'error', failure: { message: 'busy three', code: 'SERVER' } } },
+      data: { reason: { kind: 'error', error: { message: 'busy three', code: 'SERVER' } } },
    })
  })

@@ -438,7 +445,7 @@ describe('provider-routed retry policy', () => {
    expect(agent.session.events.some(event => event.type === 'llm/retry')).toBe(false)
    expect(agent.session.events.at(-1)).toMatchObject({
      type: 'turn/end',
-      data: { reason: { kind: 'error', failure: { code: 'NO_ADAPTER' } } },
+      data: { reason: { kind: 'error', error: { code: 'NO_ADAPTER' } } },
    })
  })

--- a/packages/llm/llm-retry/tests/transport-recovery.spec.ts
+++ b/packages/llm/llm-retry/tests/transport-recovery.spec.ts
@@ -110,8 +110,8 @@ describe('bounded retry through the real DeepSeek HTTP/SSE adapter', () => {
  })

  it.each([
-    ['stream_disconnect', 0] as const,
-    ['partial_disconnect', 2] as const,
+    ['stream_disconnect', 1] as const,
+    ['partial_disconnect', 3] as const,
  ])('retries %s without committing failed chunks', async (behavior, failedChunkCount) => {
    const server = await start([behavior, 'success'], {
      apiKey: 'mock-key',
@@ -130,8 +130,11 @@ describe('bounded retry through the real DeepSeek HTTP/SSE adapter', () => {

    expect(server.requests).toHaveLength(2)
    expect(server.requests[0]?.body).toEqual(server.requests[1]?.body)
+    const retryEvent = agent.session.events.find(event => event.type === 'llm/retry')
    expect(agent.session.events.filter(event =>
-      event.type === 'assistant/chunk' && event.data.turn === 1,
+      event.type === 'assistant/chunk'
+      && retryEvent !== undefined
+      && event.seq < retryEvent.seq,
    )).toHaveLength(failedChunkCount)
    expect(agent.session.events.filter(event => event.type === 'assistant/message')
      .map(event => [event.data.turn, event.data.step]))
@@ -185,12 +188,12 @@ describe('bounded retry through the real DeepSeek HTTP/SSE adapter', () => {
    expect(server.requests).toHaveLength(1)
    expect(agent.session.events.filter(event =>
      event.type === 'assistant/chunk' && event.data.turn === 1,
-    )).toHaveLength(2)
+    )).toHaveLength(3)
    expect(agent.session.events.some(event => event.type === 'assistant/message')).toBe(false)
    expect(agent.session.events.some(event => event.type === 'llm/retry')).toBe(false)
    expect(agent.session.events.at(-1)).toMatchObject({
      type: 'turn/end',
-      data: { reason: { kind: 'error', failure: { code: 'STREAM_CLOSED' } } },
+      data: { reason: { kind: 'error', error: { code: 'STREAM_CLOSED' } } },
    })
  })

@@ -232,7 +235,7 @@ describe('bounded retry through the real DeepSeek HTTP/SSE adapter', () => {
    expect(agent.session.events.filter(event => event.type === 'llm/retry')).toHaveLength(2)
    expect(agent.session.events.at(-1)).toMatchObject({
      type: 'turn/end',
-      data: { reason: { kind: 'error', failure: { code: 'TRANSPORT' } } },
+      data: { reason: { kind: 'error', error: { code: 'TRANSPORT' } } },
    })
  })
 })