Add bounded agent-step recovery retries

packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts

@@ -976,6 +976,74 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
   })
 
   describe('API error handling', () => {
+    it('retries one transient provider failure in the same agent step', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      let llmCallNumber = 0
+      const recoveryAttempts: string[] = []
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* ({
+        extraCodebuffMetadata,
+        onCostCalculated,
+      }) {
+        llmCallNumber++
+        recoveryAttempts.push(extraCodebuffMetadata?.recovery_attempt ?? '0')
+
+        if (llmCallNumber === 1) {
+          yield { type: 'text' as const, text: 'Partial failed response' }
+          await onCostCalculated?.(7)
+          loopAgentStepsBaseParams.agentState.childRunIds.push(
+            'failed-child-run',
+          )
+          throw new APICallError({
+            statusCode: 503,
+            message: 'Service unavailable',
+            url: 'https://api.codebuff.com/v1/chat/completions',
+            requestBodyValues: {},
+            responseBody: undefined,
+            isRetryable: true,
+          })
+        }
+
+        yield { type: 'text' as const, text: 'Recovered response\n\n' }
+        yield createToolCallChunk('end_turn', {})
+        return promptSuccess('recovered-message-id')
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+        waitForAgentRecovery: async () => {},
+      })
+
+      expect(result.output.type).not.toBe('error')
+      expect(llmCallNumber).toBe(2)
+      expect(recoveryAttempts).toEqual(['0', '1'])
+      expect(
+        result.agentState.messageHistory.some((message) =>
+          message.tags?.includes('AGENT_RECOVERY'),
+        ),
+      ).toBe(true)
+      expect(result.agentState.directCreditsUsed).toBe(7)
+      expect(result.agentState.childRunIds).toEqual(['failed-child-run'])
+      expect(JSON.stringify(result.agentState.messageHistory)).not.toContain(
+        'Partial failed response',
+      )
+      expect(loopAgentStepsBaseParams.addAgentStep).toHaveBeenCalledWith(
+        expect.objectContaining({
+          credits: 7,
+          childRunIds: ['failed-child-run'],
+        }),
+      )
+    })
+
     it('should propagate error code and server message from 403 APICallError responseBody', async () => {
       const llmOnlyTemplate = {
         ...mockTemplate,

packages/agent-runtime/src/run-agent-step.ts

@@ -41,6 +41,11 @@ import { getAgentPrompt } from './templates/strings'
 import { getToolSet } from './tools/prompts'
 import { processStream } from './tools/stream-parser'
 import { getAgentOutput } from './util/agent-output'
+import {
+  classifyAgentRecovery,
+  getAgentRecoveryDelayMs,
+  MAX_AGENT_STEP_RECOVERY_ATTEMPTS,
+} from './util/agent-recovery'
 import {
   createCacheDebugSnapshot,
   enrichCacheDebugSnapshotWithProviderRequest,
@@ -59,6 +64,40 @@ import {
   countTokensMessages,
 } from './util/token-counter'
 
+type AgentRecoveryWaitParams = {
+  attempt: number
+  delayMs: number
+  kind: import('./util/agent-recovery').AgentRecoveryKind
+  signal: AbortSignal
+}
+
+const waitForAgentRecovery = async ({
+  delayMs,
+  signal,
+}: AgentRecoveryWaitParams): Promise<void> => {
+  if (delayMs <= 0) return
+
+  await new Promise<void>((resolve, reject) => {
+    let timeout: ReturnType<typeof setTimeout> | undefined
+
+    const onAbort = () => {
+      if (timeout !== undefined) clearTimeout(timeout)
+      reject(new AbortError())
+    }
+
+    if (signal.aborted) {
+      onAbort()
+      return
+    }
+
+    timeout = setTimeout(() => {
+      signal.removeEventListener('abort', onAbort)
+      resolve()
+    }, delayMs)
+    signal.addEventListener('abort', onAbort, { once: true })
+  })
+}
+
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics'
 import type {
@@ -355,8 +394,7 @@ export const runAgentStep = async (
   const systemTokens = countTokens(system)
 
   let cacheDebugCorrelation:
-    | ReturnType<typeof createCacheDebugSnapshot>
-    | undefined
+    ReturnType<typeof createCacheDebugSnapshot> | undefined
   if (CACHE_DEBUG_FULL_LOGGING) {
     try {
       cacheDebugCorrelation = createCacheDebugSnapshot({
@@ -689,6 +727,8 @@ export async function loopAgentSteps(
      * to the message history as user prompts and keep the turn going, letting a
      * host "steer" a running agent without aborting or losing the current step. */
     drainSteeringMessages?: () => string[]
+    /** Override the recovery backoff in tests or hosts with their own scheduler. */
+    waitForAgentRecovery?: (params: AgentRecoveryWaitParams) => Promise<void>
     spawnParams: Record<string, any> | undefined
     startAgentRun: StartAgentRunFn
     userId: string | undefined
@@ -1179,28 +1219,109 @@ export async function loopAgentSteps(
       const creditsBefore = currentAgentState.directCreditsUsed
       const childrenBefore = currentAgentState.childRunIds.length
       llmStepNumber++
+      let recoveryAttempt = 0
+      let stepResult: Awaited<ReturnType<typeof runAgentStep>>
+
+      while (true) {
+        // runAgentStep mutates the shared state while streaming. If a provider
+        // fails after emitting partial text/tool calls, roll back the
+        // non-billable transcript before retrying so the next request does not
+        // contain a partial turn. Credits and child runs are intentionally
+        // preserved: the provider may have charged the failed attempt and a
+        // child may already have completed work.
+        const stateBeforeAttempt = {
+          agentContext: cloneDeep(currentAgentState.agentContext),
+          messageHistory: cloneDeep(currentAgentState.messageHistory),
+          output: cloneDeep(currentAgentState.output),
+          stepsRemaining: currentAgentState.stepsRemaining,
+          contextTokenCount: currentAgentState.contextTokenCount,
+        }
+
+        try {
+          stepResult = await runAgentStep({
+            ...params,
+
+            agentState: currentAgentState,
+            agentTemplate,
+            extraCodebuffMetadata: {
+              ...(params.extraCodebuffMetadata ?? {}),
+              llm_step_number: String(llmStepNumber),
+              ...(recoveryAttempt > 0 && {
+                recovery_attempt: String(recoveryAttempt),
+              }),
+            },
+            n,
+            prompt: currentPrompt,
+            runId,
+            spawnParams: currentParams,
+            system,
+            tools,
+            additionalToolDefinitions: additionalToolDefinitionsWithCache,
+          })
+          break
+        } catch (error) {
+          const recovery = classifyAgentRecovery(error)
+          if (
+            !recovery.retryable ||
+            recoveryAttempt >= MAX_AGENT_STEP_RECOVERY_ATTEMPTS
+          ) {
+            throw error
+          }
+
+          const creditsUsedAfterFailure = currentAgentState.creditsUsed
+          const directCreditsUsedAfterFailure =
+            currentAgentState.directCreditsUsed
+          const childRunIdsAfterFailure = [...currentAgentState.childRunIds]
+
+          Object.assign(initialAgentState, stateBeforeAttempt, {
+            creditsUsed: creditsUsedAfterFailure,
+            directCreditsUsed: directCreditsUsedAfterFailure,
+            childRunIds: childRunIdsAfterFailure,
+          })
+          currentAgentState = initialAgentState
+
+          recoveryAttempt++
+          const delayMs = getAgentRecoveryDelayMs(recoveryAttempt)
+          currentAgentState.messageHistory = [
+            ...currentAgentState.messageHistory,
+            userMessage({
+              content: withSystemTags(
+                `The previous model request encountered a transient ${recovery.kind} failure. Continue the same task from the preserved work; do not restart completed steps.`,
+              ),
+              tags: ['AGENT_RECOVERY'],
+              keepDuringTruncation: true,
+            }),
+          ]
+
+          logger.warn(
+            {
+              agentType,
+              agentId: currentAgentState.agentId,
+              runId,
+              llmStepNumber,
+              recoveryAttempt,
+              recoveryKind: recovery.kind,
+              statusCode: recovery.statusCode,
+              delayMs,
+            },
+            'Retrying failed agent step after a transient provider error',
+          )
+
+          await (params.waitForAgentRecovery ?? waitForAgentRecovery)({
+            attempt: recoveryAttempt,
+            delayMs,
+            kind: recovery.kind,
+            signal,
+          })
+        }
+      }
+
       const {
         agentState: newAgentState,
         shouldEndTurn: llmShouldEndTurn,
         messageId,
         nResponses: generatedResponses,
-      } = await runAgentStep({
-        ...params,
-
-        agentState: currentAgentState,
-        agentTemplate,
-        extraCodebuffMetadata: {
-          ...(params.extraCodebuffMetadata ?? {}),
-          llm_step_number: String(llmStepNumber),
-        },
-        n,
-        prompt: currentPrompt,
-        runId,
-        spawnParams: currentParams,
-        system,
-        tools,
-        additionalToolDefinitions: additionalToolDefinitionsWithCache,
-      })
+      } = stepResult
 
       if (newAgentState.runId) {
         await addAgentStep({

packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts

@@ -0,0 +1,88 @@
+import { describe, expect, it } from 'bun:test'
+
+import { AbortError } from '@codebuff/common/util/error'
+
+import {
+  classifyAgentRecovery,
+  getAgentRecoveryDelayMs,
+  MAX_AGENT_STEP_RECOVERY_ATTEMPTS,
+} from '../agent-recovery'
+
+describe('classifyAgentRecovery', () => {
+  it('retries transient network failures', () => {
+    const decision = classifyAgentRecovery(
+      Object.assign(new Error('fetch failed'), { code: 'ECONNRESET' }),
+    )
+
+    expect(decision).toEqual({ retryable: true, kind: 'network' })
+  })
+
+  it('retries idle timeouts', () => {
+    const decision = classifyAgentRecovery(
+      Object.assign(new Error('The operation timed out.'), {
+        name: 'TimeoutError',
+      }),
+    )
+
+    expect(decision).toEqual({ retryable: true, kind: 'idle-timeout' })
+  })
+
+  it('retries rate limits and server failures', () => {
+    expect(
+      classifyAgentRecovery(
+        Object.assign(new Error('busy'), { statusCode: 429 }),
+      ),
+    ).toEqual({ retryable: true, kind: 'rate-limit', statusCode: 429 })
+    expect(
+      classifyAgentRecovery(
+        Object.assign(new Error('unavailable'), { status: 503 }),
+      ),
+    ).toEqual({ retryable: true, kind: 'server', statusCode: 503 })
+  })
+
+  it('honors an explicit retryable signal without a status code', () => {
+    expect(
+      classifyAgentRecovery(
+        Object.assign(new Error('provider retry'), { isRetryable: true }),
+      ),
+    ).toEqual({ retryable: true, kind: 'server' })
+  })
+
+  it('does not retry aborts or authentication failures', () => {
+    expect(classifyAgentRecovery(new AbortError())).toEqual({
+      retryable: false,
+      reason: 'aborted',
+    })
+    expect(
+      classifyAgentRecovery(
+        Object.assign(new Error('unauthorized'), {
+          status: 401,
+          isRetryable: true,
+        }),
+      ),
+    ).toEqual({ retryable: false, reason: 'authentication', statusCode: 401 })
+  })
+
+  it('does not retry ordinary client errors', () => {
+    expect(
+      classifyAgentRecovery(
+        Object.assign(new Error('bad request'), { statusCode: 400 }),
+      ),
+    ).toEqual({ retryable: false, reason: 'client', statusCode: 400 })
+  })
+})
+
+describe('getAgentRecoveryDelayMs', () => {
+  it('uses bounded deterministic exponential backoff', () => {
+    expect(getAgentRecoveryDelayMs(0)).toBe(500)
+    expect(getAgentRecoveryDelayMs(1)).toBe(500)
+    expect(getAgentRecoveryDelayMs(2)).toBe(1000)
+    expect(getAgentRecoveryDelayMs(20)).toBe(4000)
+  })
+})
+
+describe('MAX_AGENT_STEP_RECOVERY_ATTEMPTS', () => {
+  it('allows two bounded recovery attempts', () => {
+    expect(MAX_AGENT_STEP_RECOVERY_ATTEMPTS).toBe(2)
+  })
+})