diff --git a/packages/agent-core-v2/docs/state-manifest.d.ts b/packages/agent-core-v2/docs/state-manifest.d.ts index bd4c2ea4a..ea4fd6511 100644 --- a/packages/agent-core-v2/docs/state-manifest.d.ts +++ b/packages/agent-core-v2/docs/state-manifest.d.ts @@ -1508,7 +1508,7 @@ export interface AgentStateSnapshot { 'task.notificationDelivery': readonly string[]; 'task.scheduledNotificationKeys': Set; // 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; diff --git a/packages/agent-core-v2/docs/wire-manifest.d.ts b/packages/agent-core-v2/docs/wire-manifest.d.ts index f3c4ced6f..e9616b0f7 100644 --- a/packages/agent-core-v2/docs/wire-manifest.d.ts +++ b/packages/agent-core-v2/docs/wire-manifest.d.ts @@ -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; diff --git a/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts b/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts index e40d7a542..1f1706944 100644 --- a/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts +++ b/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingOps.ts @@ -41,6 +41,15 @@ export class TokenCountingRebased extends Event2> { } export interface TokenCountingRebased extends z.infer {} +const turnRecordedSchema = sizeSchema.extend({ turnId: z.number() }); + +export class TokenCountingTurnRecorded extends Event2> { + static override readonly type = 'token_counting.turn_recorded'; + static override readonly durable = true; + static override readonly schema = turnRecordedSchema; +} +export interface TokenCountingTurnRecorded extends z.infer {} + 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 { diff --git a/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingService.ts b/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingService.ts index 072d2b4ad..8221c2918 100644 --- a/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingService.ts +++ b/packages/agent-core-v2/src/agent/tokenCounting/tokenCountingService.ts @@ -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 { diff --git a/packages/agent-core-v2/src/app/sessionLegacy/sessionLegacyService.ts b/packages/agent-core-v2/src/app/sessionLegacy/sessionLegacyService.ts index e2b0065f3..bb526b567 100644 --- a/packages/agent-core-v2/src/app/sessionLegacy/sessionLegacyService.ts +++ b/packages/agent-core-v2/src/app/sessionLegacy/sessionLegacyService.ts @@ -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, }; } diff --git a/packages/agent-core-v2/src/app/sessionLegacy/sessionProtocol.ts b/packages/agent-core-v2/src/app/sessionLegacy/sessionProtocol.ts index 46d72c756..17972bddf 100644 --- a/packages/agent-core-v2/src/app/sessionLegacy/sessionProtocol.ts +++ b/packages/agent-core-v2/src/app/sessionLegacy/sessionProtocol.ts @@ -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; diff --git a/packages/agent-core-v2/test/agent/tokenCounting/tokenCounting.test.ts b/packages/agent-core-v2/test/agent/tokenCounting/tokenCounting.test.ts index 4a724eecb..3de7ff33c 100644 --- a/packages/agent-core-v2/test/agent/tokenCounting/tokenCounting.test.ts +++ b/packages/agent-core-v2/test/agent/tokenCounting/tokenCounting.test.ts @@ -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); + }); }); diff --git a/packages/agent-core-v2/test/app/config/config.test.ts b/packages/agent-core-v2/test/app/config/config.test.ts index 7dabf6d2f..513421ce3 100644 --- a/packages/agent-core-v2/test/app/config/config.test.ts +++ b/packages/agent-core-v2/test/app/config/config.test.ts @@ -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": "