From 54365f44240fae8862244622399b1a4906d4f02b Mon Sep 17 00:00:00 2001 From: Acters Date: Wed, 29 Jul 2026 18:05:57 -0700 Subject: [PATCH 1/2] =?UTF-8?q?Kimi=20OAuth:=20auto=20context=20tier=20for?= =?UTF-8?q?=20k3=20(k3-256k=20=E2=86=92=20k3=20wire=20upgrade)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Kimi coding endpoint serves K3 as k3-256k (256k context, cheaper) and k3 (1M context, stated to consume 2x the quota — not independently verified). Switching from k3-256k to k3 seamlessly upgrades the prompt cache server-side. 256k is large enough for most tasks, so exposing a single user-facing k3 model and automatically upgrading the wire model ID once the session context passes the 256k cutoff (256k window minus the default compaction reserve) keeps the bulk of sessions on the cheaper variant and avoids compaction on long-horizon tasks. - Add Model.wireModelId; openai-completions sends it on the wire while recording the registry id on assistant messages - New k3-context-tier module derives the effective model (256k tier with wireModelId k3-256k, or 1M tier) from current context tokens - AgentSession applies the tier on model set/cycle/refresh/restore and session creation, and upgrades instead of compacting on both the threshold and overflow paths; upgrade runs even when compaction is disabled since no context reduction is involved - Cutoff is fixed at the default compaction threshold of the 256k window, so a user-lowered compaction threshold compacts first and effectively disables the upgrade - Emit context_window_upgrade session event; interactive mode shows a status line and the footer reflects the new window - OAuth discovery never surfaces k3-256k as a separate selectable model --- .../ai/src/providers/openai-completions.ts | 3 +- packages/ai/src/types.ts | 7 + packages/ai/src/utils/oauth/kimi-coding.ts | 5 +- packages/ai/test/kimi-coding-oauth.test.ts | 17 ++ .../ai/test/openai-completions-kimi.test.ts | 42 +++ packages/coding-agent/docs/providers.md | 1 + packages/coding-agent/docs/rpc.md | 1 + .../coding-agent/src/core/agent-session.ts | 95 ++++++- .../coding-agent/src/core/k3-context-tier.ts | 97 +++++++ packages/coding-agent/src/core/sdk.ts | 5 +- .../src/modes/interactive/interactive-mode.ts | 10 + .../test/k3-context-tier-session.test.ts | 261 ++++++++++++++++++ .../coding-agent/test/k3-context-tier.test.ts | 106 +++++++ 13 files changed, 639 insertions(+), 11 deletions(-) create mode 100644 packages/coding-agent/src/core/k3-context-tier.ts create mode 100644 packages/coding-agent/test/k3-context-tier-session.test.ts create mode 100644 packages/coding-agent/test/k3-context-tier.test.ts diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 517a786c..2ef1bfff 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -375,7 +375,8 @@ function buildParams(model: Model<"openai-completions">, context: Context, optio maybeAddOpenRouterAnthropicCacheControl(model, messages); const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { - model: model.id, + // Context-tiered models (e.g. Kimi K3) send a cheaper variant id on the wire. + model: model.wireModelId ?? model.id, messages, stream: true, }; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 8f92e410..e656b2ec 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -346,6 +346,13 @@ export interface Model { contextWindow: number; maxTokens: number; headers?: Record; + /** + * Model identifier sent to the API when it differs from the registry id. + * Used by context-tiered models such as Kimi K3, where dreb exposes a + * single user-facing model (`k3`) but sends the cheaper 256k variant + * (`k3-256k`) on the wire until the session context outgrows it. + */ + wireModelId?: string; /** Explicit auth transport. When unset, the provider's default and token heuristics apply. */ authMode?: AuthMode; /** Compatibility overrides for OpenAI-compatible APIs. If not set, auto-detected from baseUrl. */ diff --git a/packages/ai/src/utils/oauth/kimi-coding.ts b/packages/ai/src/utils/oauth/kimi-coding.ts index 75db0642..d2719b38 100644 --- a/packages/ai/src/utils/oauth/kimi-coding.ts +++ b/packages/ai/src/utils/oauth/kimi-coding.ts @@ -811,8 +811,11 @@ export const kimiCodingOAuthProvider: OAuthProviderInterface = { // The official client treats only the explicit "anthropic" protocol as a // separate wire format; absent and future values use the default Kimi route. + // `k3-256k` is the cheaper 256k wire variant of `k3`; dreb surfaces both + // through the single user-facing `k3` model (auto context tier), so it + // must not appear as a separate selectable model. const supportedDiscovered = discovered.filter( - (info) => info.supports_tool_use !== false && info.protocol !== "anthropic", + (info) => info.supports_tool_use !== false && info.protocol !== "anthropic" && info.id !== "k3-256k", ); const result: Model[] = []; const seen = new Set(); diff --git a/packages/ai/test/kimi-coding-oauth.test.ts b/packages/ai/test/kimi-coding-oauth.test.ts index a511ab47..5da0eb24 100644 --- a/packages/ai/test/kimi-coding-oauth.test.ts +++ b/packages/ai/test/kimi-coding-oauth.test.ts @@ -1515,6 +1515,23 @@ describe("Kimi For Coding OAuth", () => { expect(result.filter((m) => m.provider === "kimi-coding-oauth").map((m) => m.id)).toEqual(["k3"]); }); + it("never surfaces the k3-256k wire variant as a separate selectable model", () => { + const creds: OAuthCredentials & { models?: KimiModelInfo[] } = { + refresh: "r", + access: "a", + expires: Date.now() + 120_000, + models: [ + { id: "k3", display_name: "K3", context_length: 1048576 }, + { id: "k3-256k", display_name: "K3 256k", context_length: 262144 }, + ], + }; + + const result = kimiCodingOAuthProvider.modifyModels!(buildStaticModels(), creds); + const ids = result.filter((m) => m.provider === "kimi-coding-oauth").map((m) => m.id); + + expect(ids).toEqual(["k3"]); + }); + it("preserves static IDs and appends an unknown legacy model", () => { const creds: OAuthCredentials & { modelId?: string; diff --git a/packages/ai/test/openai-completions-kimi.test.ts b/packages/ai/test/openai-completions-kimi.test.ts index 147f74bd..9ff57efe 100644 --- a/packages/ai/test/openai-completions-kimi.test.ts +++ b/packages/ai/test/openai-completions-kimi.test.ts @@ -452,3 +452,45 @@ describe("openai-completions kimi thinkingFormat", () => { ]); }); }); + +describe("openai-completions wireModelId", () => { + beforeEach(() => { + mockState.lastParams = undefined; + mockState.chunks = undefined; + }); + + it("sends wireModelId on the wire while recording the registry model id", async () => { + const model: Model<"openai-completions"> = { + ...KIMI_MODEL, + provider: "kimi-coding-oauth", + id: "k3", + wireModelId: "k3-256k", + contextWindow: 262144, + }; + + const result = await streamSimple( + model, + { messages: [{ role: "user", content: "Hi", timestamp: Date.now() }] }, + { apiKey: "test" }, + ).result(); + + expect((mockState.lastParams as { model?: string }).model).toBe("k3-256k"); + expect(result.model).toBe("k3"); + }); + + it("sends the registry model id when wireModelId is unset", async () => { + const model: Model<"openai-completions"> = { + ...KIMI_MODEL, + provider: "kimi-coding-oauth", + id: "k3", + }; + + await streamSimple( + model, + { messages: [{ role: "user", content: "Hi", timestamp: Date.now() }] }, + { apiKey: "test" }, + ).result(); + + expect((mockState.lastParams as { model?: string }).model).toBe("k3"); + }); +}); diff --git a/packages/coding-agent/docs/providers.md b/packages/coding-agent/docs/providers.md index f73f4b25..0ab91fc2 100644 --- a/packages/coding-agent/docs/providers.md +++ b/packages/coding-agent/docs/providers.md @@ -52,6 +52,7 @@ Use `/logout` to clear credentials. Tokens are stored in `~/.dreb/agent/auth.jso - `/login` uses the Kimi Code OAuth subscription endpoint at `https://api.kimi.com/coding/v1`. - `KIMI_API_KEY` uses Kimi For Coding's Anthropic-compatible API at `https://api.kimi.com/coding`. - Built-in OAuth models are `kimi-for-coding` (default, 262k context), `k3` (1M context), and `kimi-for-coding-highspeed` (262k context). Model availability is plan-dependent: the `k3` 1M-context ID and the `kimi-for-coding-highspeed` ID are only exposed when the subscription includes them, and `kimi-for-coding-highspeed` runs at roughly 6× speed for 3× quota usage. On login/refresh, dreb asks the Kimi API which models the subscription is entitled to and updates context, reasoning, image, tool-use, protocol, and thinking-effort metadata. Compatible newly discovered IDs are templated conservatively; if discovery fails, the static list remains available. +- The OAuth `k3` model is context-tiered: the Kimi endpoint serves it as `k3-256k` (256k context, cheaper) and `k3` (1M context, stated to consume 2× the quota), and upgrading does not invalidate the prompt cache. The cheaper variant is exclusive to the Kimi for Coding OAuth endpoint — the pay-per-token Moonshot AI Platform does not expose it. dreb starts every session on the `k3-256k` wire ID and automatically upgrades to `k3` once the session context passes the 256k cutoff (256k minus the default compaction reserve) instead of compacting. A user-lowered compaction threshold compacts before the cutoff and thus effectively disables the upgrade. - OAuth requests use the current Kimi Code device identity contract and share its stable `~/.kimi-code/device_id`. Extra low-precedence headers can be supplied with newline-separated `KIMI_CODE_CUSTOM_HEADERS` values. - The OAuth coding endpoint accepts OpenAI-style multimodal content arrays with base64 `image_url` data URLs. dreb keeps a conservative 32k output-token cap because the managed model catalog does not currently advertise a per-model output limit. - Moonshot Open Platform uses a different base URL (`https://api.moonshot.ai/v1`); don't assume both routes expose identical behavior. diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index fdc24aad..e0cac857 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -1598,6 +1598,7 @@ Response: | `length_retry` | Response hit the token limit; retrying with a larger budget | | `auto_compaction_start` | Auto-compaction begins | | `auto_compaction_end` | Auto-compaction completes | +| `context_window_upgrade` | Wire model tier auto-upgraded (e.g. Kimi K3 256k → 1M); includes `provider`, `modelId`, `fromContextWindow`, `toContextWindow` | | `auto_retry_start` | Auto-retry begins (after transient error) | | `auto_retry_end` | Auto-retry completes (success or final failure) | | `background_agent_start` | Background subagent launched (includes `sessionDir`) | diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index ae2d6311..e7792a7b 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -66,6 +66,14 @@ import { } from "./extensions/index.js"; import { checkScriptContent, extractScriptPaths, isForbiddenCommand } from "./forbidden-commands.js"; import { type GitRepoState, getGitRepoState } from "./git-repo-state.js"; +import { + deriveK3ContextTierModel, + isK3256kTier, + K3_1M_CONTEXT_WINDOW, + K3_256K_CONTEXT_WINDOW, + K3_UPGRADE_CUTOFF_TOKENS, + shouldUpgradeK3Tier, +} from "./k3-context-tier.js"; import { log } from "./logger.js"; import type { BashExecutionMessage, CustomMessage } from "./messages.js"; import type { ModelRegistry } from "./model-registry.js"; @@ -152,7 +160,15 @@ export type AgentSessionEvent = | { type: "parent_paused_for_background_agents"; runningAgentCount: number; turnsUsed: number; turnLimit: number } | { type: "session_name_changed"; name: string } | { type: "tasks_update"; tasks: readonly SessionTask[] } - | { type: "suggest_next"; command: string }; + | { type: "suggest_next"; command: string } + | { + /** The wire model tier auto-upgraded because the session context outgrew the smaller tier (e.g. Kimi K3 256k → 1M). */ + type: "context_window_upgrade"; + provider: string; + modelId: string; + fromContextWindow: number; + toContextWindow: number; + }; /** Listener function for agent session events */ export type AgentSessionEventListener = (event: AgentSessionEvent) => void; @@ -2017,7 +2033,7 @@ export class AgentSession { const previousModel = this.model; const thinkingLevel = this._getThinkingLevelForModelSwitch(); - this.agent.setModel(model); + this.agent.setModel(this._applyContextTier(model)); this._refreshThinkingDisplay(model); this.sessionManager.appendModelChange(model.provider, model.id); this.settingsManager.setDefaultModelAndProvider(model.provider, model.id); @@ -2092,7 +2108,7 @@ export class AgentSession { const thinkingLevel = this._getThinkingLevelForModelSwitch(next.thinkingLevel); // Apply model - this.agent.setModel(next.model); + this.agent.setModel(this._applyContextTier(next.model)); this._refreshThinkingDisplay(next.model); this.sessionManager.appendModelChange(next.model.provider, next.model.id); this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id); @@ -2128,7 +2144,7 @@ export class AgentSession { } const thinkingLevel = this._getThinkingLevelForModelSwitch(); - this.agent.setModel(nextModel); + this.agent.setModel(this._applyContextTier(nextModel)); this._refreshThinkingDisplay(nextModel); this.sessionManager.appendModelChange(nextModel.provider, nextModel.id); this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id); @@ -2405,12 +2421,11 @@ export class AgentSession { */ private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise { const settings = this.settingsManager.getCompactionSettings(); - if (!settings.enabled) return; // Skip if message was aborted (user cancelled) - unless skipAbortedCheck is false if (skipAbortedCheck && assistantMessage.stopReason === "aborted") return; - const contextWindow = this.model?.contextWindow ?? 0; + let contextWindow = this.model?.contextWindow ?? 0; // Skip overflow check if the message came from a different model. // This handles the case where user switched from a smaller-context model (e.g. opus) @@ -2431,6 +2446,29 @@ export class AgentSession { // Case 1: Overflow - LLM returned context overflow error if (sameModel && isContextOverflow(assistantMessage, contextWindow)) { + // K3 auto context tier: an overflow in the 256k tier upgrades to the + // 1M tier instead of compacting — the Kimi backend grows the prompt + // cache seamlessly. This runs even when compaction is disabled since + // no context reduction is involved. + if (this._tryUpgradeK3ContextTier()) { + // Remove the error message from agent state (it IS saved to session + // for history, but we don't want it in context for the retry) + const messages = this.agent.state.messages; + if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { + this.agent.replaceMessages(messages.slice(0, -1)); + } + setTimeout(() => { + this.agent.continue().catch((err) => { + this.warnInSession( + `Agent failed to continue after context window upgrade: ${err instanceof Error ? err.message : String(err)}`, + ); + }); + }, 100); + return; + } + + if (!settings.enabled) return; + if (this._overflowRecoveryAttempted) { this._emit({ type: "auto_compaction_end", @@ -2455,6 +2493,8 @@ export class AgentSession { } // Case 2: Threshold - context is getting large + if (!settings.enabled) return; + // For error messages (no usage data), estimate from last successful response. // This ensures sessions that hit persistent API errors (e.g. 529) can still compact. let contextTokens: number; @@ -2477,11 +2517,50 @@ export class AgentSession { } else { contextTokens = calculateContextTokens(assistantMessage.usage); } + + // K3 auto context tier: reaching the 256k cutoff upgrades to the 1M tier + // instead of compacting. The cutoff is fixed at the default compaction + // threshold of the 256k window, so a user-lowered compaction threshold + // fires first and effectively disables the upgrade. + if (shouldUpgradeK3Tier(this.model, contextTokens) && this._tryUpgradeK3ContextTier()) { + contextWindow = this.model?.contextWindow ?? contextWindow; + } + if (shouldCompact(contextTokens, contextWindow, settings)) { await this._runAutoCompaction("threshold", false); } } + /** + * Apply the K3 auto context tier to a model being set on the agent. The + * user-facing `k3` model runs on the cheaper `k3-256k` wire model ID until + * the session context grows past the 256k cutoff; no-op for other models. + * See k3-context-tier.ts. + */ + private _applyContextTier(model: Model): Model { + return deriveK3ContextTierModel(model, estimateContextTokens(this.agent.state.messages).tokens); + } + + /** + * Upgrade the K3 auto context tier from 256k to 1M. The Kimi backend + * upgrades the prompt cache seamlessly, so no context is lost or + * compacted. Returns true when the upgrade was applied. + */ + private _tryUpgradeK3ContextTier(): boolean { + const model = this.model; + if (!model || !isK3256kTier(model)) return false; + const upgraded = deriveK3ContextTierModel(model, K3_UPGRADE_CUTOFF_TOKENS + 1); + this.agent.setModel(upgraded); + this._emit({ + type: "context_window_upgrade", + provider: upgraded.provider, + modelId: upgraded.id, + fromContextWindow: K3_256K_CONTEXT_WINDOW, + toContextWindow: K3_1M_CONTEXT_WINDOW, + }); + return true; + } + /** * Internal: Run auto-compaction with events. */ @@ -2742,7 +2821,7 @@ export class AgentSession { return; } - this.agent.setModel(refreshedModel); + this.agent.setModel(this._applyContextTier(refreshedModel)); this._refreshThinkingDisplay(refreshedModel); } @@ -3387,7 +3466,7 @@ export class AgentSession { (m) => m.provider === sessionContext.model!.provider && m.id === sessionContext.model!.modelId, ); if (match) { - this.agent.setModel(match); + this.agent.setModel(this._applyContextTier(match)); this._refreshThinkingDisplay(match); await this._emitModelSelect(match, previousModel, "restore"); } diff --git a/packages/coding-agent/src/core/k3-context-tier.ts b/packages/coding-agent/src/core/k3-context-tier.ts new file mode 100644 index 00000000..5ebc5f38 --- /dev/null +++ b/packages/coding-agent/src/core/k3-context-tier.ts @@ -0,0 +1,97 @@ +/** + * Kimi K3 auto context tier. + * + * The Kimi coding endpoint serves K3 under two model IDs: + * - `k3-256k`: 256k context window (cheaper) + * - `k3`: 1M context window (stated to consume 2x the quota) + * + * The pricing/quota terms are as stated by the Kimi team and not + * independently verified, but they are the rationale for auto-switching: + * 256k is still large enough for most tasks, so the first 256k of every + * session can run on the cheaper model and only sessions that genuinely + * outgrow it pay the 1M premium. + * + * `k3-256k` is exclusive to the Kimi for Coding OAuth endpoint — the + * pay-per-token Moonshot AI Platform does not expose the cheaper variant — + * so this applies solely to the `kimi-coding-oauth` provider. + * + * Per the Kimi backend team, switching from `k3-256k` to `k3` does not + * invalidate the prompt cache — the cache seamlessly upgrades from 256k to + * 1M. dreb therefore exposes a single user-selectable `k3` model and + * automatically upgrades the wire model ID once the session context grows + * past the 256k cutoff, avoiding context compaction on long-horizon tasks. + * + * The upgrade cutoff is the 256k window minus the DEFAULT compaction reserve, + * i.e. the point where auto-compaction would trigger under default settings + * for a 256k-window model. Users who lower their compaction threshold compact + * before the cutoff is ever reached, which effectively disables the upgrade. + */ + +import type { Api, Model } from "@dreb/ai"; +import { DEFAULT_COMPACTION_SETTINGS } from "./compaction/compaction.js"; + +/** Provider and user-facing model ID the tier logic applies to. */ +export const K3_PROVIDER = "kimi-coding-oauth"; +export const K3_MODEL_ID = "k3"; + +/** Wire model ID sent while in the cheaper 256k tier. */ +export const K3_256K_WIRE_MODEL_ID = "k3-256k"; + +/** Context window of the 256k tier. */ +export const K3_256K_CONTEXT_WINDOW = 262144; + +/** Context window of the 1M tier. */ +export const K3_1M_CONTEXT_WINDOW = 1048576; + +/** + * Context token count at which the wire model ID upgrades from `k3-256k` to + * `k3`: the 256k window minus the default compaction reserve. + */ +export const K3_UPGRADE_CUTOFF_TOKENS = K3_256K_CONTEXT_WINDOW - DEFAULT_COMPACTION_SETTINGS.reserveTokens; + +/** Whether the model is the user-facing Kimi K3 model subject to auto context tiers. */ +export function isK3Model(model: Model | null | undefined): boolean { + return model?.provider === K3_PROVIDER && model?.id === K3_MODEL_ID; +} + +/** Whether the model is currently in the 256k tier (wire model ID `k3-256k`). */ +export function isK3256kTier(model: Model | null | undefined): boolean { + return isK3Model(model) && model?.wireModelId === K3_256K_WIRE_MODEL_ID; +} + +/** + * Derive the effective model for the current context size. + * + * At or below the cutoff: 256k context window with wire model ID `k3-256k`. + * Above the cutoff: 1M context window with the registry model ID `k3` sent + * on the wire (the Kimi backend upgrades the cache seamlessly). + * + * Returns the input unchanged for non-K3 models and is idempotent for + * already-derived models. + */ +export function deriveK3ContextTierModel(model: Model, contextTokens: number): Model; +export function deriveK3ContextTierModel( + model: Model | undefined, + contextTokens: number, +): Model | undefined; +export function deriveK3ContextTierModel( + model: Model | undefined, + contextTokens: number, +): Model | undefined { + if (!model || !isK3Model(model)) return model; + if (contextTokens > K3_UPGRADE_CUTOFF_TOKENS) { + if (model.wireModelId === undefined && model.contextWindow === K3_1M_CONTEXT_WINDOW) return model; + const { wireModelId: _droppedWireModelId, ...rest } = model; + return { ...rest, contextWindow: K3_1M_CONTEXT_WINDOW }; + } + if (model.wireModelId === K3_256K_WIRE_MODEL_ID && model.contextWindow === K3_256K_CONTEXT_WINDOW) return model; + return { ...model, contextWindow: K3_256K_CONTEXT_WINDOW, wireModelId: K3_256K_WIRE_MODEL_ID }; +} + +/** + * Whether the session should upgrade from the 256k tier to the 1M tier. + * Only true while in the 256k tier and past the cutoff. + */ +export function shouldUpgradeK3Tier(model: Model | null | undefined, contextTokens: number): boolean { + return isK3256kTier(model) && contextTokens > K3_UPGRADE_CUTOFF_TOKENS; +} diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index d4f274c6..b8da293c 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -6,6 +6,7 @@ import { AgentSession } from "./agent-session.js"; import { AuthStorage } from "./auth-storage.js"; import { DEFAULT_THINKING_LEVEL } from "./defaults.js"; import type { ExtensionRunner, LoadExtensionsResult, ToolDefinition } from "./extensions/index.js"; +import { deriveK3ContextTierModel } from "./k3-context-tier.js"; import { convertToLlm } from "./messages.js"; import { ModelRegistry } from "./model-registry.js"; import { findInitialModel } from "./model-resolver.js"; @@ -327,7 +328,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} agent = new Agent({ initialState: { systemPrompt: "", - model, + // New sessions start with empty context, so K3 starts in the + // cheaper 256k wire tier (upgrades automatically past the cutoff). + model: deriveK3ContextTierModel(model, 0), thinkingLevel, tools: [], }, diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 4a531408..60a74466 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -2814,6 +2814,16 @@ export class InteractiveMode { } break; } + + case "context_window_upgrade": { + // Auto context tier upgrade (e.g. Kimi K3 256k → 1M); cache is + // preserved server-side and no compaction occurred. + this.showStatus( + `${event.modelId} context window upgraded: ${Math.round(event.fromContextWindow / 1024)}k → ${Math.round(event.toContextWindow / 1048576)}M (cache preserved)`, + ); + this.footer.invalidate(); + break; + } } } diff --git a/packages/coding-agent/test/k3-context-tier-session.test.ts b/packages/coding-agent/test/k3-context-tier-session.test.ts new file mode 100644 index 00000000..9dfd17e5 --- /dev/null +++ b/packages/coding-agent/test/k3-context-tier-session.test.ts @@ -0,0 +1,261 @@ +import { existsSync, mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { Agent } from "@dreb/agent-core"; +import { type AssistantMessage, findModel } from "@dreb/ai"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { AgentSession, type AgentSessionEvent } from "../src/core/agent-session.js"; +import { AuthStorage } from "../src/core/auth-storage.js"; +import { + K3_1M_CONTEXT_WINDOW, + K3_256K_CONTEXT_WINDOW, + K3_256K_WIRE_MODEL_ID, + K3_UPGRADE_CUTOFF_TOKENS, +} from "../src/core/k3-context-tier.js"; +import { ModelRegistry } from "../src/core/model-registry.js"; +import { SessionManager } from "../src/core/session-manager.js"; +import { SettingsManager } from "../src/core/settings-manager.js"; +import { createTestResourceLoader } from "./utilities.js"; + +vi.mock("../src/core/compaction/index.js", () => ({ + calculateContextTokens: (usage: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + totalTokens?: number; + }) => usage.totalTokens ?? usage.input + usage.output + usage.cacheRead + usage.cacheWrite, + collectEntriesForBranchSummary: () => ({ entries: [], commonAncestorId: null }), + compact: async () => ({ + summary: "compacted", + firstKeptEntryId: "entry-1", + tokensBefore: 100, + details: {}, + }), + estimateContextTokens: ( + messages: Array<{ + role: string; + usage?: { input: number; output: number; cacheRead: number; cacheWrite: number; totalTokens?: number }; + stopReason?: string; + }>, + ) => { + // Walk backwards to find last non-error, non-aborted assistant with usage + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]; + if (msg.role === "assistant" && msg.stopReason !== "error" && msg.stopReason !== "aborted" && msg.usage) { + const tokens = + msg.usage.totalTokens ?? msg.usage.input + msg.usage.output + msg.usage.cacheRead + msg.usage.cacheWrite; + return { tokens, usageTokens: tokens, trailingTokens: 0, lastUsageIndex: i }; + } + } + return { tokens: 0, usageTokens: 0, trailingTokens: 0, lastUsageIndex: null }; + }, + generateBranchSummary: async () => ({ summary: "", aborted: false, readFiles: [], modifiedFiles: [] }), + prepareCompaction: () => ({ dummy: true }), + shouldCompact: ( + contextTokens: number, + contextWindow: number, + settings: { enabled: boolean; reserveTokens: number }, + ) => settings.enabled && contextTokens > contextWindow - settings.reserveTokens, +})); + +function assistantMessage(overrides: Partial = {}): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text: "done" }], + api: "openai-completions", + provider: "kimi-coding-oauth", + model: "k3", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + ...overrides, + }; +} + +describe("AgentSession K3 auto context tier", () => { + let session: AgentSession; + let tempDir: string; + let settingsManager: SettingsManager; + let events: AgentSessionEvent[]; + + beforeEach(() => { + tempDir = join(tmpdir(), `dreb-k3-tier-test-${Date.now()}`); + mkdirSync(tempDir, { recursive: true }); + events = []; + + const model = findModel("anthropic", "sonnet")!; + const agent = new Agent({ + initialState: { + model, + systemPrompt: "Test", + tools: [], + }, + }); + + const sessionManager = SessionManager.inMemory(); + settingsManager = SettingsManager.create(tempDir, tempDir); + const authStorage = AuthStorage.create(join(tempDir, "auth.json")); + authStorage.setRuntimeApiKey("kimi-coding-oauth", "test-key"); + const modelRegistry = new ModelRegistry(authStorage, tempDir); + + session = new AgentSession({ + agent, + sessionManager, + settingsManager, + cwd: tempDir, + modelRegistry, + resourceLoader: createTestResourceLoader(), + }); + session.subscribe((event) => { + events.push(event); + }); + }); + + afterEach(() => { + session.dispose(); + vi.useRealTimers(); + vi.restoreAllMocks(); + if (tempDir && existsSync(tempDir)) { + rmSync(tempDir, { recursive: true }); + } + }); + + async function selectK3() { + const k3 = findModel("kimi-coding-oauth", "k3")!; + await session.setModel(k3); + } + + async function checkCompaction(message: AssistantMessage) { + await ( + session as unknown as { + _checkCompaction: (msg: AssistantMessage, skipAbortedCheck?: boolean) => Promise; + } + )._checkCompaction(message); + } + + it("starts the user-facing k3 model in the cheaper 256k wire tier", async () => { + await selectK3(); + + expect(session.model?.id).toBe("k3"); + expect(session.model?.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + }); + + it("leaves non-K3 models untouched", async () => { + expect(session.model?.id).toContain("claude"); + expect(session.model?.wireModelId).toBeUndefined(); + }); + + it("upgrades to the 1M tier instead of compacting at the default threshold", async () => { + await selectK3(); + + await checkCompaction( + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 1 } }), + ); + + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + expect(session.model?.wireModelId).toBeUndefined(); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(true); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + }); + + it("stays in the 256k tier below the cutoff", async () => { + await selectK3(); + + await checkCompaction( + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS - 1000 } }), + ); + + expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + expect(session.model?.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(false); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + }); + + it("compacts first when the user lowers the compaction threshold, effectively disabling the upgrade", async () => { + await selectK3(); + // Apply after setModel: setModel persists the default model, which + // rebuilds settings from the on-disk layers. + settingsManager.applyOverrides({ compaction: { reserveTokens: 100000 } }); + + // 200k tokens: below the 256k upgrade cutoff (245760) but past the + // user-lowered compaction threshold (262144 - 100000 = 162144). + await checkCompaction(assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: 200000 } })); + + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(true); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(false); + expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + }); + + it("upgrades instead of compacting when the 256k tier overflows", async () => { + vi.useFakeTimers(); + await selectK3(); + const continueSpy = vi.spyOn(session.agent, "continue").mockResolvedValue(); + + await checkCompaction( + assistantMessage({ + stopReason: "error", + errorMessage: "exceeded model token limit: 262144 (requested: 300000)", + }), + ); + await vi.advanceTimersByTimeAsync(100); + + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + expect(session.model?.wireModelId).toBeUndefined(); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(true); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + expect(continueSpy).toHaveBeenCalledTimes(1); + }); + + it("still upgrades on overflow when compaction is disabled", async () => { + vi.useFakeTimers(); + await selectK3(); + // Apply after setModel: setModel persists the default model, which + // rebuilds settings from the on-disk layers. + settingsManager.applyOverrides({ compaction: { enabled: false } }); + vi.spyOn(session.agent, "continue").mockResolvedValue(); + + await checkCompaction( + assistantMessage({ + stopReason: "error", + errorMessage: "exceeded model token limit: 262144 (requested: 300000)", + }), + ); + await vi.advanceTimersByTimeAsync(100); + + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(true); + }); + + it("does not upgrade on overflow once already in the 1M tier", async () => { + vi.useFakeTimers(); + await selectK3(); + // Pre-seed a large context so setModel derives the 1M tier directly. + session.agent.replaceMessages([ + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 5000 } }), + ]); + await selectK3(); + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + + const upgradeEventsBefore = events.filter((e) => e.type === "context_window_upgrade").length; + await checkCompaction( + assistantMessage({ + stopReason: "error", + errorMessage: "exceeded model token limit: 1048576 (requested: 1100000)", + }), + ); + await vi.advanceTimersByTimeAsync(100); + + // 1M-tier overflow falls back to the normal compaction recovery path. + expect(events.filter((e) => e.type === "context_window_upgrade")).toHaveLength(upgradeEventsBefore); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(true); + }); +}); diff --git a/packages/coding-agent/test/k3-context-tier.test.ts b/packages/coding-agent/test/k3-context-tier.test.ts new file mode 100644 index 00000000..41d04cf5 --- /dev/null +++ b/packages/coding-agent/test/k3-context-tier.test.ts @@ -0,0 +1,106 @@ +import type { Model } from "@dreb/ai"; +import { describe, expect, it } from "vitest"; +import { + deriveK3ContextTierModel, + isK3Model, + isK3256kTier, + K3_1M_CONTEXT_WINDOW, + K3_256K_CONTEXT_WINDOW, + K3_256K_WIRE_MODEL_ID, + K3_UPGRADE_CUTOFF_TOKENS, + shouldUpgradeK3Tier, +} from "../src/core/k3-context-tier.js"; + +function k3Model(overrides: Partial> = {}): Model<"openai-completions"> { + return { + id: "k3", + name: "Kimi K3", + api: "openai-completions", + provider: "kimi-coding-oauth", + baseUrl: "https://api.kimi.com/coding/v1", + reasoning: true, + input: ["text", "image"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: K3_1M_CONTEXT_WINDOW, + maxTokens: 32768, + ...overrides, + }; +} + +describe("k3-context-tier", () => { + it("anchors the upgrade cutoff to the default compaction threshold of the 256k window", () => { + expect(K3_256K_CONTEXT_WINDOW).toBe(262144); + expect(K3_1M_CONTEXT_WINDOW).toBe(1048576); + expect(K3_UPGRADE_CUTOFF_TOKENS).toBe(262144 - 16384); + }); + + it("identifies only the user-facing kimi-coding-oauth k3 model", () => { + expect(isK3Model(k3Model())).toBe(true); + expect(isK3Model(k3Model({ id: "kimi-for-coding" }))).toBe(false); + expect(isK3Model(k3Model({ provider: "kimi-coding" as Model<"openai-completions">["provider"] }))).toBe(false); + expect(isK3Model(undefined)).toBe(false); + }); + + it("passes non-K3 models through unchanged", () => { + const other = k3Model({ id: "kimi-for-coding", contextWindow: 262144 }); + expect(deriveK3ContextTierModel(other, 0)).toBe(other); + expect(deriveK3ContextTierModel(other, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(other); + expect(shouldUpgradeK3Tier(other, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(false); + }); + + it("derives the 256k tier for small contexts", () => { + const derived = deriveK3ContextTierModel(k3Model(), 1000); + expect(derived.id).toBe("k3"); + expect(derived.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + expect(derived.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + expect(isK3256kTier(derived)).toBe(true); + }); + + it("derives the 1M tier past the cutoff", () => { + const derived = deriveK3ContextTierModel(k3Model(), K3_UPGRADE_CUTOFF_TOKENS + 1); + expect(derived.id).toBe("k3"); + expect(derived.wireModelId).toBeUndefined(); + expect(derived.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + expect(isK3256kTier(derived)).toBe(false); + }); + + it("treats the cutoff as inclusive of the cheaper tier", () => { + const atCutoff = deriveK3ContextTierModel(k3Model(), K3_UPGRADE_CUTOFF_TOKENS); + expect(isK3256kTier(atCutoff)).toBe(true); + expect(shouldUpgradeK3Tier(atCutoff, K3_UPGRADE_CUTOFF_TOKENS)).toBe(false); + }); + + it("upgrades a derived 256k-tier model to the 1M tier", () => { + const tier256k = deriveK3ContextTierModel(k3Model(), 0); + expect(shouldUpgradeK3Tier(tier256k, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(true); + + const upgraded = deriveK3ContextTierModel(tier256k, K3_UPGRADE_CUTOFF_TOKENS + 1); + expect(upgraded.wireModelId).toBeUndefined(); + expect(upgraded.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + }); + + it("does not upgrade the 1M tier or other models", () => { + const tier1m = deriveK3ContextTierModel(k3Model(), K3_UPGRADE_CUTOFF_TOKENS + 1); + expect(shouldUpgradeK3Tier(tier1m, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(false); + expect(shouldUpgradeK3Tier(undefined, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(false); + }); + + it("derives the 256k tier again when context drops below the cutoff (fresh session / post-compaction)", () => { + const tier1m = deriveK3ContextTierModel(k3Model(), K3_UPGRADE_CUTOFF_TOKENS + 1); + const redegraded = deriveK3ContextTierModel(tier1m, 5000); + expect(redegraded.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + expect(redegraded.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + }); + + it("is idempotent for already-derived models", () => { + const once256k = deriveK3ContextTierModel(k3Model(), 0); + expect(deriveK3ContextTierModel(once256k, 0)).toBe(once256k); + + const once1m = deriveK3ContextTierModel(k3Model(), K3_UPGRADE_CUTOFF_TOKENS + 1); + expect(deriveK3ContextTierModel(once1m, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(once1m); + }); + + it("passes undefined through", () => { + expect(deriveK3ContextTierModel(undefined, 0)).toBeUndefined(); + }); +}); From 0fe577efc8fbf222cab9e6814a2abd3f5ccb3ba4 Mon Sep 17 00:00:00 2001 From: Acters Date: Wed, 29 Jul 2026 20:05:05 -0700 Subject: [PATCH 2/2] Kimi OAuth: address k3 context tier review findings - Threshold-path upgrade now fires even when compaction is disabled: the upgrade is model-capability management, not context reduction - Resume re-derives the tier from restored session context instead of always starting in the 256k tier; compacted sessions return to the cheaper 256k wire tier (auto and manual compaction paths) - models.json contextWindow overrides disable automatic tiering instead of being silently replaced - A user-lowered compaction threshold preempts the upgrade even when a single response jumps straight past the cutoff - Extract duplicated last-assistant-message removal into a helper and drop a redundant upgrade guard - Tests: SDK bootstrap/resume, post-compaction tier return, override passthrough, jump precedence, pre-prompt path, non-K3 disabled-overflow regression, interactive status line, exact event payload assertions --- packages/coding-agent/docs/providers.md | 2 +- .../coding-agent/src/core/agent-session.ts | 53 +++++--- .../coding-agent/src/core/k3-context-tier.ts | 11 +- packages/coding-agent/src/core/sdk.ts | 10 +- .../test/interactive-mode-status.test.ts | 21 ++++ .../test/k3-context-tier-sdk.test.ts | 101 +++++++++++++++ .../test/k3-context-tier-session.test.ts | 118 +++++++++++++++++- .../coding-agent/test/k3-context-tier.test.ts | 15 +++ 8 files changed, 307 insertions(+), 24 deletions(-) create mode 100644 packages/coding-agent/test/k3-context-tier-sdk.test.ts diff --git a/packages/coding-agent/docs/providers.md b/packages/coding-agent/docs/providers.md index 0ab91fc2..c877804c 100644 --- a/packages/coding-agent/docs/providers.md +++ b/packages/coding-agent/docs/providers.md @@ -52,7 +52,7 @@ Use `/logout` to clear credentials. Tokens are stored in `~/.dreb/agent/auth.jso - `/login` uses the Kimi Code OAuth subscription endpoint at `https://api.kimi.com/coding/v1`. - `KIMI_API_KEY` uses Kimi For Coding's Anthropic-compatible API at `https://api.kimi.com/coding`. - Built-in OAuth models are `kimi-for-coding` (default, 262k context), `k3` (1M context), and `kimi-for-coding-highspeed` (262k context). Model availability is plan-dependent: the `k3` 1M-context ID and the `kimi-for-coding-highspeed` ID are only exposed when the subscription includes them, and `kimi-for-coding-highspeed` runs at roughly 6× speed for 3× quota usage. On login/refresh, dreb asks the Kimi API which models the subscription is entitled to and updates context, reasoning, image, tool-use, protocol, and thinking-effort metadata. Compatible newly discovered IDs are templated conservatively; if discovery fails, the static list remains available. -- The OAuth `k3` model is context-tiered: the Kimi endpoint serves it as `k3-256k` (256k context, cheaper) and `k3` (1M context, stated to consume 2× the quota), and upgrading does not invalidate the prompt cache. The cheaper variant is exclusive to the Kimi for Coding OAuth endpoint — the pay-per-token Moonshot AI Platform does not expose it. dreb starts every session on the `k3-256k` wire ID and automatically upgrades to `k3` once the session context passes the 256k cutoff (256k minus the default compaction reserve) instead of compacting. A user-lowered compaction threshold compacts before the cutoff and thus effectively disables the upgrade. +- The OAuth `k3` model is context-tiered: the Kimi endpoint serves it as `k3-256k` (256k context, cheaper) and `k3` (1M context, stated to consume 2× the quota), and upgrading does not invalidate the prompt cache. The cheaper variant is exclusive to the Kimi for Coding OAuth endpoint — the pay-per-token Moonshot AI Platform does not expose it. dreb starts every session on the `k3-256k` wire ID and automatically upgrades to `k3` once the session context passes the 256k cutoff (256k minus the default compaction reserve) instead of compacting. A user-lowered compaction threshold compacts before the cutoff and thus effectively disables the upgrade. Customizing `k3`'s context window via `models.json` disables automatic tiering entirely, and compacted or fresh sessions return to the cheaper tier. - OAuth requests use the current Kimi Code device identity contract and share its stable `~/.kimi-code/device_id`. Extra low-precedence headers can be supplied with newline-separated `KIMI_CODE_CUSTOM_HEADERS` values. - The OAuth coding endpoint accepts OpenAI-style multimodal content arrays with base64 `image_url` data URLs. dreb keeps a conservative 32k output-token cap because the managed model catalog does not currently advertise a per-model output limit. - Moonshot Open Platform uses a different base URL (`https://api.moonshot.ai/v1`); don't assume both routes expose identical behavior. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index e7792a7b..bc868945 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2367,6 +2367,11 @@ export class AgentSession { const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.replaceMessages(sessionContext.messages); + // Re-derive the K3 context tier: the compacted context is small again, + // so the session returns to the cheaper 256k wire tier. + if (this.model) { + this.agent.setModel(this._applyContextTier(this.model)); + } // Get the saved compaction entry for the extension event const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as @@ -2453,10 +2458,7 @@ export class AgentSession { if (this._tryUpgradeK3ContextTier()) { // Remove the error message from agent state (it IS saved to session // for history, but we don't want it in context for the retry) - const messages = this.agent.state.messages; - if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { - this.agent.replaceMessages(messages.slice(0, -1)); - } + this._removeLastAssistantMessage(); setTimeout(() => { this.agent.continue().catch((err) => { this.warnInSession( @@ -2484,17 +2486,12 @@ export class AgentSession { this._overflowRecoveryAttempted = true; // Remove the error message from agent state (it IS saved to session for history, // but we don't want it in context for the retry) - const messages = this.agent.state.messages; - if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { - this.agent.replaceMessages(messages.slice(0, -1)); - } + this._removeLastAssistantMessage(); await this._runAutoCompaction("overflow", true); return; } // Case 2: Threshold - context is getting large - if (!settings.enabled) return; - // For error messages (no usage data), estimate from last successful response. // This ensures sessions that hit persistent API errors (e.g. 529) can still compact. let contextTokens: number; @@ -2519,13 +2516,26 @@ export class AgentSession { } // K3 auto context tier: reaching the 256k cutoff upgrades to the 1M tier - // instead of compacting. The cutoff is fixed at the default compaction - // threshold of the 256k window, so a user-lowered compaction threshold - // fires first and effectively disables the upgrade. - if (shouldUpgradeK3Tier(this.model, contextTokens) && this._tryUpgradeK3ContextTier()) { - contextWindow = this.model?.contextWindow ?? contextWindow; + // instead of compacting. This is model-capability management, not context + // reduction, so it applies even when compaction is disabled. The cutoff is + // fixed at the default compaction threshold of the 256k window; a + // user-lowered compaction threshold takes precedence and effectively + // disables the upgrade. + if (shouldUpgradeK3Tier(this.model, contextTokens)) { + // A user-lowered compaction threshold takes precedence over the + // upgrade: if the user's compact point for the 256k window sits below + // the default cutoff and is already exceeded, compact instead. + const userCompactPoint = K3_256K_CONTEXT_WINDOW - settings.reserveTokens; + const userThresholdPreempts = + settings.enabled && userCompactPoint < K3_UPGRADE_CUTOFF_TOKENS && contextTokens > userCompactPoint; + if (!userThresholdPreempts) { + this._tryUpgradeK3ContextTier(); + contextWindow = this.model?.contextWindow ?? contextWindow; + } } + if (!settings.enabled) return; + if (shouldCompact(contextTokens, contextWindow, settings)) { await this._runAutoCompaction("threshold", false); } @@ -2541,6 +2551,14 @@ export class AgentSession { return deriveK3ContextTierModel(model, estimateContextTokens(this.agent.state.messages).tokens); } + /** Remove the last message from agent state when it is an assistant message. */ + private _removeLastAssistantMessage(): void { + const messages = this.agent.state.messages; + if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { + this.agent.replaceMessages(messages.slice(0, -1)); + } + } + /** * Upgrade the K3 auto context tier from 256k to 1M. The Kimi backend * upgrades the prompt cache seamlessly, so no context is lost or @@ -2648,6 +2666,11 @@ export class AgentSession { const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.replaceMessages(sessionContext.messages); + // Re-derive the K3 context tier: the compacted context is small again, + // so the session returns to the cheaper 256k wire tier. + if (this.model) { + this.agent.setModel(this._applyContextTier(this.model)); + } // Get the saved compaction entry for the extension event const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as diff --git a/packages/coding-agent/src/core/k3-context-tier.ts b/packages/coding-agent/src/core/k3-context-tier.ts index 5ebc5f38..e575293d 100644 --- a/packages/coding-agent/src/core/k3-context-tier.ts +++ b/packages/coding-agent/src/core/k3-context-tier.ts @@ -67,7 +67,9 @@ export function isK3256kTier(model: Model | null | undefined): boolean { * on the wire (the Kimi backend upgrades the cache seamlessly). * * Returns the input unchanged for non-K3 models and is idempotent for - * already-derived models. + * already-derived models. A K3 model whose context window was customized + * (e.g. a models.json override) is also returned unchanged — automatic + * tiering never silently replaces user-configured limits. */ export function deriveK3ContextTierModel(model: Model, contextTokens: number): Model; export function deriveK3ContextTierModel( @@ -79,12 +81,15 @@ export function deriveK3ContextTierModel( contextTokens: number, ): Model | undefined { if (!model || !isK3Model(model)) return model; + const isStock = model.contextWindow === K3_1M_CONTEXT_WINDOW && model.wireModelId === undefined; + const isDerived256k = model.contextWindow === K3_256K_CONTEXT_WINDOW && model.wireModelId === K3_256K_WIRE_MODEL_ID; + if (!isStock && !isDerived256k) return model; if (contextTokens > K3_UPGRADE_CUTOFF_TOKENS) { - if (model.wireModelId === undefined && model.contextWindow === K3_1M_CONTEXT_WINDOW) return model; + if (isStock) return model; const { wireModelId: _droppedWireModelId, ...rest } = model; return { ...rest, contextWindow: K3_1M_CONTEXT_WINDOW }; } - if (model.wireModelId === K3_256K_WIRE_MODEL_ID && model.contextWindow === K3_256K_CONTEXT_WINDOW) return model; + if (isDerived256k) return model; return { ...model, contextWindow: K3_256K_CONTEXT_WINDOW, wireModelId: K3_256K_WIRE_MODEL_ID }; } diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index b8da293c..fa7eea2c 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -4,6 +4,7 @@ import type { Message, Model } from "@dreb/ai"; import { getAgentDir, getDocsPath } from "../config.js"; import { AgentSession } from "./agent-session.js"; import { AuthStorage } from "./auth-storage.js"; +import { estimateContextTokens } from "./compaction/index.js"; import { DEFAULT_THINKING_LEVEL } from "./defaults.js"; import type { ExtensionRunner, LoadExtensionsResult, ToolDefinition } from "./extensions/index.js"; import { deriveK3ContextTierModel } from "./k3-context-tier.js"; @@ -328,9 +329,12 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} agent = new Agent({ initialState: { systemPrompt: "", - // New sessions start with empty context, so K3 starts in the - // cheaper 256k wire tier (upgrades automatically past the cutoff). - model: deriveK3ContextTierModel(model, 0), + // K3 auto context tier: new sessions start on the cheaper 256k wire + // tier; a resumed session derives the tier from its restored context. + model: deriveK3ContextTierModel( + model, + hasExistingSession ? estimateContextTokens(existingSession.messages).tokens : 0, + ), thinkingLevel, tools: [], }, diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index d5522024..c5480f0e 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -272,6 +272,27 @@ describe("InteractiveMode working indicator", () => { expect(inlineStatuses.at(-1)).toBeNull(); }); + test("context_window_upgrade shows a status line and invalidates the footer", async () => { + const { fakeThis } = createWorkingFakeThis(); + fakeThis.chatContainer = new Container(); + fakeThis.lastStatusSpacer = undefined; + fakeThis.lastStatusText = undefined; + fakeThis.showStatus = (...args: unknown[]) => + (InteractiveMode as any).prototype.showStatus.call(fakeThis, ...args); + + await dispatchEvent(fakeThis, { + type: "context_window_upgrade", + provider: "kimi-coding-oauth", + modelId: "k3", + fromContextWindow: 262144, + toContextWindow: 1048576, + }); + + const rendered = fakeThis.chatContainer.children.flatMap((child: any) => child.render(120)).join("\n"); + expect(rendered).toContain("k3 context window upgraded: 256k → 1M (cache preserved)"); + expect(fakeThis.footer.invalidate).toHaveBeenCalled(); + }); + test("non-agent inline loaders can update and clear branch summary, dream, and compaction statuses", () => { const { fakeThis, inlineStatuses } = createWorkingFakeThis(); diff --git a/packages/coding-agent/test/k3-context-tier-sdk.test.ts b/packages/coding-agent/test/k3-context-tier-sdk.test.ts new file mode 100644 index 00000000..58dc0298 --- /dev/null +++ b/packages/coding-agent/test/k3-context-tier-sdk.test.ts @@ -0,0 +1,101 @@ +import { existsSync, mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { type AssistantMessage, findModel } from "@dreb/ai"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.js"; +import { + K3_1M_CONTEXT_WINDOW, + K3_256K_CONTEXT_WINDOW, + K3_256K_WIRE_MODEL_ID, + K3_UPGRADE_CUTOFF_TOKENS, +} from "../src/core/k3-context-tier.js"; +import { DefaultResourceLoader } from "../src/core/resource-loader.js"; +import { createAgentSession } from "../src/core/sdk.js"; +import { SessionManager } from "../src/core/session-manager.js"; +import { SettingsManager } from "../src/core/settings-manager.js"; + +function bigAssistantMessage(): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text: "..." }], + api: "openai-completions", + provider: "kimi-coding-oauth", + model: "k3", + usage: { + input: 240000, + output: 7000, + cacheRead: 0, + cacheWrite: 0, + totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 5000, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +describe("createAgentSession K3 auto context tier", () => { + let tempDir: string; + let agentDir: string; + + beforeEach(() => { + tempDir = join(tmpdir(), `dreb-k3-sdk-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); + agentDir = join(tempDir, "agent"); + mkdirSync(agentDir, { recursive: true }); + }); + + afterEach(() => { + if (tempDir && existsSync(tempDir)) { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + + async function createSession(sessionManager: SessionManager) { + const settingsManager = SettingsManager.create(tempDir, agentDir); + const authStorage = AuthStorage.create(join(agentDir, "auth.json")); + authStorage.setRuntimeApiKey("kimi-coding-oauth", "test-key"); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + model: findModel("kimi-coding-oauth", "k3")!, + settingsManager, + sessionManager, + authStorage, + resourceLoader, + }); + return session; + } + + it("starts a fresh k3 session on the cheaper 256k wire tier", async () => { + const session = await createSession(SessionManager.inMemory()); + + expect(session.model?.id).toBe("k3"); + expect(session.model?.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + + session.dispose(); + }); + + it("resumes straight into the 1M tier when restored context already exceeds the cutoff", async () => { + const sessionManager = SessionManager.create(tempDir); + sessionManager.appendModelChange("kimi-coding-oauth", "k3"); + sessionManager.appendMessage(bigAssistantMessage()); + + const session = await createSession(sessionManager); + + expect(session.model?.id).toBe("k3"); + expect(session.model?.wireModelId).toBeUndefined(); + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + + session.dispose(); + }); +}); diff --git a/packages/coding-agent/test/k3-context-tier-session.test.ts b/packages/coding-agent/test/k3-context-tier-session.test.ts index 9dfd17e5..e457afe4 100644 --- a/packages/coding-agent/test/k3-context-tier-session.test.ts +++ b/packages/coding-agent/test/k3-context-tier-session.test.ts @@ -163,7 +163,30 @@ describe("AgentSession K3 auto context tier", () => { expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); expect(session.model?.wireModelId).toBeUndefined(); - expect(events.some((e) => e.type === "context_window_upgrade")).toBe(true); + const upgradeEvents = events.filter((e) => e.type === "context_window_upgrade"); + expect(upgradeEvents).toHaveLength(1); + expect(upgradeEvents[0]).toMatchObject({ + provider: "kimi-coding-oauth", + modelId: "k3", + fromContextWindow: K3_256K_CONTEXT_WINDOW, + toContextWindow: K3_1M_CONTEXT_WINDOW, + }); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + }); + + it("upgrades on the threshold path even when compaction is disabled", async () => { + await selectK3(); + // Apply after setModel: setModel persists the default model, which + // rebuilds settings from the on-disk layers. + settingsManager.applyOverrides({ compaction: { enabled: false } }); + + await checkCompaction( + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 1 } }), + ); + + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + expect(session.model?.wireModelId).toBeUndefined(); + expect(events.filter((e) => e.type === "context_window_upgrade")).toHaveLength(1); expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); }); @@ -195,6 +218,39 @@ describe("AgentSession K3 auto context tier", () => { expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); }); + it("honors the lowered threshold even when one response jumps straight past the cutoff", async () => { + await selectK3(); + settingsManager.applyOverrides({ compaction: { reserveTokens: 100000 } }); + + // 250k tokens: past the 256k upgrade cutoff AND past the user-lowered + // threshold (162144). The user's threshold preempts the upgrade. + await checkCompaction(assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: 250000 } })); + + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(true); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(false); + expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + }); + + it("upgrades via the pre-prompt check for an aborted response (skipAbortedCheck=false)", async () => { + await selectK3(); + + await ( + session as unknown as { + _checkCompaction: (msg: AssistantMessage, skipAbortedCheck?: boolean) => Promise; + } + )._checkCompaction( + assistantMessage({ + stopReason: "aborted", + usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 1 }, + }), + false, + ); + + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + expect(events.filter((e) => e.type === "context_window_upgrade")).toHaveLength(1); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + }); + it("upgrades instead of compacting when the 256k tier overflows", async () => { vi.useFakeTimers(); await selectK3(); @@ -221,7 +277,7 @@ describe("AgentSession K3 auto context tier", () => { // Apply after setModel: setModel persists the default model, which // rebuilds settings from the on-disk layers. settingsManager.applyOverrides({ compaction: { enabled: false } }); - vi.spyOn(session.agent, "continue").mockResolvedValue(); + const continueSpy = vi.spyOn(session.agent, "continue").mockResolvedValue(); await checkCompaction( assistantMessage({ @@ -233,6 +289,64 @@ describe("AgentSession K3 auto context tier", () => { expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); expect(events.some((e) => e.type === "context_window_upgrade")).toBe(true); + expect(continueSpy).toHaveBeenCalledTimes(1); + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + }); + + it("returns to the cheaper 256k tier after auto-compaction", async () => { + await selectK3(); + // Upgrade to the 1M tier first. + await checkCompaction( + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 1 } }), + ); + expect(session.model?.contextWindow).toBe(K3_1M_CONTEXT_WINDOW); + + // Context keeps growing past the 1M compaction threshold (1M - reserve). + await checkCompaction( + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_1M_CONTEXT_WINDOW - 1000 } }), + ); + + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(true); + expect(session.model?.contextWindow).toBe(K3_256K_CONTEXT_WINDOW); + expect(session.model?.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + }); + + it("leaves a k3 model with a user-overridden context window untouched", async () => { + const k3 = { ...findModel("kimi-coding-oauth", "k3")!, contextWindow: 131072 }; + await session.setModel(k3); + + expect(session.model?.contextWindow).toBe(131072); + expect(session.model?.wireModelId).toBeUndefined(); + + await checkCompaction( + assistantMessage({ usage: { ...assistantMessage().usage, totalTokens: K3_UPGRADE_CUTOFF_TOKENS + 1 } }), + ); + expect(session.model?.contextWindow).toBe(131072); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(false); + }); + + it("does nothing for a non-K3 overflow when compaction is disabled (gate-move regression)", async () => { + vi.useFakeTimers(); + // Default session model is anthropic; disable compaction after any model persistence. + settingsManager.applyOverrides({ compaction: { enabled: false } }); + const continueSpy = vi.spyOn(session.agent, "continue").mockResolvedValue(); + const modelBefore = session.model; + + await checkCompaction( + assistantMessage({ + provider: "anthropic", + model: modelBefore!.id, + api: "anthropic-messages", + stopReason: "error", + errorMessage: "prompt is too long: 250000 tokens > 200000 maximum", + }), + ); + await vi.advanceTimersByTimeAsync(100); + + expect(events.some((e) => e.type === "auto_compaction_start")).toBe(false); + expect(events.some((e) => e.type === "context_window_upgrade")).toBe(false); + expect(continueSpy).not.toHaveBeenCalled(); + expect(session.model).toBe(modelBefore); }); it("does not upgrade on overflow once already in the 1M tier", async () => { diff --git a/packages/coding-agent/test/k3-context-tier.test.ts b/packages/coding-agent/test/k3-context-tier.test.ts index 41d04cf5..a58e1d9a 100644 --- a/packages/coding-agent/test/k3-context-tier.test.ts +++ b/packages/coding-agent/test/k3-context-tier.test.ts @@ -103,4 +103,19 @@ describe("k3-context-tier", () => { it("passes undefined through", () => { expect(deriveK3ContextTierModel(undefined, 0)).toBeUndefined(); }); + + it("never replaces a user-customized context window (models.json override)", () => { + const overridden = k3Model({ contextWindow: 131072 }); + expect(deriveK3ContextTierModel(overridden, 0)).toBe(overridden); + expect(deriveK3ContextTierModel(overridden, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(overridden); + expect(shouldUpgradeK3Tier(overridden, K3_UPGRADE_CUTOFF_TOKENS + 1)).toBe(false); + }); + + it("resumes tiering after a derived model returns to the stock 1M shape", () => { + const tier256k = deriveK3ContextTierModel(k3Model(), 0); + const backTo1m = deriveK3ContextTierModel(tier256k, K3_UPGRADE_CUTOFF_TOKENS + 1); + // The 1M tier is the stock shape again, so a fresh derivation keeps working. + const redegraded = deriveK3ContextTierModel(backTo1m, 100); + expect(redegraded.wireModelId).toBe(K3_256K_WIRE_MODEL_ID); + }); });