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

@@ -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.

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,
     });
   }
 

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<LLMChatResponse> {
@@ -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,

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,
+  });
 }

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', () => {

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<VideoURLPart>;
 }

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 {

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 {

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',

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', () => {