From 4e458d63643a56a2fb1ba9f908c774e56eef1c75 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 27 May 2026 11:20:26 +0800 Subject: [PATCH] refactor: share LLM retry classification (#92) --- .changeset/share-llm-retry-classification.md | 7 ++++++ .../agent-core/src/agent/compaction/full.ts | 18 ++----------- packages/agent-core/src/agent/index.ts | 19 ++++++++++++++ packages/agent-core/src/agent/turn/index.ts | 16 +----------- .../agent-core/src/agent/turn/kosong-llm.ts | 21 +++------------- packages/agent-core/src/loop/llm.ts | 4 +-- packages/kosong/src/errors.ts | 10 ++++++++ packages/kosong/src/index.ts | 1 + packages/kosong/test/errors.test.ts | 25 +++++++++++++++++++ 9 files changed, 71 insertions(+), 50 deletions(-) create mode 100644 .changeset/share-llm-retry-classification.md diff --git a/.changeset/share-llm-retry-classification.md b/.changeset/share-llm-retry-classification.md new file mode 100644 index 000000000..209793cfb --- /dev/null +++ b/.changeset/share-llm-retry-classification.md @@ -0,0 +1,7 @@ +--- +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kosong": patch +"@moonshot-ai/kimi-code": patch +--- + +Use one retry classification for transient LLM failures across regular turns and compaction. diff --git a/packages/agent-core/src/agent/compaction/full.ts b/packages/agent-core/src/agent/compaction/full.ts index bf0f87da5..4f20cb265 100644 --- a/packages/agent-core/src/agent/compaction/full.ts +++ b/packages/agent-core/src/agent/compaction/full.ts @@ -6,10 +6,8 @@ import { toKimiErrorPayload, } from '#/errors'; import { - APIConnectionError, APIEmptyResponseError, - APIStatusError, - APITimeoutError, + isRetryableGenerateError, inputTotal, type GenerateResult, type Message, @@ -470,7 +468,7 @@ export class FullCompaction { const summary = extractCompactionSummary(response); return { response, summary, retryCount }; } catch (error) { - if (attempt >= maxAttempts || !isRetryableCompactionError(error)) { + if (attempt >= maxAttempts || !isRetryableGenerateError(error)) { throw error; } retryCount += 1; @@ -547,15 +545,3 @@ function compactionTelemetryTrigger( } return trigger; } - -function isRetryableCompactionError(error: unknown): boolean { - if (error instanceof APIConnectionError || error instanceof APITimeoutError) { - return true; - } - if (error instanceof APIEmptyResponseError) { - return true; - } - if (!(error instanceof APIStatusError)) return false; - const statusCode = (error as { readonly statusCode?: unknown }).statusCode; - return typeof statusCode === 'number' && [429, 500, 502, 503, 504].includes(statusCode); -} diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index 67feb6e75..b70f943c9 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/packages/agent-core/src/agent/index.ts @@ -50,9 +50,11 @@ import { ToolManager } from './tool/index'; import { TurnFlow } from './turn'; import { GENERATE_REQUEST_LOG_CONTEXT, + KosongLLM, type GenerateOptionsWithRequestLog, } from './turn/kosong-llm'; import { UsageRecorder } from './usage'; +import { resolveCompletionBudget } from '../utils/completion-budget'; export type { AgentRecord, AgentRecordPersistence } from './records'; export type { BuiltinTool, ToolInfo, ToolSource, UserToolRegistration } from './tool'; @@ -178,6 +180,23 @@ export class Agent { }; } + get llm(): KosongLLM { + const model = this.config.model; + const provider = this.config.provider.withThinking(this.config.thinkingLevel); + const loopControl = this.providerManager?.config.loopControl; + const completionBudgetConfig = resolveCompletionBudget({ + reservedContextSize: loopControl?.reservedContextSize, + }); + return new KosongLLM({ + provider, + modelName: model, + systemPrompt: this.config.systemPrompt, + capability: this.config.modelCapabilities, + generate: this.generate, + completionBudgetConfig, + }); + } + private logLlmRequest( provider: ChatProvider, systemPrompt: string, diff --git a/packages/agent-core/src/agent/turn/index.ts b/packages/agent-core/src/agent/turn/index.ts index d8756a788..4f767d625 100644 --- a/packages/agent-core/src/agent/turn/index.ts +++ b/packages/agent-core/src/agent/turn/index.ts @@ -32,11 +32,9 @@ import { import type { AgentEvent, TurnEndedEvent } from '../../rpc'; import type { TelemetryPropertyValue } from '../../telemetry'; import { abortable } from '../../utils/abort'; -import { resolveCompletionBudget } from '../../utils/completion-budget'; import { USER_PROMPT_ORIGIN, type PromptOrigin } from '../context'; import { renderUserPromptHookBlockResult, renderUserPromptHookResult } from '../hooks'; import { canonicalTelemetryArgs, isPlainRecord } from './canonical-args'; -import { KosongLLM } from './kosong-llm'; import { ToolCallDeduplicator } from './tool-dedup'; interface ActiveTurn { @@ -364,24 +362,12 @@ export class TurnFlow { while (true) { signal.throwIfAborted(); const model = this.agent.config.model; - const provider = this.agent.config.provider.withThinking(this.agent.config.thinkingLevel); const loopControl = this.agent.providerManager?.config.loopControl; - const completionBudgetConfig = resolveCompletionBudget({ - reservedContextSize: loopControl?.reservedContextSize, - }); - try { const result = await runTurn({ turnId: String(turnId), signal, - llm: new KosongLLM({ - provider, - modelName: model, - systemPrompt: this.agent.config.systemPrompt, - capability: this.agent.config.modelCapabilities, - generate: this.agent.generate, - completionBudgetConfig, - }), + llm: this.agent.llm, buildMessages: () => this.agent.context.messages, dispatchEvent: this.buildDispatchEvent(turnId), tools: this.agent.tools.loopTools, diff --git a/packages/agent-core/src/agent/turn/kosong-llm.ts b/packages/agent-core/src/agent/turn/kosong-llm.ts index ca522eb82..ef164dcca 100644 --- a/packages/agent-core/src/agent/turn/kosong-llm.ts +++ b/packages/agent-core/src/agent/turn/kosong-llm.ts @@ -16,12 +16,9 @@ */ import { - APIConnectionError, - APIEmptyResponseError, - APIStatusError, - APITimeoutError, emptyUsage, generate as kosongGenerate, + isRetryableGenerateError, type ChatProvider, type GenerateCallbacks, type Message, @@ -81,10 +78,6 @@ export class KosongLLM implements LLM { } async chat(params: LLMChatParams): Promise { - return this.chatOnce(params); - } - - private async chatOnce(params: LLMChatParams): Promise { const callbacks = buildKosongCallbacks(params); // Compute and apply the per-request completion budget against a @@ -121,8 +114,8 @@ export class KosongLLM implements LLM { const response: LLMChatResponse = { toolCalls: [...result.message.toolCalls], - ...(result.finishReason !== null ? { providerFinishReason: result.finishReason } : {}), - ...(result.rawFinishReason !== null ? { rawFinishReason: result.rawFinishReason } : {}), + providerFinishReason: result.finishReason ?? undefined, + rawFinishReason: result.rawFinishReason ?? undefined, usage: result.usage ?? emptyUsage(), }; @@ -130,13 +123,7 @@ export class KosongLLM implements LLM { } isRetryableError(error: unknown): boolean { - if (error instanceof APIConnectionError || error instanceof APITimeoutError) { - return true; - } - if (error instanceof APIEmptyResponseError) { - return true; - } - return error instanceof APIStatusError && [429, 500, 502, 503, 504].includes(error.statusCode); + return isRetryableGenerateError(error); } } diff --git a/packages/agent-core/src/loop/llm.ts b/packages/agent-core/src/loop/llm.ts index df8314593..cd56206e2 100644 --- a/packages/agent-core/src/loop/llm.ts +++ b/packages/agent-core/src/loop/llm.ts @@ -57,8 +57,8 @@ export interface LLMChatParams { export interface LLMChatResponse { toolCalls: ToolCall[]; - providerFinishReason?: FinishReason | undefined; - rawFinishReason?: string | undefined; + providerFinishReason?: FinishReason; + rawFinishReason?: string; usage: TokenUsage; } diff --git a/packages/kosong/src/errors.ts b/packages/kosong/src/errors.ts index 721e2abfe..ba104a2e4 100644 --- a/packages/kosong/src/errors.ts +++ b/packages/kosong/src/errors.ts @@ -64,6 +64,16 @@ export class APIEmptyResponseError extends ChatProviderError { } } +export function isRetryableGenerateError(error: unknown): boolean { + if (error instanceof APIConnectionError || error instanceof APITimeoutError) { + return true; + } + if (error instanceof APIEmptyResponseError) { + return true; + } + return error instanceof APIStatusError && [429, 500, 502, 503, 504].includes(error.statusCode); +} + const CONTEXT_OVERFLOW_MESSAGE_PATTERNS = [ /context[ _-]?length/, /(?:context[ _-]?window.*exceed|exceed.*context[ _-]?window)/, diff --git a/packages/kosong/src/index.ts b/packages/kosong/src/index.ts index d3ac01521..35401dc79 100644 --- a/packages/kosong/src/index.ts +++ b/packages/kosong/src/index.ts @@ -60,6 +60,7 @@ export { APIStatusError, APITimeoutError, ChatProviderError, + isRetryableGenerateError, } from './errors'; /** diff --git a/packages/kosong/test/errors.test.ts b/packages/kosong/test/errors.test.ts index 3c1cc68e0..414b82312 100644 --- a/packages/kosong/test/errors.test.ts +++ b/packages/kosong/test/errors.test.ts @@ -5,6 +5,7 @@ import { APIStatusError, APITimeoutError, ChatProviderError, + isRetryableGenerateError, normalizeAPIStatusError, } from '#/errors'; import { describe, expect, it } from 'vitest'; @@ -84,6 +85,30 @@ describe('APIContextOverflowError', () => { }); }); +describe('isRetryableGenerateError', () => { + it('matches transient provider errors and empty generate responses', () => { + expect(isRetryableGenerateError(new APIConnectionError('conn'))).toBe(true); + expect(isRetryableGenerateError(new APITimeoutError('timeout'))).toBe(true); + expect(isRetryableGenerateError(new APIEmptyResponseError('empty'))).toBe(true); + }); + + it.each([429, 500, 502, 503, 504])('treats HTTP %i as retryable', (statusCode) => { + expect(isRetryableGenerateError(new APIStatusError(statusCode, 'retryable'))).toBe(true); + }); + + it.each([400, 401, 403, 404, 422])('treats HTTP %i as non-retryable', (statusCode) => { + expect(isRetryableGenerateError(new APIStatusError(statusCode, 'non-retryable'))).toBe(false); + }); + + it('does not retry context overflow or unknown errors', () => { + expect( + isRetryableGenerateError(new APIContextOverflowError(400, 'Context length exceeded')), + ).toBe(false); + expect(isRetryableGenerateError(new Error('boom'))).toBe(false); + expect(isRetryableGenerateError('boom')).toBe(false); + }); +}); + describe('error hierarchy instanceof checks', () => { it('all error types are instanceof ChatProviderError', () => { const errors = [