From 76c643bcb6da447c8c47728b4f58512a7a11cfa6 Mon Sep 17 00:00:00 2001 From: qer Date: Fri, 26 Jun 2026 19:12:04 +0800 Subject: [PATCH] feat: cap completion tokens to remaining context window for chat-completions (#1131) * feat: cap completion tokens to remaining context window for chat-completions * test: cover dynamic completion budget for kimi and openai-legacy * fix: leave compaction budget uncapped to avoid one-token summaries * fix: add hookCount to plugins-selector test mocks --- .changeset/dynamic-completion-budget.md | 7 +++++ packages/agent-core/src/agent/index.ts | 1 + .../agent-core/src/agent/turn/kosong-llm.ts | 9 +++++++ .../agent-core/src/utils/completion-budget.ts | 6 ++++- .../test/agent/config-state.test.ts | 6 +++-- packages/kosong/src/provider.ts | 26 ++++++++++++++++++- packages/kosong/src/providers/kimi.ts | 16 ++++++++++-- .../kosong/src/providers/openai-legacy.ts | 24 +++++++++++++++-- packages/kosong/test/kimi.test.ts | 13 ++++++++++ packages/kosong/test/openai-legacy.test.ts | 14 ++++++++++ 10 files changed, 114 insertions(+), 8 deletions(-) create mode 100644 .changeset/dynamic-completion-budget.md diff --git a/.changeset/dynamic-completion-budget.md b/.changeset/dynamic-completion-budget.md new file mode 100644 index 000000000..70e833f8d --- /dev/null +++ b/.changeset/dynamic-completion-budget.md @@ -0,0 +1,7 @@ +--- +"@moonshot-ai/kosong": minor +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kimi-code": patch +--- + +Cap completion tokens to the remaining context window for chat-completions providers, avoiding context-overflow and invalid max_tokens errors. diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index ff49b36be..4e733a80c 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/packages/agent-core/src/agent/index.ts @@ -244,6 +244,7 @@ export class Agent { capability: this.config.modelCapabilities, generate: this.generate, completionBudgetConfig, + usedContextTokens: () => this.context.tokenCount, }); } diff --git a/packages/agent-core/src/agent/turn/kosong-llm.ts b/packages/agent-core/src/agent/turn/kosong-llm.ts index fb6b34074..edd501d7e 100644 --- a/packages/agent-core/src/agent/turn/kosong-llm.ts +++ b/packages/agent-core/src/agent/turn/kosong-llm.ts @@ -55,6 +55,12 @@ export interface KosongLLMConfig { * final cap is applied to each request. */ readonly completionBudgetConfig?: CompletionBudgetConfig | undefined; + /** + * Returns the number of context tokens already consumed by the latest + * completed step (API-reported input + output). Used by chat-completions + * providers to size the completion budget to the remaining context window. + */ + readonly usedContextTokens?: (() => number) | undefined; } export class KosongLLM implements LLM { @@ -65,6 +71,7 @@ export class KosongLLM implements LLM { private readonly provider: ChatProvider; private readonly generate: GenerateFn; private readonly completionBudgetConfig: CompletionBudgetConfig | undefined; + private readonly usedContextTokens: (() => number) | undefined; constructor(config: KosongLLMConfig) { this.provider = config.provider; @@ -73,6 +80,7 @@ export class KosongLLM implements LLM { this.capability = config.capability; this.generate = config.generate ?? kosongGenerate; this.completionBudgetConfig = config.completionBudgetConfig; + this.usedContextTokens = config.usedContextTokens; } async chat(params: LLMChatParams): Promise { @@ -98,6 +106,7 @@ export class KosongLLM implements LLM { provider: this.provider, budget: this.completionBudgetConfig, capability: this.capability, + usedContextTokens: this.usedContextTokens?.(), }); const options: GenerateOptionsWithRequestLogFields = { signal: params.signal, diff --git a/packages/agent-core/src/utils/completion-budget.ts b/packages/agent-core/src/utils/completion-budget.ts index ceb086ef2..55e62a6ca 100644 --- a/packages/agent-core/src/utils/completion-budget.ts +++ b/packages/agent-core/src/utils/completion-budget.ts @@ -79,6 +79,7 @@ export function applyCompletionBudget(args: { readonly provider: ChatProvider; readonly budget: CompletionBudgetConfig | undefined; readonly capability: ModelCapability | undefined; + readonly usedContextTokens?: number; }): ChatProvider { if (args.budget === undefined) return args.provider; if (args.provider.withMaxCompletionTokens === undefined) return args.provider; @@ -86,5 +87,8 @@ export function applyCompletionBudget(args: { budget: args.budget, capability: args.capability, }); - return args.provider.withMaxCompletionTokens(cap); + return args.provider.withMaxCompletionTokens(cap, { + usedContextTokens: args.usedContextTokens, + maxContextTokens: args.capability?.max_context_tokens, + }); } diff --git a/packages/agent-core/test/agent/config-state.test.ts b/packages/agent-core/test/agent/config-state.test.ts index ddccdef52..f200108a7 100644 --- a/packages/agent-core/test/agent/config-state.test.ts +++ b/packages/agent-core/test/agent/config-state.test.ts @@ -74,7 +74,7 @@ describe('ConfigState model capabilities', () => { }); }); - it('uses model max output size as the LLM completion cap', async () => { + it('clamps the LLM completion cap to 128k for openai-compatible providers', async () => { let requestMaxTokens: unknown; const ctx = testAgent({ generate: async (provider) => { @@ -121,7 +121,9 @@ describe('ConfigState model capabilities', () => { signal: new AbortController().signal, }); - expect(requestMaxTokens).toBe(384000); + // maxOutputSize (384000) is clamped to the 128k ceiling applied to + // non-Kimi chat-completions providers. + expect(requestMaxTokens).toBe(131072); }); it('uses session id as a provider prompt cache hint without storing it on Agent', () => { diff --git a/packages/kosong/src/provider.ts b/packages/kosong/src/provider.ts index becf996b0..1782f03e4 100644 --- a/packages/kosong/src/provider.ts +++ b/packages/kosong/src/provider.ts @@ -13,6 +13,22 @@ import type { TokenUsage } from './usage'; */ export type ThinkingEffort = 'off' | 'low' | 'medium' | 'high' | 'xhigh' | 'max'; +/** + * Optional context passed to {@link ChatProvider.withMaxCompletionTokens} so a + * provider can tighten the caller-supplied cap to its own transport + * constraints. + */ +export interface MaxCompletionTokensOptions { + /** + * Tokens already consumed by the current context (API-reported input + + * output of the latest completed step). Chat-completions providers use it + * to size the cap to the remaining context window. + */ + readonly usedContextTokens?: number; + /** Model context-window size in tokens (`max_context_size`). */ + readonly maxContextTokens?: number; +} + /** * Normalized finish-reason signal indicating why a generation stopped. * @@ -155,11 +171,19 @@ export interface ChatProvider { * budget clamped to `maxCompletionTokens`. Optional because not every * backend benefits from a client-computed cap. * + * When `options` are provided, implementations may further tighten the cap + * based on their own transport constraints — e.g. chat-completions + * endpoints size the cap to the remaining context window + * (`maxContextTokens - usedContextTokens`) and/or clamp to a fixed ceiling. + * * Implementations MUST NOT mutate or replace internal HTTP clients on the * returned clone — the clone is expected to share transport state with the * original. See `KimiChatProvider._clone()` for the rationale. */ - withMaxCompletionTokens?(maxCompletionTokens: number): ChatProvider; + withMaxCompletionTokens?( + maxCompletionTokens: number, + options?: MaxCompletionTokensOptions, + ): ChatProvider; /** Upload a video and return a content part that can be sent to this provider. */ uploadVideo?(input: string | VideoUploadInput, options?: GenerateOptions): Promise; } diff --git a/packages/kosong/src/providers/kimi.ts b/packages/kosong/src/providers/kimi.ts index c37c2d2a4..3a120c078 100644 --- a/packages/kosong/src/providers/kimi.ts +++ b/packages/kosong/src/providers/kimi.ts @@ -4,6 +4,7 @@ import type { ChatProvider, FinishReason, GenerateOptions, + MaxCompletionTokensOptions, ProviderRequestAuth, StreamedMessage, ThinkingEffort, @@ -527,8 +528,19 @@ export class KimiChatProvider implements ChatProvider { return this._withGenerationKwargs(kwargs); } - withMaxCompletionTokens(maxCompletionTokens: number): KimiChatProvider { - return this._withGenerationKwargs({ max_completion_tokens: maxCompletionTokens }); + withMaxCompletionTokens( + maxCompletionTokens: number, + options?: MaxCompletionTokensOptions, + ): KimiChatProvider { + let cap = maxCompletionTokens; + if ( + options?.usedContextTokens !== undefined && + options?.maxContextTokens !== undefined && + options.maxContextTokens > 0 + ) { + cap = Math.min(cap, options.maxContextTokens - options.usedContextTokens); + } + return this._withGenerationKwargs({ max_completion_tokens: Math.max(1, cap) }); } withExtraBody(extraBody: ExtraBody): KimiChatProvider { diff --git a/packages/kosong/src/providers/openai-legacy.ts b/packages/kosong/src/providers/openai-legacy.ts index c9987a5f5..6caa55e24 100644 --- a/packages/kosong/src/providers/openai-legacy.ts +++ b/packages/kosong/src/providers/openai-legacy.ts @@ -3,6 +3,7 @@ import type { ChatProvider, FinishReason, GenerateOptions, + MaxCompletionTokensOptions, ProviderRequestAuth, StreamedMessage, ThinkingEffort, @@ -46,6 +47,13 @@ import { // arms can be overridden by an explicit `reasoningKey` on the provider config. const KNOWN_REASONING_KEYS = ['reasoning_content', 'reasoning_details', 'reasoning'] as const; const DEFAULT_OUTBOUND_REASONING_KEY = KNOWN_REASONING_KEYS[0]; + +/** + * Hard upper bound on `max_tokens` for OpenAI-compatible chat-completions + * endpoints. Many third-party providers reject `max_tokens` above this limit + * (the documented range is `[1, 131072]`). + */ +const CHAT_COMPLETIONS_MAX_OUTPUT_TOKENS_CEILING = 128 * 1024; const OPENAI_CHAT_TOOL_CALL_ID_POLICY: ToolCallIdPolicy = { normalize: (id) => sanitizeToolCallId(id, 64), maxLength: 64, @@ -586,8 +594,20 @@ export class OpenAILegacyChatProvider implements ChatProvider { return clone; } - withMaxCompletionTokens(maxCompletionTokens: number): OpenAILegacyChatProvider { - return this.withGenerationKwargs(completionTokenKwargs(this._model, maxCompletionTokens)); + withMaxCompletionTokens( + maxCompletionTokens: number, + options?: MaxCompletionTokensOptions, + ): OpenAILegacyChatProvider { + let cap = maxCompletionTokens; + if ( + options?.usedContextTokens !== undefined && + options?.maxContextTokens !== undefined && + options.maxContextTokens > 0 + ) { + cap = Math.min(cap, options.maxContextTokens - options.usedContextTokens); + } + cap = Math.min(cap, CHAT_COMPLETIONS_MAX_OUTPUT_TOKENS_CEILING); + return this.withGenerationKwargs(completionTokenKwargs(this._model, Math.max(1, cap))); } private _clone(): OpenAILegacyChatProvider { diff --git a/packages/kosong/test/kimi.test.ts b/packages/kosong/test/kimi.test.ts index f309b2030..e0ffcf249 100644 --- a/packages/kosong/test/kimi.test.ts +++ b/packages/kosong/test/kimi.test.ts @@ -641,6 +641,19 @@ describe('KimiChatProvider', () => { expect(body['max_tokens']).toBeUndefined(); }); + it('withMaxCompletionTokens sizes the cap to the remaining context window', async () => { + const provider = createProvider().withMaxCompletionTokens(100000, { + usedContextTokens: 30000, + maxContextTokens: 100000, + }); + const history: Message[] = [ + { role: 'user', content: [{ type: 'text', text: 'Hi' }], toolCalls: [] }, + ]; + const body = await captureRequestBody(provider, '', [], history); + + expect(body['max_completion_tokens']).toBe(70000); + }); + it('passes constructor generation kwargs into the request body', async () => { const provider = new KimiChatProvider({ model: 'kimi-k2-turbo-preview', diff --git a/packages/kosong/test/openai-legacy.test.ts b/packages/kosong/test/openai-legacy.test.ts index 8335a8228..f02efd048 100644 --- a/packages/kosong/test/openai-legacy.test.ts +++ b/packages/kosong/test/openai-legacy.test.ts @@ -641,6 +641,20 @@ describe('OpenAILegacyChatProvider', () => { expect(body['max_tokens']).toBe(1024); expect(body['max_completion_tokens']).toBeUndefined(); }); + + it('withMaxCompletionTokens clamps to the 128k ceiling', async () => { + const provider = createProvider().withMaxCompletionTokens(1000000, { + usedContextTokens: 30000, + maxContextTokens: 1000000, + }); + const history: Message[] = [ + { role: 'user', content: [{ type: 'text', text: 'Hi' }], toolCalls: [] }, + ]; + const body = await captureRequestBody(provider, '', [], history); + + // 1000000 - 30000 = 970000, clamped to 131072 + expect(body['max_tokens']).toBe(131072); + }); }); describe('maxTokens option', () => {