From 83e257348e466361f539d0d212d87944316e7f52 Mon Sep 17 00:00:00 2001 From: viniciusdebruin Date: Tue, 25 Aug 2026 10:54:19 -0300 Subject: [PATCH 1/2] Add bounded agent-step recovery retries --- .../src/__tests__/loop-agent-steps.test.ts | 51 +++++++ packages/agent-runtime/src/run-agent-step.ts | 133 +++++++++++++++--- .../src/util/__tests__/agent-recovery.test.ts | 81 +++++++++++ .../agent-runtime/src/util/agent-recovery.ts | 84 +++++++++++ 4 files changed, 330 insertions(+), 19 deletions(-) create mode 100644 packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts create mode 100644 packages/agent-runtime/src/util/agent-recovery.ts diff --git a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts index a3d24e38f8..c68c2d7a94 100644 --- a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts +++ b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts @@ -976,6 +976,57 @@ 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, + }) { + llmCallNumber++ + recoveryAttempts.push(extraCodebuffMetadata?.recovery_attempt ?? '0') + + if (llmCallNumber === 1) { + 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) + }) + it('should propagate error code and server message from 403 APICallError responseBody', async () => { const llmOnlyTemplate = { ...mockTemplate, diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index 0b4b0f254e..f2475031e7 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/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 => { + if (delayMs <= 0) return + + await new Promise((resolve, reject) => { + let timeout: ReturnType | 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 - | undefined + ReturnType | 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 spawnParams: Record | undefined startAgentRun: StartAgentRunFn userId: string | undefined @@ -1179,28 +1219,83 @@ export async function loopAgentSteps( const creditsBefore = currentAgentState.directCreditsUsed const childrenBefore = currentAgentState.childRunIds.length llmStepNumber++ + let recoveryAttempt = 0 + let stepResult: Awaited> + + while (true) { + 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 + } + + 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({ diff --git a/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts b/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts new file mode 100644 index 0000000000..0dd753d4f3 --- /dev/null +++ b/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it } from 'bun:test' + +import { AbortError } from '@codebuff/common/util/error' + +import { + classifyAgentRecovery, + getAgentRecoveryDelayMs, +} 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) + }) +}) diff --git a/packages/agent-runtime/src/util/agent-recovery.ts b/packages/agent-runtime/src/util/agent-recovery.ts new file mode 100644 index 0000000000..61f943aa99 --- /dev/null +++ b/packages/agent-runtime/src/util/agent-recovery.ts @@ -0,0 +1,84 @@ +import { + extractApiErrorDetails, + getErrorObject, + isAbortError, + isFetchIdleTimeoutError, + isTransientNetworkError, +} from '@codebuff/common/util/error' + +/** One additional full agent-step attempt after the provider's own retries. */ +export const MAX_AGENT_STEP_RECOVERY_ATTEMPTS = 1 + +export type AgentRecoveryKind = + 'network' | 'idle-timeout' | 'rate-limit' | 'server' + +export type AgentRecoveryDecision = + | { + retryable: true + kind: AgentRecoveryKind + statusCode?: number + } + | { + retryable: false + reason: 'aborted' | 'authentication' | 'client' | 'unknown' + statusCode?: number + } + +const isRetryableStatusCode = (statusCode: number): boolean => + statusCode === 408 || + statusCode === 425 || + statusCode === 429 || + (statusCode >= 500 && statusCode <= 599) + +/** + * Classifies errors that escaped the provider SDK's internal retry loop. + * + * Only failures that are safe to repeat without changing the user's request + * are marked retryable. Abort and authentication failures must reach the + * normal error path immediately. + */ +export function classifyAgentRecovery(error: unknown): AgentRecoveryDecision { + if (isAbortError(error)) { + return { retryable: false, reason: 'aborted' } + } + + if (isFetchIdleTimeoutError(error)) { + return { retryable: true, kind: 'idle-timeout' } + } + + if (isTransientNetworkError(error)) { + return { retryable: true, kind: 'network' } + } + + const statusCode = extractApiErrorDetails(error).statusCode + if (statusCode !== undefined && isRetryableStatusCode(statusCode)) { + return { + retryable: true, + kind: statusCode === 429 ? 'rate-limit' : 'server', + statusCode, + } + } + + if (statusCode === 401 || statusCode === 403) { + return { retryable: false, reason: 'authentication', statusCode } + } + + if (statusCode !== undefined && statusCode >= 400 && statusCode < 500) { + return { retryable: false, reason: 'client', statusCode } + } + + // Some providers mark an error retryable without exposing an HTTP status. + // Trust that explicit signal, but do not infer retryability from an + // otherwise unknown error or override a known client/authentication error. + if (getErrorObject(error).isRetryable === true) { + return { retryable: true, kind: 'server', statusCode } + } + + return { retryable: false, reason: 'unknown' } +} + +/** Deterministic exponential backoff for the outer recovery attempt. */ +export function getAgentRecoveryDelayMs(attempt: number): number { + const normalizedAttempt = Math.max(1, Math.floor(attempt)) + return Math.min(4_000, 500 * 2 ** (normalizedAttempt - 1)) +} From c488f2e5de83264b1736e6edf381c00efd17bb10 Mon Sep 17 00:00:00 2001 From: viniciusdebruin <39091387+viniciusdebruin@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:08:43 -0300 Subject: [PATCH 2/2] fix: preserve state across agent recovery retries --- .../src/__tests__/loop-agent-steps.test.ts | 17 ++++++++++++ packages/agent-runtime/src/run-agent-step.ts | 26 +++++++++++++++++++ .../src/util/__tests__/agent-recovery.test.ts | 7 +++++ .../agent-runtime/src/util/agent-recovery.ts | 4 +-- 4 files changed, 52 insertions(+), 2 deletions(-) diff --git a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts index c68c2d7a94..a8ee896254 100644 --- a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts +++ b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts @@ -990,11 +990,17 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => 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', @@ -1025,6 +1031,17 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => 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 () => { diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index f2475031e7..34f05329b1 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -1223,6 +1223,20 @@ export async function loopAgentSteps( let stepResult: Awaited> 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, @@ -1254,6 +1268,18 @@ export async function loopAgentSteps( 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 = [ diff --git a/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts b/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts index 0dd753d4f3..4f48d09812 100644 --- a/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts +++ b/packages/agent-runtime/src/util/__tests__/agent-recovery.test.ts @@ -5,6 +5,7 @@ import { AbortError } from '@codebuff/common/util/error' import { classifyAgentRecovery, getAgentRecoveryDelayMs, + MAX_AGENT_STEP_RECOVERY_ATTEMPTS, } from '../agent-recovery' describe('classifyAgentRecovery', () => { @@ -79,3 +80,9 @@ describe('getAgentRecoveryDelayMs', () => { 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) + }) +}) diff --git a/packages/agent-runtime/src/util/agent-recovery.ts b/packages/agent-runtime/src/util/agent-recovery.ts index 61f943aa99..1bdeec5aa3 100644 --- a/packages/agent-runtime/src/util/agent-recovery.ts +++ b/packages/agent-runtime/src/util/agent-recovery.ts @@ -6,8 +6,8 @@ import { isTransientNetworkError, } from '@codebuff/common/util/error' -/** One additional full agent-step attempt after the provider's own retries. */ -export const MAX_AGENT_STEP_RECOVERY_ATTEMPTS = 1 +/** Two additional full agent-step attempts after the provider's own retries. */ +export const MAX_AGENT_STEP_RECOVERY_ATTEMPTS = 2 export type AgentRecoveryKind = 'network' | 'idle-timeout' | 'rate-limit' | 'server'