mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-07-10 01:39:25 +00:00
refactor(agent-core-v2): replace contextSize getStatus with ranged get
IAgentContextSizeService.getStatus() becomes get(start?, end?) returning { size, measured, estimated } where size = measured + estimated. start/end follow Array.prototype.slice semantics (default whole context, negative indices count from the end, inverted range is empty). Update consumers, the context:status action binding, tests, and the edge-exposure cheatsheet.
This commit is contained in:
parent
1a7b26472b
commit
cf9bca54e4
13 changed files with 124 additions and 55 deletions
|
|
@ -103,7 +103,7 @@ Read = `GET`, write = `POST`. `sid` = `session_id`, `aid` = `agent_id`.
|
|||
| `tasks` | `list` / `get` / `readOutput` | IBackgroundService.* | GET |
|
||||
| `tasks` | `stop` / `detach` | IBackgroundService.* | POST |
|
||||
| `usage` | `status` | IUsageService.status | GET |
|
||||
| `context` | `status` | IContextSizeService.getStatus | GET |
|
||||
| `context` | `status` | IAgentContextSizeService.get | GET |
|
||||
| `swarm` | `isActive` | ISwarmService.isActive | GET |
|
||||
| `swarm` | `enter` / `exit` | ISwarmService.* | POST |
|
||||
| `permission` | `getMode` | IPermissionModeService.mode | GET |
|
||||
|
|
|
|||
|
|
@ -2,15 +2,16 @@ import { createDecorator } from '#/_base/di/instantiation';
|
|||
import type { Message } from '#/app/llmProtocol/message';
|
||||
import type { TokenUsage } from '#/app/llmProtocol/usage';
|
||||
|
||||
export interface ContextSizeStatus {
|
||||
readonly contextTokens: number;
|
||||
readonly contextTokensWithPending: number;
|
||||
export interface ContextSize {
|
||||
readonly size: number;
|
||||
readonly measured: number;
|
||||
readonly estimated: number;
|
||||
}
|
||||
|
||||
export interface IAgentContextSizeService {
|
||||
readonly _serviceBrand: undefined;
|
||||
|
||||
getStatus(): ContextSizeStatus;
|
||||
get(start?: number, end?: number): ContextSize;
|
||||
measured(input: readonly Message[], output: readonly Message[], usage: TokenUsage): void;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,25 +6,29 @@
|
|||
* `wire.dispatch(contextSizeMeasured(...))` (called by `llmRequester` after each
|
||||
* measured exchange), and emits the `contextTokens` slice of
|
||||
* `agent.status.updated` live through `wire.signal` when the measured value
|
||||
* changes. `getStatus().contextTokens` is the deterministic measured value
|
||||
* (replay-safe); `contextTokensWithPending` adds the live token estimate of the
|
||||
* not-yet-measured tail, computed at read time from the surviving
|
||||
* `contextMemory` messages beyond the measured prefix — the sparse
|
||||
* `measuredPrefixTokens` / per-message `estimates` are deliberately not
|
||||
* persisted (see `contextSizeOps`). Bound at Agent scope.
|
||||
* changes. `get(start?, end?)` returns `{ size, measured, estimated }` for the
|
||||
* context-message range `[start, end)`, resolved like `Array.prototype.slice`
|
||||
* (defaulting to the whole context; negative indices count back from the end;
|
||||
* an inverted range is empty): `measured`
|
||||
* is the deterministic measured value of the measured-prefix portion
|
||||
* (replay-safe; the exact aggregate is only known for the full prefix, so
|
||||
* sub-ranges fall back to a per-message estimate), `estimated` is the live token
|
||||
* estimate of the not-yet-measured portion, and `size = measured + estimated`.
|
||||
* The sparse `measuredPrefixTokens` / per-message `estimates` are deliberately
|
||||
* not persisted (see `contextSizeOps`). Bound at Agent scope.
|
||||
*/
|
||||
|
||||
import { Disposable } from '#/_base/di/lifecycle';
|
||||
import { InstantiationType } from '#/_base/di/extensions';
|
||||
import { LifecycleScope, registerScopedService } from '#/_base/di/scope';
|
||||
import { estimateTokensForMessage } from '#/_base/utils/tokens';
|
||||
import { estimateTokensForMessages } from '#/_base/utils/tokens';
|
||||
import type { ContextMessage } from '#/agent/contextMemory';
|
||||
import { IAgentContextMemoryService } from '#/agent/contextMemory';
|
||||
import type { Message } from '#/app/llmProtocol/message';
|
||||
import type { TokenUsage } from '#/app/llmProtocol/usage';
|
||||
import { IAgentWireService, type IWireService } from '#/wire';
|
||||
|
||||
import { IAgentContextSizeService, type ContextSizeStatus } from './contextSize';
|
||||
import { IAgentContextSizeService, type ContextSize } from './contextSize';
|
||||
import { ContextSizeModel, contextSizeMeasured } from './contextSizeOps';
|
||||
|
||||
export class AgentContextSizeService extends Disposable implements IAgentContextSizeService {
|
||||
|
|
@ -39,13 +43,23 @@ export class AgentContextSizeService extends Disposable implements IAgentContext
|
|||
super();
|
||||
}
|
||||
|
||||
getStatus(): ContextSizeStatus {
|
||||
const measured = this.wire.getModel(ContextSizeModel);
|
||||
const pendingTokens = estimateTail(this.context.get(), measured.length);
|
||||
return {
|
||||
contextTokens: measured.tokens,
|
||||
contextTokensWithPending: measured.tokens + pendingTokens,
|
||||
};
|
||||
get(start?: number, end?: number): ContextSize {
|
||||
const context = this.context.get();
|
||||
const model = this.wire.getModel(ContextSizeModel);
|
||||
// Mirrors `Array.prototype.slice`: defaults to the whole context, negative
|
||||
// indices count back from the end, and an inverted range is empty.
|
||||
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);
|
||||
// The measured-prefix total is the only deterministic measured value; use it
|
||||
// when the range covers the whole prefix, otherwise estimate the sub-range.
|
||||
const measured =
|
||||
from === 0 && measuredEnd === model.length
|
||||
? model.tokens
|
||||
: estimateTokensForMessages(context.slice(from, measuredEnd));
|
||||
const estimated = estimateTokensForMessages(context.slice(estimatedStart, to));
|
||||
return { size: measured + estimated, measured, estimated };
|
||||
}
|
||||
|
||||
measured(input: readonly Message[], output: readonly Message[], usage: TokenUsage): void {
|
||||
|
|
@ -78,16 +92,9 @@ function tokenUsageTotal(usage: TokenUsage): number {
|
|||
return usage.inputCacheRead + usage.inputCacheCreation + usage.inputOther + usage.output;
|
||||
}
|
||||
|
||||
function estimateTail(
|
||||
context: readonly ContextMessage[],
|
||||
measuredLength: number,
|
||||
): number {
|
||||
let total = 0;
|
||||
for (let index = measuredLength; index < context.length; index += 1) {
|
||||
const message = context[index];
|
||||
if (message !== undefined) total += estimateTokensForMessage(message);
|
||||
}
|
||||
return total;
|
||||
function normalizeSliceIndex(index: number, length: number): number {
|
||||
if (index < 0) return Math.max(length + index, 0);
|
||||
return Math.min(index, length);
|
||||
}
|
||||
|
||||
registerScopedService(
|
||||
|
|
|
|||
|
|
@ -590,8 +590,6 @@ function isTodoItem(value: unknown): value is TodoItem {
|
|||
);
|
||||
}
|
||||
|
||||
export { AgentFullCompactionService as FullCompaction };
|
||||
|
||||
// Construct eagerly (not delayed): the service registers turn and loop hooks
|
||||
// (onLaunched / beforeStep / afterStep / onError) that drive auto
|
||||
// compaction. With delayed instantiation the eager `accessor.get(IAgentFullCompactionService)`
|
||||
|
|
|
|||
|
|
@ -285,7 +285,7 @@ export class AgentLLMRequesterService implements IAgentLLMRequesterService {
|
|||
// independently and would be squeezed to nothing at high water marks.
|
||||
usedContextTokens:
|
||||
overrides.messages === undefined
|
||||
? this.contextSize.getStatus().contextTokens
|
||||
? this.contextSize.get().measured
|
||||
: undefined,
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -216,7 +216,7 @@ export class AgentMicroCompactionService
|
|||
private contextSizeRatio(): number {
|
||||
const maxContextTokens = this.profile.getModelCapabilities().max_context_tokens;
|
||||
if (maxContextTokens === undefined || maxContextTokens <= 0) return 1;
|
||||
return this.contextSize.getStatus().contextTokensWithPending / maxContextTokens;
|
||||
return this.contextSize.get().size / maxContextTokens;
|
||||
}
|
||||
|
||||
private measureEffect(
|
||||
|
|
|
|||
|
|
@ -380,7 +380,7 @@ export class AgentRPCService implements IAgentRPCService {
|
|||
getContext(_payload: EmptyPayload) {
|
||||
return {
|
||||
history: this.context.get(),
|
||||
tokenCount: this.contextSize.getStatus().contextTokens,
|
||||
tokenCount: this.contextSize.get().measured,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -404,7 +404,7 @@ export class SessionLegacyService implements ISessionLegacyService {
|
|||
const model = profile.getModel();
|
||||
const caps = profile.getModelCapabilities() as { max_context_tokens?: number };
|
||||
const maxTokens = caps.max_context_tokens ?? 0;
|
||||
const tokens = contextSize.getStatus().contextTokens;
|
||||
const tokens = contextSize.get().measured;
|
||||
const planData = await plan.status();
|
||||
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -545,13 +545,13 @@ describe('Agent context', () => {
|
|||
|
||||
it('includes new user messages as pending until the next usage update', () => {
|
||||
ctx.appendAssistantTextWithUsage(1, 'previous answer', 1_000);
|
||||
expect(contextSize.getStatus().contextTokens).toBe(1_000);
|
||||
expect(contextSize.get().measured).toBe(1_000);
|
||||
|
||||
ctx.appendUserMessage([{ type: 'text', text: 'next user prompt'.repeat(20) }]);
|
||||
|
||||
const pendingMessages = context.get().slice(-1);
|
||||
expect(contextSize.getStatus().contextTokensWithPending).toBe(
|
||||
contextSize.getStatus().contextTokens + estimateTokensForMessages(pendingMessages),
|
||||
expect(contextSize.get().size).toBe(
|
||||
contextSize.get().measured + estimateTokensForMessages(pendingMessages),
|
||||
);
|
||||
});
|
||||
|
||||
|
|
@ -582,24 +582,86 @@ describe('Agent context', () => {
|
|||
]);
|
||||
|
||||
const pendingMessages = context.get().slice(-1);
|
||||
expect(contextSize.getStatus().contextTokens).toBe(1_280);
|
||||
expect(contextSize.getStatus().contextTokensWithPending).toBe(
|
||||
expect(contextSize.get().measured).toBe(1_280);
|
||||
expect(contextSize.get().size).toBe(
|
||||
1_280 + estimateTokensForMessages(pendingMessages),
|
||||
);
|
||||
});
|
||||
|
||||
it('keeps zero-usage steps pending instead of zeroing tokenCount', () => {
|
||||
ctx.appendAssistantTextWithUsage(1, 'previous answer', 1_000);
|
||||
expect(contextSize.getStatus().contextTokens).toBe(1_000);
|
||||
expect(contextSize.get().measured).toBe(1_000);
|
||||
|
||||
ctx.appendUserMessage([{ type: 'text', text: 'next prompt' }]);
|
||||
|
||||
expect(contextSize.getStatus().contextTokens).toBe(1_000);
|
||||
expect(contextSize.getStatus().contextTokensWithPending).toBeGreaterThanOrEqual(
|
||||
contextSize.getStatus().contextTokens,
|
||||
expect(contextSize.get().measured).toBe(1_000);
|
||||
expect(contextSize.get().size).toBeGreaterThanOrEqual(
|
||||
contextSize.get().measured,
|
||||
);
|
||||
});
|
||||
|
||||
it('get(start, end) returns the size of a context-message range', () => {
|
||||
ctx.appendAssistantTextWithUsage(1, 'previous answer', 1_000);
|
||||
// The measured prefix covers the user + assistant pair (2 messages, 1_000 tokens).
|
||||
expect(contextSize.get()).toEqual({ size: 1_000, measured: 1_000, estimated: 0 });
|
||||
|
||||
ctx.appendUserMessage([{ type: 'text', text: 'pending one'.repeat(20) }]);
|
||||
ctx.appendUserMessage([{ type: 'text', text: 'pending two'.repeat(20) }]);
|
||||
|
||||
const messages = context.get();
|
||||
const tailEstimate = estimateTokensForMessages(messages.slice(2));
|
||||
|
||||
// Whole context: measured prefix + estimated tail.
|
||||
expect(contextSize.get()).toEqual({
|
||||
size: 1_000 + tailEstimate,
|
||||
measured: 1_000,
|
||||
estimated: tailEstimate,
|
||||
});
|
||||
|
||||
// A range fully inside the pending tail is purely estimated.
|
||||
const firstPending = estimateTokensForMessages(messages.slice(2, 3));
|
||||
expect(contextSize.get(2, 3)).toEqual({
|
||||
size: firstPending,
|
||||
measured: 0,
|
||||
estimated: firstPending,
|
||||
});
|
||||
|
||||
// The full measured prefix uses the deterministic aggregate.
|
||||
expect(contextSize.get(0, 2)).toEqual({ size: 1_000, measured: 1_000, estimated: 0 });
|
||||
|
||||
// A sub-range of the prefix falls back to a per-message estimate.
|
||||
const prefixHead = estimateTokensForMessages(messages.slice(0, 1));
|
||||
expect(contextSize.get(0, 1)).toEqual({
|
||||
size: prefixHead,
|
||||
measured: prefixHead,
|
||||
estimated: 0,
|
||||
});
|
||||
|
||||
// A range spanning the measured/tail boundary splits both sides.
|
||||
const assistant = estimateTokensForMessages(messages.slice(1, 2));
|
||||
expect(contextSize.get(1, 3)).toEqual({
|
||||
size: assistant + firstPending,
|
||||
measured: assistant,
|
||||
estimated: firstPending,
|
||||
});
|
||||
|
||||
// Negative indices resolve like `Array.prototype.slice`.
|
||||
expect(contextSize.get(-2)).toEqual({
|
||||
size: tailEstimate,
|
||||
measured: 0,
|
||||
estimated: tailEstimate,
|
||||
});
|
||||
expect(contextSize.get(0, -2)).toEqual({ size: 1_000, measured: 1_000, estimated: 0 });
|
||||
expect(contextSize.get(-3, -1)).toEqual({
|
||||
size: assistant + firstPending,
|
||||
measured: assistant,
|
||||
estimated: firstPending,
|
||||
});
|
||||
|
||||
// An inverted range is empty.
|
||||
expect(contextSize.get(-1, -3)).toEqual({ size: 0, measured: 0, estimated: 0 });
|
||||
});
|
||||
|
||||
it('undo only counts real user prompts, skipping task notifications', () => {
|
||||
ctx.appendAssistantText(1, 'first response');
|
||||
ctx.appendAssistantText(2, 'second response');
|
||||
|
|
|
|||
|
|
@ -1180,7 +1180,7 @@ export class AgentTestContext {
|
|||
const microCompaction = this.get(IAgentMicroCompactionService);
|
||||
void microCompaction;
|
||||
void swarm.isActive;
|
||||
contextSize.getStatus();
|
||||
contextSize.get();
|
||||
usage.status();
|
||||
toolStore.data();
|
||||
tasks.list(false);
|
||||
|
|
@ -1249,7 +1249,7 @@ export class AgentTestContext {
|
|||
const contextSize = this.get(IAgentContextSizeService);
|
||||
return {
|
||||
history: context.get(),
|
||||
tokenCount: contextSize.getStatus().contextTokens,
|
||||
tokenCount: contextSize.get().measured,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -626,11 +626,11 @@ describe('MicroCompaction', () => {
|
|||
const rawPending = ctx.context.get().slice(-1);
|
||||
const projectedPending = (ctx.get(IAgentMicroCompactionService) as any).compact(rawPending);
|
||||
expect(textOf(projectedPending[0])).toBe(DEFAULT_MARKER);
|
||||
expect(ctx.contextSize.getStatus().contextTokensWithPending).toBe(
|
||||
ctx.contextSize.getStatus().contextTokens + estimateTokensForMessages(rawPending),
|
||||
expect(ctx.contextSize.get().size).toBe(
|
||||
ctx.contextSize.get().measured + estimateTokensForMessages(rawPending),
|
||||
);
|
||||
expect(ctx.contextSize.getStatus().contextTokensWithPending).toBeGreaterThan(
|
||||
ctx.contextSize.getStatus().contextTokens + estimateTokensForMessages(projectedPending),
|
||||
expect(ctx.contextSize.get().size).toBeGreaterThan(
|
||||
ctx.contextSize.get().measured + estimateTokensForMessages(projectedPending),
|
||||
);
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -273,9 +273,10 @@ describe('AgentRecords persistence metadata', () => {
|
|||
]);
|
||||
|
||||
expect(context.get()).toHaveLength(1);
|
||||
expect(contextSize.getStatus()).toEqual({
|
||||
contextTokens: 42,
|
||||
contextTokensWithPending: 42,
|
||||
expect(contextSize.get()).toEqual({
|
||||
size: 42,
|
||||
measured: 42,
|
||||
estimated: 0,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -189,7 +189,7 @@ export const actionMap: Record<ScopeKind, Record<string, ActionTarget>> = {
|
|||
|
||||
'usage:status': { service: IAgentUsageService, method: 'status', readonly: true },
|
||||
|
||||
'context:status': { service: IAgentContextSizeService, method: 'getStatus', readonly: true },
|
||||
'context:status': { service: IAgentContextSizeService, method: 'get', readonly: true },
|
||||
|
||||
'swarm:isActive': { service: IAgentSwarmService, method: 'isActive', readonly: true },
|
||||
'swarm:enter': { service: IAgentSwarmService, method: 'enter' },
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue