From 3b6e746396733356213992a9a3fecd794edef087 Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Thu, 16 Jul 2026 17:05:36 +0800 Subject: [PATCH 1/3] fix(agent-core-v2): keep context size readings on the measured path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The step fold creates the assistant message in the context before the exchange finishes (a skeleton at step.begin, filled by content.part folds during streaming), and the input array llmRequester passes to contextSize.measured() is that same live array — it already includes the output. Taking input.length + output.length therefore counted the folded output twice, storing a measured prefix length one past the live context. The inflated length permanently failed get()'s measured fast path, so reads silently fell back to per-message estimates (e.g. ~50 tokens shown for a ~29k-token "hi"). Take the live context length as the measured prefix instead (input and context are identical under the identity guard), and clamp the measured prefix to the context length in get() so a bad record can never knock reads off the measured path again. Add contextSize tests driving real turns that assert the wire model, get(), and rpc getContext against the exchange totals. --- .../agent/contextSize/contextSizeService.ts | 22 +++- .../agent/contextSize/contextSize.test.ts | 110 ++++++++++++++++++ 2 files changed, 127 insertions(+), 5 deletions(-) create mode 100644 packages/agent-core-v2/test/agent/contextSize/contextSize.test.ts diff --git a/packages/agent-core-v2/src/agent/contextSize/contextSizeService.ts b/packages/agent-core-v2/src/agent/contextSize/contextSizeService.ts index 67082d30fc..ef1b0c55c1 100644 --- a/packages/agent-core-v2/src/agent/contextSize/contextSizeService.ts +++ b/packages/agent-core-v2/src/agent/contextSize/contextSizeService.ts @@ -46,12 +46,17 @@ export class AgentContextSizeService extends Disposable implements IAgentContext get(start?: number, end?: number): ContextSize { const context = this.context.get(); const model = this.wire.getModel(ContextSizeModel); + // Defensive clamp: the measured prefix can never be longer than the live + // context. An op written against a mutated message array once inflated + // `model.length` past `context.length`, silently knocking every read off + // the measured path onto the per-message estimate branch. + const measuredLength = Math.min(model.length, context.length); const from = normalizeSliceIndex(start ?? 0, context.length); const to = normalizeSliceIndex(end ?? context.length, context.length); - const measuredEnd = Math.min(to, model.length); - const estimatedStart = Math.max(from, model.length); + const measuredEnd = Math.min(to, measuredLength); + const estimatedStart = Math.max(from, measuredLength); const measured = - from === 0 && measuredEnd === model.length + from === 0 && measuredEnd === measuredLength ? model.tokens : estimateTokensForMessages(context.slice(from, measuredEnd)); const estimated = estimateTokensForMessages(context.slice(estimatedStart, to)); @@ -59,8 +64,15 @@ export class AgentContextSizeService extends Disposable implements IAgentContext } measured(input: readonly Message[], output: readonly Message[], usage: TokenUsage): void { - if (!matchesContext(input, this.context.get())) return; - const length = input.length + output.length; + const context = this.context.get(); + if (!matchesContext(input, context)) return; + // The fold of the step's loop events creates the assistant message in the + // context BEFORE the exchange finishes (a skeleton at `step.begin`, filled + // by `content.part` folds during streaming), and `input` is that same live + // array — so it already includes `output` here. The measured prefix is the + // whole current context; `input.length + output.length` would count the + // folded output twice. + const length = context.length; const tokens = tokenUsageTotal(usage); this.wire.dispatch(contextSizeMeasured({ length, tokens })); this.emitIfChanged(); diff --git a/packages/agent-core-v2/test/agent/contextSize/contextSize.test.ts b/packages/agent-core-v2/test/agent/contextSize/contextSize.test.ts new file mode 100644 index 0000000000..a6b7e47663 --- /dev/null +++ b/packages/agent-core-v2/test/agent/contextSize/contextSize.test.ts @@ -0,0 +1,110 @@ +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { IAgentContextMemoryService, IAgentProfileService } from '#/index'; +import { IAgentContextSizeService } from '#/agent/contextSize/contextSize'; +import { ContextSizeModel, contextSizeMeasured } from '#/agent/contextSize/contextSizeOps'; +import type { TokenUsage } from '#/app/llmProtocol/usage'; +import { IAgentUsageService } from '#/agent/usage/usage'; +import { IWireService } from '#/wire/wire'; + +import { createTestAgent, type TestAgentContext } from '../../harness'; + +function totalOf(usage: TokenUsage | undefined): number { + if (usage === undefined) return 0; + return usage.inputOther + usage.output + usage.inputCacheRead + usage.inputCacheCreation; +} + +describe('Agent context size', () => { + let ctx: TestAgentContext; + let context: IAgentContextMemoryService; + let contextSize: IAgentContextSizeService; + let profile: IAgentProfileService; + let usage: IAgentUsageService; + let wire: IWireService; + + beforeEach(() => { + ctx = createTestAgent(); + context = ctx.get(IAgentContextMemoryService); + contextSize = ctx.get(IAgentContextSizeService); + profile = ctx.get(IAgentProfileService); + usage = ctx.get(IAgentUsageService); + wire = ctx.get(IWireService); + }); + + afterEach(async () => { + try { + await ctx.expectResumeMatches(); + } finally { + await ctx.dispose(); + } + }); + + it('adopts the exchange totals as the measured context size after a turn', async () => { + profile.update({ activeToolNames: [] }); + + ctx.mockNextResponse({ type: 'text', text: 'Hi there!' }); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'hi' }] }); + await ctx.untilTurnEnd(); + + const exchangeTotal = totalOf(usage.status().total); + expect(exchangeTotal).toBeGreaterThan(0); + expect(context.get()).toHaveLength(2); + + // The assistant message is folded into the context before the exchange + // finishes, so the measured prefix must match the live history — an + // inflated length silently knocks `get()` off the measured path onto the + // per-message estimate branch (found as `tokenCount` reading ~50 while + // the provider reported ~29k for a system-prompt-heavy "hi"). + expect(wire.getModel(ContextSizeModel)).toEqual({ + length: context.get().length, + tokens: exchangeTotal, + }); + + const size = contextSize.get(); + expect(size.measured).toBe(exchangeTotal); + expect(size.estimated).toBe(0); + expect(size.size).toBe(exchangeTotal); + expect((await ctx.rpc.getContext({})).tokenCount).toBe(exchangeTotal); + }); + + it('repoints the measured size at the last exchange across turns', async () => { + profile.update({ activeToolNames: [] }); + + ctx.mockNextResponse({ type: 'text', text: 'first' }); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'hi' }] }); + await ctx.untilTurnEnd(); + + ctx.mockNextResponse({ type: 'text', text: 'second reply, a longer one' }); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'again' }] }); + await ctx.untilTurnEnd(); + + const lastExchangeTotal = totalOf(usage.status().currentTurn); + expect(lastExchangeTotal).toBeGreaterThan(0); + expect(context.get()).toHaveLength(4); + + expect(wire.getModel(ContextSizeModel)).toEqual({ + length: context.get().length, + tokens: lastExchangeTotal, + }); + expect(contextSize.get().measured).toBe(lastExchangeTotal); + expect((await ctx.rpc.getContext({})).tokenCount).toBe(lastExchangeTotal); + }); + + it('estimates the not-yet-measured tail instead of dropping it', () => { + ctx.appendUserMessage([{ type: 'text', text: 'hello world, not measured yet' }]); + + const size = contextSize.get(); + expect(size.measured).toBe(0); + expect(size.estimated).toBeGreaterThan(0); + expect(size.size).toBe(size.estimated); + }); + + it('tolerates a stored measured prefix longer than the live context', () => { + ctx.appendUserMessage([{ type: 'text', text: 'only one message' }]); + + // A corrupt/overshooting record must not push reads onto the estimate + // branch; the measured total is clamped to the live context instead. + wire.dispatch(contextSizeMeasured({ length: 5, tokens: 1234 })); + expect(contextSize.get().measured).toBe(1234); + }); +}); From 351ea1ee90a92ff174671375a15dc6292b1e772e Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Thu, 16 Jul 2026 17:05:50 +0800 Subject: [PATCH 2/3] feat(klient): add context-usage example tracing a fresh session Polls agent.getContext()/agent.getUsage() and streams agent events for a new session after one "hi" against a real server, printing a timeline of when the context/token readings move, plus a final consistency check comparing the measured tokenCount to cumulative usage. Model seeding is optional via KIMI_EXAMPLE_* env; the server token resolves from /server.token like the v2 e2e helpers. --- packages/klient/examples/context-usage.ts | 213 ++++++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 packages/klient/examples/context-usage.ts diff --git a/packages/klient/examples/context-usage.ts b/packages/klient/examples/context-usage.ts new file mode 100644 index 0000000000..460c9e7b4f --- /dev/null +++ b/packages/klient/examples/context-usage.ts @@ -0,0 +1,213 @@ +/** + * Trace how the reported "Context size" evolves on a brand-new session after + * a single "hi" prompt, against a real running `kap-server` over the HTTP + * transport. + * + * What gets sampled, all through the klient facade: + * - `agent.getContext()` → `{ history, tokenCount }` — `tokenCount` is the + * last MEASURED exchange total (`contextSize.get().measured` server-side); + * it is 0 until the first LLM response lands and stays flat between turns. + * - `agent.getUsage()` → accumulated token usage (`byModel` / `currentTurn` + * / `total`), recorded per request. + * - `agent.status.updated` events — the live `contextTokens` / `usage` + * slices that feed the TUI footer. + * + * A 250 ms poll diffs (history length, tokenCount, usage.total) and prints a + * line only when something changed, so the output is a timeline of exactly + * when the Context size reading moves — and when it does NOT. + * + * Run against a local dev server (auth token resolves from + * `/server.token`, same as the v2 e2e helpers): + * pnpm dev:kap-server + * pnpm -C packages/klient exec tsx examples/context-usage.ts + * + * Env: + * KIMI_SERVER_URL — default `http://127.0.0.1:58627` + * KIMI_SERVER_TOKEN — default: contents of `/server.token` + * KIMI_EXAMPLE_MODEL — seed a throwaway model with this gateway model id… + * KIMI_EXAMPLE_API_KEY — …and this API key (otherwise the server's default model is used) + * KIMI_EXAMPLE_BASE_URL — optional gateway base URL for the seeded model + * KIMI_EXAMPLE_PROTOCOL — optional wire protocol for the seeded model (default `openai`) + */ +import { readFile } from 'node:fs/promises'; +import { homedir } from 'node:os'; +import { join } from 'node:path'; + +import { createKlient } from '@moonshot-ai/klient/http'; + +const BASE = process.env['KIMI_SERVER_URL'] ?? 'http://127.0.0.1:58627'; +const SEEDED_MODEL_ID = 'klient-example-model'; + +interface TokenUsage { + inputOther: number; + output: number; + inputCacheRead: number; + inputCacheCreation: number; +} + +function usageTotal(usage: TokenUsage | undefined): number | undefined { + if (usage === undefined) return undefined; + return usage.inputOther + usage.output + usage.inputCacheRead + usage.inputCacheCreation; +} + +async function resolveToken(): Promise { + const explicit = process.env['KIMI_SERVER_TOKEN']; + if (explicit !== undefined && explicit.length > 0) return explicit; + const home = process.env['KIMI_CODE_HOME'] ?? join(homedir(), '.kimi-code'); + try { + const token = (await readFile(join(home, 'server.token'), 'utf8')).trim(); + return token.length > 0 ? token : undefined; + } catch { + return undefined; + } +} + +const tick = (ms: number): Promise => + new Promise((resolve) => { + setTimeout(resolve, ms); + }); + +async function main(): Promise { + const token = await resolveToken(); + const klient = createKlient({ url: BASE, token }); + + const session = await klient.global.sessions.create({ workDir: process.cwd() }); + console.log('[session] created ->', session.id); + const agent = klient.session(session.id).agent('main'); + + // Optional throwaway model seeding; otherwise the server's default model. + const seedModel = process.env['KIMI_EXAMPLE_MODEL']; + const seedKey = process.env['KIMI_EXAMPLE_API_KEY']; + if (seedModel !== undefined && seedKey !== undefined) { + await klient.global.models.set({ + id: SEEDED_MODEL_ID, + config: { + model: seedModel, + apiKey: seedKey, + baseUrl: process.env['KIMI_EXAMPLE_BASE_URL'], + protocol: (process.env['KIMI_EXAMPLE_PROTOCOL'] ?? 'openai') as 'openai', + maxContextSize: 262_144, + }, + }); + await agent.setModel(SEEDED_MODEL_ID); + } + console.log('[model] bound ->', await agent.getModel()); + + const startedAt = Date.now(); + const elapsed = (): string => `+${String(Date.now() - startedAt).padStart(6)}ms`; + + // Live status slices (what the TUI footer consumes), as they arrive. + agent.events.on('agent.status.updated', (event) => { + const slice: Record = {}; + if ('contextTokens' in event) slice['contextTokens'] = event['contextTokens']; + if ('maxContextTokens' in event) slice['maxContextTokens'] = event['maxContextTokens']; + if ('contextUsage' in event) slice['contextUsage'] = event['contextUsage']; + if ('phase' in event) slice['phase'] = event['phase']; + const usage = event['usage'] as { total?: TokenUsage } | undefined; + if (usage !== undefined) slice['usage.total'] = usageTotal(usage.total); + console.log(`[event] ${elapsed()} agent.status.updated ->`, JSON.stringify(slice)); + }); + agent.events.on('turn.started', (event) => { + console.log(`[event] ${elapsed()} turn.started -> turnId=${String(event.turnId)}`); + }); + agent.events.on('turn.ended', (event) => { + console.log(`[event] ${elapsed()} turn.ended -> reason=${event.reason}`); + }); + agent.events.on('error', (event) => { + console.log(`[event] ${elapsed()} error ->`, JSON.stringify(event)); + }); + agent.events.onError((error) => { + console.log(`[event-err] ${elapsed()} ${error.message.split('\n')[0] ?? error.message}`); + }); + + const completed = new Promise<'completed' | 'failed' | 'timeout'>((resolve) => { + const timer = setTimeout(() => { + sub.dispose(); + resolve('timeout'); + }, 120_000); + const sub = agent.events.on('prompt.completed', (event) => { + clearTimeout(timer); + sub.dispose(); + console.log( + `[event] ${elapsed()} prompt.completed -> reason=${event.reason ?? 'unknown'}`, + ); + resolve(event.reason === 'failed' ? 'failed' : 'completed'); + }); + }); + + // Diff-polled snapshot of the RPC-visible readings. + let lastKey = ''; + const snapshot = async (tag: string): Promise => { + const [ctx, usage] = await Promise.all([agent.getContext(), agent.getUsage()]); + const total = usageTotal(usage.total); + const turn = usageTotal(usage.currentTurn); + const key = `${String(ctx.history.length)}/${String(ctx.tokenCount)}/${String(total)}/${String(turn)}`; + if (key === lastKey) return; + lastKey = key; + console.log( + `[poll] ${elapsed()} ${tag}`.padEnd(46), + `history=${String(ctx.history.length)} tokenCount(measured)=${String(ctx.tokenCount)}` + + ` usage.total=${String(total)} usage.currentTurn=${String(turn)}`, + ); + }; + + let polling = true; + const pollLoop = (async (): Promise => { + while (polling) { + try { + await snapshot(''); + } catch { + // transient RPC failure during the turn — keep polling + } + await tick(250); + } + })(); + + await tick(300); // let the event subscriptions cross the lazily opened WS + + await snapshot('created (pre-prompt)'); + console.log(`[prompt] ${elapsed()} sending "hi"`); + await agent.prompt({ input: [{ type: 'text', text: 'hi' }] }); + + const outcome = await completed; + polling = false; + await pollLoop; + lastKey = ''; // force the final line even if nothing moved since the last poll tick + await snapshot('after prompt.completed'); + + const ctx = await agent.getContext(); + const usage = await agent.getUsage(); + const total = usageTotal(usage.total); + console.log('---'); + console.log('[result] outcome ->', outcome); + console.log('[result] history messages ->', ctx.history.length); + console.log('[result] tokenCount (measured) ->', ctx.tokenCount); + console.log('[result] usage.total ->', JSON.stringify(usage.total)); + console.log('[result] usage.byModel ->', JSON.stringify(usage.byModel)); + console.log( + `[check] tokenCount vs usage.total -> ${String(ctx.tokenCount)} vs ${String(total)}`, + ); + console.log( + '[note] reading guide:\n' + + ' - tokenCount is 0 until the first measured exchange lands, then it\n' + + ' should equal THAT exchange\'s total (input + output); new messages\n' + + ' appended between turns are the unmeasured tail.\n' + + ' - after one covered exchange on a fresh session, cumulative\n' + + ' usage.total and tokenCount should roughly agree; a large gap means\n' + + ' the measured total never made it onto the wire model and the reading\n' + + ' silently fell back to per-message estimates.\n' + + ' - outcome "timeout" means the turn finished its work but the\n' + + ' prompt.completed event never reached the client over the WS.', + ); + console.log('[note] session left open for inspection in the TUI ->', session.id); + + await klient.close(); + if (outcome === 'failed') process.exit(1); +} + +try { + await main(); +} catch (error) { + console.error(error); + process.exit(1); +} From 2090b65b6533ed603a884a798efb9d93a49b6023 Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Thu, 16 Jul 2026 18:26:06 +0800 Subject: [PATCH 3/3] chore: add changeset for the context size fix --- .changeset/tidy-humans-melt.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/tidy-humans-melt.md diff --git a/.changeset/tidy-humans-melt.md b/.changeset/tidy-humans-melt.md new file mode 100644 index 0000000000..76d652b845 --- /dev/null +++ b/.changeset/tidy-humans-melt.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix the context size indicator under-reporting the model's actual context usage.