vrr/kimi-code
2
0
Fork 0
mirror of https://github.com/MoonshotAI/kimi-code.git synced 2026-08-21 14:47:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

fix(kap-server): serve real session usage in snapshot and persist per-turn context readings (#3094)

Browse source 
* fix(kap-server): serve real session usage in snapshot and persist per-turn context readings

* fix(kap-server): omit unknown session usage fields instead of reporting zero
This commit is contained in:
Haozhe 2026-08-19 22:06:53 +08:00 • committed by GitHub
parent 056f02c2de
commit 4ff06f17e3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 388 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

2
packages/agent-core-v2/docs/state-manifest.d.ts vendored
View file

@ -1508,7 +1508,7 @@ export interface AgentStateSnapshot {
'task.notificationDelivery': readonly string[];
'task.scheduledNotificationKeys': Set<string>;
// src/agent/tokenCounting/tokenCountingOps.ts
// replayable · durable — folds: TokenCountingMeasured, TokenCountingTruncated, TokenCountingRebased
// replayable · durable — folds: TokenCountingMeasured, TokenCountingTruncated, TokenCountingRebased, TokenCountingTurnRecorded
'tokenCounting': /* TokenCountingState — packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts */ {
readonly anchors: readonly /* TokenAnchor — packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts */ {
readonly length: number;

15
packages/agent-core-v2/docs/wire-manifest.d.ts vendored
View file

@ -24,7 +24,7 @@
// cross-reducers), blobs (the folding states whose blob codec offloads inline
// media to blob storage), owner (the source file declaring the class).
// Index (52 record types)
// Index (53 record types)
// config.update profile src/agent/profile/profileOps.ts
// context.append_loop_event contextMemory, turn src/agent/contextMemory/contextEvents.ts
// context.append_message contextMemory, goalForkNotice, plan, task.notificationDelivery, todo src/agent/contextMemory/contextEvents.ts
@ -65,6 +65,7 @@
// token_counting.measured tokenCounting src/agent/tokenCounting/tokenCountingOps.ts
// token_counting.rebased tokenCounting src/agent/tokenCounting/tokenCountingOps.ts
// token_counting.truncated tokenCounting src/agent/tokenCounting/tokenCountingOps.ts
// token_counting.turn_recorded tokenCounting src/agent/tokenCounting/tokenCountingOps.ts
// tools.register_user_tool userTool src/agent/userTool/userToolOps.ts
// tools.reset_active_tools profile.activeTools src/agent/profile/profileOps.ts
// tools.set_active_tools profile.activeTools src/agent/profile/profileOps.ts
@ -575,6 +576,17 @@ interface TokenCountingTruncatedPayload {
tokens: number;
}
/**
* states: tokenCounting
* owner: src/agent/tokenCounting/tokenCountingOps.ts
*/
interface TokenCountingTurnRecordedPayload {
_name: 'token_counting.turn_recorded';
length: number;
tokens: number;
turnId: number;
}
/**
* states: userTool
* owner: src/agent/userTool/userToolOps.ts
@ -788,6 +800,7 @@ interface WirePayloadMap {
"token_counting.measured": TokenCountingMeasuredPayload;
"token_counting.rebased": TokenCountingRebasedPayload;
"token_counting.truncated": TokenCountingTruncatedPayload;
"token_counting.turn_recorded": TokenCountingTurnRecordedPayload;
"tools.register_user_tool": ToolsRegisterUserToolPayload;
"tools.reset_active_tools": ToolsResetActiveToolsPayload;
"tools.set_active_tools": ToolsSetActiveToolsPayload;

22
packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts
View file

@ -41,6 +41,15 @@ export class TokenCountingRebased extends Event2<z.infer<typeof rebaseSchema>> {
}
export interface TokenCountingRebased extends z.infer<typeof rebaseSchema> {}
const turnRecordedSchema = sizeSchema.extend({ turnId: z.number() });
export class TokenCountingTurnRecorded extends Event2<z.infer<typeof turnRecordedSchema>> {
static override readonly type = 'token_counting.turn_recorded';
static override readonly durable = true;
static override readonly schema = turnRecordedSchema;
}
export interface TokenCountingTurnRecorded extends z.infer<typeof turnRecordedSchema> {}
function anchorsEqual(a: readonly TokenAnchor[], b: readonly TokenAnchor[]): boolean {
return a.length === b.length && a.every((anchor, i) => anchor === b[i]);
}
@ -79,6 +88,19 @@ export const tokenCountingKey = defineState(
s.tokens = tokens;
}
ctx.emit(new AgentStatusUpdated({ contextTokens: s.tokens }));
})
.on(TokenCountingTurnRecorded, (s, e, ctx) => {
const length = normalizeAnchorLength(e.length);
const tokens = Math.max(0, e.tokens);
const pinned = s.anchors.some((a) => a.length === length);
const anchors = pinned
? s.anchors
: [...s.anchors.filter((a) => a.length < length), { length, tokens, measured: false }];
if (!(s.tokens === tokens && anchorsEqual(s.anchors, anchors))) {
s.anchors = anchors;
s.tokens = tokens;
}
ctx.emit(new AgentStatusUpdated({ contextTokens: s.tokens }));
});
function normalizeAnchorLength(length: number): number {

15
packages/agent-core-v2/src/agent/tokenCounting/tokenCountingService.ts
View file

@ -2,8 +2,10 @@ import { Disposable } from '#/_base/di/lifecycle';
import { LifecycleScope } from '#/app/scopes';
import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
import { IConfigService } from '#/app/config/config';
import { IEventBus } from '#/app/event/eventBus';
import { contextMemoryKey } from '#/agent/contextMemory/contextOps';
import type { ContextMessage } from '#/agent/contextMemory/types';
import { TurnEnded } from '#/agent/loop/turnOps';
import type { Message } from '#/kosong/contract/message';
import type { Tool } from '#/kosong/contract/tool';
import {
@ -25,6 +27,7 @@ import {
} from './tokenCounting';
import {
TokenCountingMeasured,
TokenCountingTurnRecorded,
tokenCountingKey,
type TokenAnchor,
} from './tokenCountingOps';
@ -38,9 +41,21 @@ export class AgentTokenCountingService extends Disposable implements IAgentToken
@IEventDispatcher private readonly dispatcher: IEventDispatcher,
@IConfigService private readonly config: IConfigService,
@IAgentStateService private readonly agentState: IAgentStateService,
@IEventBus private readonly eventBus: IEventBus,
) {
super();
this.agentState.contributeState(tokenCountingKey);
this._register(
this.eventBus.subscribe(TurnEnded, (e) => {
void this.dispatcher.dispatch(
new TokenCountingTurnRecorded({
turnId: e.turnId,
length: this.context().length,
tokens: this.statusSize(),
}),
);
}),
);
}
get strategy(): TokenCountingStrategy {

2
packages/agent-core-v2/src/app/sessionLegacy/sessionLegacyService.ts
View file

@ -87,7 +87,7 @@ export class SessionLegacyService implements ISessionLegacyService {
swarm_mode: swarm.isActive,
context_tokens: tokens,
max_context_tokens: maxTokens > 0 ? maxTokens : undefined,
context_usage: maxTokens > 0 ? Math.min(1, tokens / maxTokens) : 0,
context_usage: maxTokens > 0 ? Math.min(1, tokens / maxTokens) : undefined,
};
}

2
packages/agent-core-v2/src/app/sessionLegacy/sessionProtocol.ts
View file

@ -77,6 +77,6 @@ export const sessionStatusResponseSchema = z.object({
swarm_mode: z.boolean(),
context_tokens: z.number().int().nonnegative(),
max_context_tokens: z.number().int().nonnegative().optional(),
context_usage: z.number().min(0).max(1),
context_usage: z.number().min(0).max(1).optional(),
});
export type SessionStatusResponse = z.infer<typeof sessionStatusResponseSchema>;

69
packages/agent-core-v2/test/agent/tokenCounting/tokenCounting.test.ts
View file

@ -4,6 +4,7 @@ import { IAgentContextMemoryService, IAgentProfileService } from '#/index';
import { IAgentStateService } from '#/agent/state/agentState';
import { IAgentTokenCountingService } from '#/agent/tokenCounting/tokenCounting';
import { TokenCountingMeasured, tokenCountingKey } from '#/agent/tokenCounting/tokenCountingOps';
import { TurnEnded } from '#/agent/loop/turnOps';
import { estimateTokensForMessages } from '#/kosong/contract/tokens';
import type { TokenUsage } from '#/kosong/contract/usage';
import { IAgentUsageService } from '#/agent/usage/usage';
@ -244,4 +245,72 @@ describe('Agent token counting', () => {
Math.max(tokenCounting.get().size, tokenCounting.latestMeasured()),
);
});
it('journals the reported size as a durable record at every turn end', async () => {
const persistence = new InMemoryWireRecordPersistence();
const live = createTestAgent({ persistence });
try {
live.get(IAgentProfileService).update({ activeToolNames: [] });
live.mockNextResponse({ type: 'text', text: 'Hi there!' });
await live.rpc.prompt({ input: [{ type: 'text', text: 'hi' }] });
await live.untilTurnEnd();
const counting = live.get(IAgentTokenCountingService);
const reported = counting.statusSize();
expect(reported).toBeGreaterThan(0);
await live.get(IWireService).flush();
const records = persistence.records.filter(
(record) => record.type === 'token_counting.turn_recorded',
);
expect(records).toHaveLength(1);
expect(records[0]).toMatchObject({
length: live.get(IAgentContextMemoryService).get().length,
tokens: reported,
});
expect(live.get(IAgentStateService).get(tokenCountingKey).anchors).toEqual([
{ length: 2, tokens: reported, measured: true },
]);
} finally {
await live.dispose();
}
});
it('pins the reported size at turn end when no measured anchor covers it', async () => {
ctx.appendUserMessage([{ type: 'text', text: 'unmeasured tail' }]);
const expected = tokenCounting.statusSize();
expect(expected).toBeGreaterThan(0);
expect(ctx.agentState.get(tokenCountingKey).anchors).toEqual([]);
await ctx.dispatcher.dispatch(new TurnEnded({ turnId: 1, reason: 'completed' }));
expect(ctx.agentState.get(tokenCountingKey).anchors).toEqual([
{ length: 1, tokens: expected, measured: false },
]);
expect(tokenCounting.statusSize()).toBe(expected);
});
it('drops the pinned turn reading on compaction', async () => {
ctx.appendUserMessage([{ type: 'text', text: 'unmeasured tail' }]);
await ctx.dispatcher.dispatch(new TurnEnded({ turnId: 1, reason: 'completed' }));
expect(ctx.agentState.get(tokenCountingKey).anchors).toHaveLength(1);
context.applyCompaction({
summary: 'summary of the tail',
compactedCount: 1,
tokensBefore: 100,
summaryOutputTokens: 50,
});
const history = context.get();
const anchors = ctx.agentState.get(tokenCountingKey).anchors;
expect(anchors).toHaveLength(1);
expect(anchors[0]).toEqual({
length: history.length,
tokens: tokenCounting.get().size,
measured: false,
});
expect(tokenCounting.statusSize()).toBe(tokenCounting.get().size);
});
});

50
packages/agent-core-v2/test/app/config/config.test.ts
View file

@ -393,30 +393,32 @@ describe('Agent config', () => {
await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Start a fresh turn' }] });
expect(await ctx.untilTurnEnd()).toMatchInlineSnapshot(`
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "lastTurn": { "turnId": 0, "reason": "completed", "at": "<time>" }, "background": [] }
[emit] prompt.completed { "time": "<time>", "promptId": "<msg-1>", "finishedAt": "<time>", "reason": "completed" }
[wire] prompt.accepted { "promptId": "<msg-2>", "time": "<time>" }
[wire] turn.prompt { "input": [ { "type": "text", "text": "Start a fresh turn" } ], "origin": { "kind": "user" }, "time": "<time>" }
[emit] turn.started { "time": "<time>", "turnId": 1, "origin": { "kind": "user" }, "prompt": "Start a fresh turn" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "running", "step": 0, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[emit] context.spliced { "time": "<time>", "start": 4, "deleteCount": 0, "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Start a fresh turn" } ], "toolCalls": [], "origin": { "kind": "user" }, "id": "<msg-2>" } ] }
[wire] context.append_message { "message": { "role": "user", "content": [ { "type": "text", "text": "Start a fresh turn" } ], "toolCalls": [], "origin": { "kind": "user" }, "id": "<msg-2>" }, "time": "<time>" }
[emit] turn.step.started { "time": "<time>", "turnId": 1, "step": 1, "stepId": "<uuid-6>" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "running", "step": 1, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[wire] context.append_loop_event { "event": { "type": "step.begin", "uuid": "<uuid-6>", "turnId": "1", "step": 1 }, "time": "<time>" }
[wire] llm.request { "kind": "loop", "provider": "openai", "model": "changed-model", "modelAlias": "changed-model", "thinkingEffort": "off", "maxTokens": 1000000, "toolSelect": false, "systemPromptHash": "7617cb8b42659214c397a1d7505fce204b673b078a10de8bcccc697d88dcda56", "toolsHash": "4f53cda18c2baa0c0354bb5f9a3ecbe5ed12ab4d8e11ba873c2f11161202b945", "messageCount": 5, "turnStep": "1.1", "time": "<time>" }
[emit] assistant.delta { "time": "<time>", "turnId": 1, "delta": "Now the changed config is active." }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "streaming", "stream": "assistant", "step": 1, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[wire] usage.record { "model": "changed-model", "usage": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "usageScope": "turn", "time": "<time>" }
[emit] agent.status.updated { "time": "<time>", "usage": { "byModel": { "mock-model": { "inputOther": 40, "output": 30, "inputCacheRead": 0, "inputCacheCreation": 0 }, "changed-model": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 } }, "total": { "inputOther": 90, "output": 42, "inputCacheRead": 0, "inputCacheCreation": 0 }, "currentTurn": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 } } }
[wire] token_counting.measured { "length": 6, "tokens": 62, "time": "<time>" }
[emit] agent.status.updated { "time": "<time>", "contextTokens": 62 }
[emit] turn.step.completed { "time": "<time>", "turnId": 1, "step": 1, "stepId": "<uuid-6>", "usage": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "end_turn", "providerFinishReason": "completed", "rawFinishReason": "stop" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "running", "step": 1, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[wire] context.append_loop_event { "event": { "type": "content.part", "uuid": "<uuid-7>", "turnId": "1", "step": 1, "stepUuid": "<uuid-6>", "part": { "type": "text", "text": "Now the changed config is active." } }, "time": "<time>" }
[wire] context.append_loop_event { "event": { "type": "step.end", "uuid": "<uuid-6>", "turnId": "1", "step": 1, "finishReason": "end_turn", "usage": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "messageId": "mock-3", "providerFinishReason": "completed", "rawFinishReason": "stop" }, "time": "<time>" }
[wire] turn.ended { "turnId": 1, "reason": "completed", "time": "<time>" }
[emit] turn.ended { "time": "<time>", "turnId": 1, "reason": "completed" }
[wire] token_counting.turn_recorded { "turnId": 0, "length": 4, "tokens": 44, "time": "<time>" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "lastTurn": { "turnId": 0, "reason": "completed", "at": "<time>" }, "background": [] }
[emit] agent.status.updated { "time": "<time>", "contextTokens": 44 }
[emit] prompt.completed { "time": "<time>", "promptId": "<msg-1>", "finishedAt": "<time>", "reason": "completed" }
[wire] prompt.accepted { "promptId": "<msg-2>", "time": "<time>" }
[wire] turn.prompt { "input": [ { "type": "text", "text": "Start a fresh turn" } ], "origin": { "kind": "user" }, "time": "<time>" }
[emit] turn.started { "time": "<time>", "turnId": 1, "origin": { "kind": "user" }, "prompt": "Start a fresh turn" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "running", "step": 0, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[emit] context.spliced { "time": "<time>", "start": 4, "deleteCount": 0, "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Start a fresh turn" } ], "toolCalls": [], "origin": { "kind": "user" }, "id": "<msg-2>" } ] }
[wire] context.append_message { "message": { "role": "user", "content": [ { "type": "text", "text": "Start a fresh turn" } ], "toolCalls": [], "origin": { "kind": "user" }, "id": "<msg-2>" }, "time": "<time>" }
[emit] turn.step.started { "time": "<time>", "turnId": 1, "step": 1, "stepId": "<uuid-6>" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "running", "step": 1, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[wire] context.append_loop_event { "event": { "type": "step.begin", "uuid": "<uuid-6>", "turnId": "1", "step": 1 }, "time": "<time>" }
[wire] llm.request { "kind": "loop", "provider": "openai", "model": "changed-model", "modelAlias": "changed-model", "thinkingEffort": "off", "maxTokens": 1000000, "toolSelect": false, "systemPromptHash": "7617cb8b42659214c397a1d7505fce204b673b078a10de8bcccc697d88dcda56", "toolsHash": "4f53cda18c2baa0c0354bb5f9a3ecbe5ed12ab4d8e11ba873c2f11161202b945", "messageCount": 5, "turnStep": "1.1", "time": "<time>" }
[emit] assistant.delta { "time": "<time>", "turnId": 1, "delta": "Now the changed config is active." }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "streaming", "stream": "assistant", "step": 1, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[wire] usage.record { "model": "changed-model", "usage": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "usageScope": "turn", "time": "<time>" }
[emit] agent.status.updated { "time": "<time>", "usage": { "byModel": { "mock-model": { "inputOther": 40, "output": 30, "inputCacheRead": 0, "inputCacheCreation": 0 }, "changed-model": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 } }, "total": { "inputOther": 90, "output": 42, "inputCacheRead": 0, "inputCacheCreation": 0 }, "currentTurn": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 } } }
[wire] token_counting.measured { "length": 6, "tokens": 62, "time": "<time>" }
[emit] agent.status.updated { "time": "<time>", "contextTokens": 62 }
[emit] turn.step.completed { "time": "<time>", "turnId": 1, "step": 1, "stepId": "<uuid-6>", "usage": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "finishReason": "end_turn", "providerFinishReason": "completed", "rawFinishReason": "stop" }
[emit] agent.activity.updated { "time": "<time>", "lifecycle": "ready", "turn": { "turnId": 1, "origin": { "kind": "user" }, "phase": "running", "step": 1, "ending": false, "pendingApprovals": [], "activeToolCalls": [], "since": "<time>" }, "background": [] }
[wire] context.append_loop_event { "event": { "type": "content.part", "uuid": "<uuid-7>", "turnId": "1", "step": 1, "stepUuid": "<uuid-6>", "part": { "type": "text", "text": "Now the changed config is active." } }, "time": "<time>" }
[wire] context.append_loop_event { "event": { "type": "step.end", "uuid": "<uuid-6>", "turnId": "1", "step": 1, "finishReason": "end_turn", "usage": { "inputOther": 50, "output": 12, "inputCacheRead": 0, "inputCacheCreation": 0 }, "messageId": "mock-3", "providerFinishReason": "completed", "rawFinishReason": "stop" }, "time": "<time>" }
[wire] turn.ended { "turnId": 1, "reason": "completed", "time": "<time>" }
[emit] turn.ended { "time": "<time>", "turnId": 1, "reason": "completed" }
`);
expect(ctx.lastLlmInput()).toMatchInlineSnapshot(`
system: "Changed system prompt."

1
packages/agent-core-v2/test/index.test.ts
View file

@ -97,6 +97,7 @@ const V2_RECORD_TYPES: ReadonlySet<string> = new Set([
'cron.add',
'cron.delete',
'cron.cursor',
'token_counting.turn_recorded',
]);
describe('v1 wire vocabulary', () => {

54
packages/agent-core-v2/test/tool/tool.test.ts
View file

File diff suppressed because one or more lines are too long

6
packages/kap-server/src/protocol/session.ts
View file

@ -15,10 +15,10 @@ export const sessionUsageSchema = z.object({
output_tokens: z.number().int().nonnegative(),
cache_read_tokens: z.number().int().nonnegative(),
cache_creation_tokens: z.number().int().nonnegative(),
total_cost_usd: z.number().nonnegative(),
total_cost_usd: z.number().nonnegative().optional(),
context_tokens: z.number().int().nonnegative(),
context_limit: z.number().int().nonnegative(),
turn_count: z.number().int().nonnegative(),
context_limit: z.number().int().nonnegative().optional(),
turn_count: z.number().int().nonnegative().optional(),
});
export type SessionUsage = z.infer<typeof sessionUsageSchema>;

34
packages/kap-server/src/routes/snapshot.ts
View file

@ -19,6 +19,11 @@ import {
type InFlightTurn,
type SessionSnapshotResponse,
} from '../protocol/rest-snapshot';
import { emptySessionUsage, type SessionUsage } from '../protocol/session';
import {
readLegacyStatus,
type LegacyStatusSnapshot,
} from '../services/legacyStatus/legacyStatus';
import { loadMessageHistory } from '../services/messages/messageHistory';
import { type SessionEventBroadcaster } from '../transport/ws/v1/sessionEventBroadcaster';
import { toWireApproval } from './approvals';
@ -104,13 +109,19 @@ async function assembleSnapshot(
const workspace = await core.accessor.get(IWorkspaceService).get(workspaceId);
const cwd = workspace?.root ?? '';
const meta = await handle.accessor.get(ISessionMetadata).read();
const session = toWireSession(
{ ...meta, workspaceId },
cwd,
resolveSessionFacts(core, sessionId),
);
const main = await ensureMainAgent(handle);
const status = readLegacyStatus(main);
const session = {
...toWireSession(
{ ...meta, workspaceId },
cwd,
resolveSessionFacts(core, sessionId),
),
agent_config: { model: status?.model ?? '' },
usage: toSnapshotUsage(status),
};
const all = await loadMessageHistory(core, main, sessionId, meta.createdAt);
const hasMore = all.length > SNAPSHOT_MESSAGE_PAGE_SIZE;
const items = all.slice(-SNAPSHOT_MESSAGE_PAGE_SIZE);
@ -147,6 +158,19 @@ function readCurrentPromptId(main: IAgentScopeHandle | undefined): string | unde
}
}
function toSnapshotUsage(status: LegacyStatusSnapshot | undefined): SessionUsage {
if (status === undefined) return emptySessionUsage();
const total = status.usage?.total;
return {
input_tokens: total?.inputOther ?? 0,
output_tokens: total?.output ?? 0,
cache_read_tokens: total?.inputCacheRead ?? 0,
cache_creation_tokens: total?.inputCacheCreation ?? 0,
context_tokens: status.contextTokens,
context_limit: status.maxContextTokens,
};
}
function attachCurrentPromptIdToInFlight(
inFlightTurn: InFlightTurn | null,
currentPromptId: string | undefined,

174
packages/kap-server/test/snapshot.test.ts
View file

@ -10,7 +10,10 @@ import {
IAppendLogStore,
IEventBus,
IAgentLifecycleService,
IAgentProfileService,
IAgentPromptService,
IAgentTokenCountingService,
IAgentUsageService,
ISessionInteractionService,
ISessionContext,
ISessionIndex,
@ -24,6 +27,7 @@ import {
resumeSessionById,
} from '@moonshot-ai/agent-core-v2';
import { sessionSnapshotResponseSchema } from '../src/protocol/rest-snapshot';
import { emptySessionUsage } from '../src/protocol/session';
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
import { registerSnapshotRoutes } from '../src/routes/snapshot';
@ -59,6 +63,22 @@ describe('server-v2 snapshot route enrichment', () => {
[IWireService, { flush: async () => {} }],
[IAgentScopeContext, { scope: () => 'scope/sess_snapshot' }],
[IAgentBlobService, { loadParts: async (parts: unknown) => parts }],
[
IAgentProfileService,
{
getModelCapabilities: () => ({ max_input_tokens: 262144 }),
getModel: () => 'kimi-for-test',
},
],
[
IAgentUsageService,
{
status: () => ({
total: { inputOther: 120, output: 34, inputCacheRead: 56, inputCacheCreation: 7 },
}),
},
],
[IAgentTokenCountingService, { statusSize: () => 4321 }],
]),
};
const session = {
@ -184,6 +204,15 @@ describe('server-v2 snapshot route enrichment', () => {
assistant_text: 'Hello',
current_prompt_id: promptId,
});
expect(snap.session.usage).toEqual({
input_tokens: 120,
output_tokens: 34,
cache_read_tokens: 56,
cache_creation_tokens: 7,
context_tokens: 4321,
context_limit: 262144,
});
expect(snap.session.agent_config.model).toBe('kimi-for-test');
expect(snap.subagents).toEqual([
expect.objectContaining({
id: 'agent-1',
@ -195,6 +224,125 @@ describe('server-v2 snapshot route enrichment', () => {
}),
]);
});
it('keeps the placeholder usage when the main agent exposes no status services', async () => {
const sessionId = 'sess_snapshot_degraded';
const workspaceId = 'wd_snapshot_abcdef012345';
const now = Date.parse('2026-01-01T00:00:00.000Z');
const main = {
accessor: fakeAccessor([
[IAgentContextMemoryService, { get: () => [] }],
[IWireService, { flush: async () => {} }],
[IAgentScopeContext, { scope: () => 'scope/sess_snapshot_degraded' }],
[IAgentBlobService, { loadParts: async (parts: unknown) => parts }],
[IAgentProfileService, undefined],
[IAgentUsageService, undefined],
[IAgentTokenCountingService, undefined],
]),
};
const session = {
accessor: fakeAccessor([
[ISessionContext, { workspaceId }],
[
ISessionMetadata,
{
read: async () => ({
id: sessionId,
title: 'Snapshot degraded',
createdAt: now,
updatedAt: now,
archived: false,
}),
},
],
[IAgentLifecycleService, { get: () => main, create: async () => main }],
[ISessionInteractionService, { listPending: () => [] }],
]),
};
const handler = {
accessor: fakeAccessor([
[
ISessionLifecycleService,
{ resume: async () => session, get: () => undefined },
],
]),
};
const core = {
accessor: fakeAccessor([
[
ISessionIndex,
{
get: async () => ({
id: sessionId,
workspaceId,
cwd: '/workspace',
createdAt: now,
updatedAt: now,
archived: false,
}),
},
],
[
ISessionManager,
{
resume: async () => session,
get: () => undefined,
list: () => [],
},
],
[IWorkspaceService, { get: async () => ({ root: '/workspace' }) }],
[ITelemetryService, { withContext: () => ({ track2: () => {} }) }],
[
IAppendLogStore,
{
read: async function* () {},
},
],
]),
};
const broadcaster = {
getSnapshotState: async () => ({
seq: 1,
epoch: 'ep_snapshot',
inFlightTurn: null,
subagents: [],
}),
};
let routeHandler:
| ((
req: { id: string; params: { session_id: string } },
reply: { send(payload: unknown): unknown },
) => Promise<void> | void)
| undefined;
registerSnapshotRoutes(
{
get: (_path, _options, handler) => {
routeHandler = handler;
},
},
{
core: core as never,
broadcaster: broadcaster as never,
},
);
let payload: unknown;
await routeHandler?.(
{ id: 'req_snapshot_degraded', params: { session_id: sessionId } },
{
send: (value) => {
payload = value;
},
},
);
const body = payload as { code: number; data: unknown };
expect(body.code).toBe(0);
const snap = sessionSnapshotResponseSchema.parse(body.data);
expect(snap.session.usage).toEqual(emptySessionUsage());
expect(snap.session.agent_config.model).toBe('');
});
});
describe('server-v2 GET /api/v1/sessions/:id/snapshot', () => {
@ -283,6 +431,32 @@ describe('server-v2 GET /api/v1/sessions/:id/snapshot', () => {
});
});
it('serves the real usage ledger instead of the zero placeholder', async () => {
const sid = await createSession();
await ensureMainAgent(sid);
const session = getLiveSessionById(server!.core.accessor, sid);
const main = session!.accessor.get(IAgentLifecycleService).get('main')!;
main.accessor.get(IAgentUsageService).record('kimi-for-test', {
inputOther: 120,
output: 34,
inputCacheRead: 56,
inputCacheCreation: 7,
});
main.accessor.get(IAgentContextMemoryService).append({
role: 'user',
content: [{ type: 'text', text: 'hello' }],
toolCalls: [],
});
const snap = await snapshot(sid);
expect(snap.session.usage.input_tokens).toBe(120);
expect(snap.session.usage.output_tokens).toBe(34);
expect(snap.session.usage.cache_read_tokens).toBe(56);
expect(snap.session.usage.cache_creation_tokens).toBe(7);
expect(snap.session.usage.context_tokens).toBeGreaterThan(0);
expect(snap.session.usage.context_limit).toBeUndefined();
});
it('returns 404 for an unknown session', async () => {
const res = await fetch(`${base}/api/v1/sessions/sess_does_not_exist/snapshot`, {
headers: authHeaders(server as RunningServer),

2
packages/protocol/src/rest/session.ts
View file

@ -143,7 +143,7 @@ export const sessionStatusResponseSchema = z.object({
/** Omitted when the context limit is unknown — 0 is the engine's "unknown"
* marker, never a real limit. */
max_context_tokens: z.number().int().nonnegative().optional(),
context_usage: z.number().min(0).max(1),
context_usage: z.number().min(0).max(1).optional(),
});
export type SessionStatusResponse = z.infer<typeof sessionStatusResponseSchema>;

6
packages/protocol/src/session.ts
View file

@ -12,10 +12,10 @@ export const sessionUsageSchema = z.object({
output_tokens: z.number().int().nonnegative(),
cache_read_tokens: z.number().int().nonnegative(),
cache_creation_tokens: z.number().int().nonnegative(),
total_cost_usd: z.number().nonnegative(),
total_cost_usd: z.number().nonnegative().optional(),
context_tokens: z.number().int().nonnegative(),
context_limit: z.number().int().nonnegative(),
turn_count: z.number().int().nonnegative(),
context_limit: z.number().int().nonnegative().optional(),
turn_count: z.number().int().nonnegative().optional(),
});
export type SessionUsage = z.infer<typeof sessionUsageSchema>;