refactor(token-meter): fold the surface once for the meter and the projection

`contextBreakdown.messageTokens` and `measure().surfaceTokens` answer the same
question in the same heuristic vocabulary, and the panel's composition rows are
only honest while they agree. Each owner carried its own copy of the positional
fold — same pricing, same `{seq, tokens}` node list, same replace-range lookup
and guard, differing only in mutable versus immutable application — so an edit
to either one would have moved the panel away from `measure()` with both sides
still green. The duplication gate caught the shared 62 tokens.

`src/surface-fold.ts` now owns `foldSurfaceTokens`: total, allocation-fresh,
returning the event's price, the next surface, and the signed total delta. The
service assigns that result where it used to prepare a commit closure, which
keeps its validate-before-mutate replay transaction intact — the fold throws
before any state is touched, so a malformed event still fails identically on
every retry. `_prepareSurfaceMutation` and `_estimateSurfaceEvent` go away with
it, and the projection's apply drops to one call.

Covers the identity with a session that appends and then compacts, asserting
the projection figure equals the service surface at each boundary; the test
fails when either side of the fold is perturbed.
This commit is contained in:
Yichen Jiang
2026-08-05 16:29:51 +08:00
parent 46562b2c30
commit e62cbe12e4
12 changed files with 138 additions and 100 deletions

View File

@@ -8,7 +8,7 @@ import { Context, Service } from 'cordis'
import z from 'schemastery'
import { BlockAssembler, deepFreeze } from '@deepseek-ai/dsh-llm'
import type { Message, TokenUsage } from '@deepseek-ai/dsh-llm'
import type { EpochHeader, Session, SessionEvent, SurfaceEvent } from '@deepseek-ai/dsh-session'
import type { EpochHeader, Session, SessionEvent } from '@deepseek-ai/dsh-session'
import { canonicalHeader, headerEquals, isSurfaceEvent } from '@deepseek-ai/dsh-session'
// Type-only: resolves the optional projection registry Context seam.
import type {} from '@deepseek-ai/dsh-session-projection'
@@ -21,6 +21,7 @@ import type {
import { contextBreakdownProjectionDefinition } from './breakdown-projection.ts'
import { contextPressureProjectionDefinition, tokenUsageProjectionDefinition } from './usage-projection.ts'
import { estimateContent, estimateHeader, estimateMessage, ROLE_OVERHEAD } from './estimate.ts'
import { foldSurfaceTokens } from './surface-fold.ts'
export type * from './types.ts'
@@ -39,11 +40,6 @@ interface ReplayState {
anchor: MeasurementAnchor | undefined
}
interface PreparedSurfaceMutation {
readonly tokens: number
commit(state: ReplayState): void
}
/** Sum disjoint provider usage buckets without double-counting reasoning output. */
function usageTokens(usage: TokenUsage): number {
return usage.inputTokens
@@ -219,7 +215,7 @@ export class TokenMeterService extends Service {
}
const surface = isSurfaceEvent(event)
? this._prepareSurfaceMutation(session, state, event)
? foldSurfaceTokens(state.surface, event)
: undefined
if (event.type === 'assistant/message') {
@@ -266,53 +262,13 @@ export class TokenMeterService extends Service {
state.header = nextHeader
state.stepStart = nextStepStart
if (surface !== undefined) surface.commit(state)
if (surface !== undefined) {
state.surface = surface.nodes
state.surfaceTokens += surface.deltaTokens
}
state.anchor = nextAnchor
}
/** Validate one surface operation and return its allocation-light commit. */
private _prepareSurfaceMutation(
session: Session,
state: ReplayState,
event: SurfaceEvent,
): PreparedSurfaceMutation {
const tokens = this._estimateSurfaceEvent(session, event)
const op = event.surfaceOp
if (op === 'append') {
return {
tokens,
commit(target) {
target.surface.push({ seq: event.seq, tokens })
target.surfaceTokens += tokens
},
}
}
const startIdx = state.surface.findIndex(node => node.seq === op.start)
const endIdx = state.surface.findIndex(node => node.seq === op.end)
if (startIdx === -1 || endIdx === -1 || startIdx > endIdx) {
throw new Error(
`token meter: replace at seq ${event.seq} has invalid current range ${op.start}-${op.end}`,
)
}
const removedTokens = state.surface
.slice(startIdx, endIdx + 1)
.reduce((total, node) => total + node.tokens, 0)
return {
tokens,
commit(target) {
target.surface.splice(startIdx, endIdx - startIdx + 1, { seq: event.seq, tokens })
target.surfaceTokens += tokens - removedTokens
},
}
}
/** Price one current surface event exactly as it projects to a request. */
private _estimateSurfaceEvent(session: Session, event: SurfaceEvent): number {
const message = session.deriveEventMessage(event)
return message === null ? 0 : this.estimateMessage(message)
}
/**
* Reassemble provider output from exact chunk provenance for a usage anchor.
* Missing legacy provenance conservatively treats the durable output as the