mengzhx/deepseek-harness
1
0
Fork 0
mirror of https://github.com/deepseek-ai/deepseek-harness.git synced 2026-09-14 04:01:35 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
5b671d7df3df08b3dee8cb681d8977cfd2f61fee
deepseek-harness/packages/llm/token-meter/tests/turn-usage.spec.ts
T

419 lines
15 KiB
TypeScript
Raw Blame History

import { describe, expect, it } from 'vitest'
import type { StreamChunk, TokenUsage } from '@deepseek-ai/dsh-llm'
import type { SessionEvent } from '@deepseek-ai/dsh-session'
import { deriveTurnTokenUsage } from '../src/turn-usage.ts'
function event(seq: number, type: string, data: unknown): SessionEvent {
return { seq, time: seq, type, data } as unknown as SessionEvent
}
type UsageOverrides = { [Key in keyof TokenUsage]?: TokenUsage[Key] | undefined }
function usage(overrides: UsageOverrides = {}): TokenUsage {
const value = {
inputTokens: 100,
outputTokens: 20,
totalTokens: 170,
cacheReadTokens: 50,
...overrides,
}
return Object.fromEntries(Object.entries(value).filter(([, entry]) => entry !== undefined)) as unknown as TokenUsage
}
function message(
seq: number,
tokenUsage?: TokenUsage,
provider = 'deepseek',
model = 'deepseek-chat',
step = 1,
streamTokenUsage = tokenUsage,
) {
return event(seq, 'assistant/message', {
turn: 1,
step,
stream: [
{ type: 'chunk', time: seq, chunk: { type: 'block-start', index: 0, blockType: 'text' } },
...(streamTokenUsage === undefined
? []
: [{ type: 'chunk' as const, time: seq, chunk: { type: 'usage' as const, usage: streamTokenUsage } }]),
],
message: {
id: `message-${seq}`,
role: 'assistant',
content: [{ type: 'text', text: 'done' }],
source: { kind: 'model', provider, model },
},
...tokenUsage === undefined ? {} : { usage: tokenUsage },
})
}
function attempt(seq: number, chunks: readonly StreamChunk[], step = 1): SessionEvent {
return event(seq, 'assistant/attempt', {
turn: 1,
step,
stream: chunks.map((chunk, index) => ({ type: 'chunk', time: seq + index, chunk })),
})
}
function completeAttempt(...middle: readonly SessionEvent[]): SessionEvent[] {
return [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
...middle,
event(90, 'step/end', { turn: 1, step: 1 }),
event(91, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]
}
describe('deriveTurnTokenUsage', () => {
it('preserves authoritative totals and explicit optional buckets', () => {
expect(deriveTurnTokenUsage(completeAttempt(message(3, usage({
cacheWriteTokens: 0,
reasoningTokens: 8,
}))))).toEqual({
uncachedInputTokens: 100,
outputTokens: 20,
totalTokens: 170,
cacheReadTokens: 50,
cacheWriteTokens: 0,
reasoningTokens: 8,
routes: [{ provider: 'deepseek', model: 'deepseek-chat' }],
})
})
it('derives an exact total only when both cache buckets are present', () => {
expect(deriveTurnTokenUsage(completeAttempt(message(3, usage({
totalTokens: undefined,
inputTokens: 10,
outputTokens: 4,
cacheReadTokens: 2,
cacheWriteTokens: 1,
}))))?.totalTokens).toBe(17)
expect(deriveTurnTokenUsage(completeAttempt(message(3, usage({
totalTokens: undefined,
cacheWriteTokens: undefined,
}))))).toBeUndefined()
})
it('lets final message usage replace the latest streaming sample', () => {
const result = deriveTurnTokenUsage(completeAttempt(
message(
4,
usage({ inputTokens: 30, outputTokens: 5, totalTokens: 45, cacheReadTokens: 10 }),
'deepseek',
'deepseek-chat',
1,
usage(),
),
))
expect(result).toMatchObject({ uncachedInputTokens: 30, outputTokens: 5, totalTokens: 45 })
})
it('keeps the latest streaming sample when the final message omits usage', () => {
const result = deriveTurnTokenUsage(completeAttempt(
message(4, undefined, 'deepseek', 'deepseek-chat', 1, usage()),
))
expect(result).toMatchObject({ uncachedInputTokens: 100, outputTokens: 20, totalTokens: 170 })
})
it('counts an error-finished attempt once across its retry boundary', () => {
const events = completeAttempt(
attempt(3, [
{ type: 'usage', usage: usage() },
{ type: 'finish', reason: { kind: 'error', failure: { code: 'HTTP', message: 'failed' } } },
]),
event(5, 'llm/retry', { turn: 1, step: 1 }),
event(6, 'llm/retry-started', { turn: 1, step: 1, retry: 1 }),
message(7, usage({ inputTokens: 40, outputTokens: 10, totalTokens: 70, cacheReadTokens: 20 })),
)
expect(deriveTurnTokenUsage(events)).toEqual({
uncachedInputTokens: 140,
outputTokens: 30,
totalTokens: 240,
cacheReadTokens: 70,
})
})
it('does not invent an attempt for a scheduled retry that never started', () => {
const result = deriveTurnTokenUsage(completeAttempt(
attempt(3, [
{ type: 'usage', usage: usage() },
{ type: 'finish', reason: { kind: 'error', failure: { code: 'HTTP', message: 'failed' } } },
]),
event(4, 'llm/retry', { turn: 1, step: 1 }),
))
expect(result).toMatchObject({ totalTokens: 170 })
})
it('fails closed for missing lifecycle or missing attempt usage', () => {
expect(deriveTurnTokenUsage([
event(1, 'turn/start', { turn: 1 }),
message(2, usage()),
event(3, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
])).toBeUndefined()
expect(deriveTurnTokenUsage(completeAttempt(message(3)))).toBeUndefined()
})
it.each([
['negative', usage({ inputTokens: -1 })],
['fractional', usage({ outputTokens: 1.5 })],
['unsafe', usage({ totalTokens: Number.MAX_SAFE_INTEGER + 1 })],
['invalid cache read', usage({ cacheReadTokens: -1 })],
['invalid cache write', usage({ cacheWriteTokens: 1.5 })],
['negative exact prompt', usage({ outputTokens: 20, totalTokens: 10, cacheReadTokens: undefined })],
['total below known prompt', usage({ totalTokens: 160 })],
['contradictory complete buckets', usage({ totalTokens: 171, cacheWriteTokens: 0 })],
['reasoning exceeds output', usage({ reasoningTokens: 21 })],
['prompt bucket overflow', usage({
inputTokens: Number.MAX_SAFE_INTEGER,
outputTokens: 0,
totalTokens: Number.MAX_SAFE_INTEGER,
cacheReadTokens: 1,
})],
['derived total overflow', usage({
inputTokens: Number.MAX_SAFE_INTEGER,
outputTokens: 1,
totalTokens: undefined,
cacheReadTokens: 0,
cacheWriteTokens: 0,
})],
])('fails closed for %s usage', (_label, invalidUsage) => {
expect(deriveTurnTokenUsage(completeAttempt(message(3, invalidUsage)))).toBeUndefined()
})
it('omits optional aggregates and routes unless every attempt reports them', () => {
const events = [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
message(3, usage({ totalTokens: 175, cacheWriteTokens: 5, reasoningTokens: 2 })),
event(4, 'step/end', { turn: 1, step: 1 }),
event(5, 'step/start', { turn: 1, step: 2 }),
event(6, 'assistant/message', {
turn: 1,
step: 2,
message: {
id: 'message-6', role: 'assistant', content: [],
source: { kind: 'model', provider: '', model: '' },
},
usage: usage({ cacheReadTokens: undefined, cacheWriteTokens: undefined, reasoningTokens: undefined }),
}),
event(7, 'step/end', { turn: 1, step: 2 }),
event(8, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]
expect(deriveTurnTokenUsage(events)).toEqual({ uncachedInputTokens: 200, outputTokens: 40, totalTokens: 345 })
})
it('sums multiple steps and preserves distinct attributed routes', () => {
const events = [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
message(3, usage()),
event(4, 'step/end', { turn: 1, step: 1 }),
event(5, 'step/start', { turn: 1, step: 2 }),
message(6, usage(), 'openai', 'gpt-5', 2),
event(7, 'step/end', { turn: 1, step: 2 }),
event(8, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]
expect(deriveTurnTokenUsage(events)).toEqual({
uncachedInputTokens: 200,
outputTokens: 40,
totalTokens: 340,
cacheReadTokens: 100,
routes: [
{ provider: 'deepseek', model: 'deepseek-chat' },
{ provider: 'openai', model: 'gpt-5' },
],
})
})
it('fails closed when aggregation overflows a safe integer', () => {
const half = Math.floor(Number.MAX_SAFE_INTEGER / 2) + 1
const attempt = usage({ inputTokens: 0, outputTokens: 0, cacheReadTokens: undefined, totalTokens: half })
const events = [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
message(3, attempt),
event(4, 'step/end', { turn: 1, step: 1 }),
event(5, 'step/start', { turn: 1, step: 2 }),
event(6, 'assistant/message', {
turn: 1,
step: 2,
message: {
id: 'message-6', role: 'assistant', content: [],
source: { kind: 'model', provider: 'deepseek', model: 'deepseek-chat' },
},
usage: attempt,
}),
event(7, 'step/end', { turn: 1, step: 2 }),
event(8, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]
expect(deriveTurnTokenUsage(events)).toBeUndefined()
})
it.each([
['uncached input', usage({
inputTokens: Math.floor(Number.MAX_SAFE_INTEGER / 2) + 1,
outputTokens: 0,
cacheReadTokens: undefined,
totalTokens: Math.floor(Number.MAX_SAFE_INTEGER / 2) + 1,
})],
['output', usage({
inputTokens: 0,
outputTokens: Math.floor(Number.MAX_SAFE_INTEGER / 2) + 1,
cacheReadTokens: undefined,
totalTokens: Math.floor(Number.MAX_SAFE_INTEGER / 2) + 1,
})],
])('fails closed when aggregate %s overflows', (_label, attempt) => {
expect(deriveTurnTokenUsage([
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
message(3, attempt),
event(4, 'step/end', { turn: 1, step: 1 }),
event(5, 'step/start', { turn: 1, step: 1 }),
message(6, attempt),
event(7, 'step/end', { turn: 1, step: 1 }),
event(8, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
])).toBeUndefined()
})
it('closes a sampled attempt at step/end', () => {
expect(deriveTurnTokenUsage(completeAttempt(
attempt(3, [
{ type: 'usage', usage: usage() },
{ type: 'finish', reason: { kind: 'stop' } },
]),
event(5, 'tool/call', { turn: 1, step: 1 }),
))).toMatchObject({ totalTokens: 170 })
})
it('accepts an aborted finish after observing usage', () => {
expect(deriveTurnTokenUsage(completeAttempt(
attempt(3, [
{ type: 'usage', usage: usage() },
{ type: 'finish', reason: { kind: 'aborted', failure: { message: 'aborted', code: 'ABORTED' } } },
]),
))).toMatchObject({ totalTokens: 170 })
})
it.each([
['empty turn', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]],
['duplicate turn start', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'turn/start', { turn: 1 }),
]],
['wrong turn end', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'turn/end', { turn: 2, reason: { kind: 'completed' } }),
]],
['turn end during an open attempt', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
event(3, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]],
['duplicate turn end', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
event(3, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
]],
['event after turn end', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'turn/end', { turn: 1, reason: { kind: 'completed' } }),
event(3, 'step/start', { turn: 1, step: 1 }),
]],
['wrong-turn step start', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 2, step: 1 }),
]],
['nested step start', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
event(3, 'step/start', { turn: 1, step: 2 }),
]],
['retry start without a scheduled retry', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'llm/retry-started', { turn: 1, step: 1, retry: 1 }),
]],
['retry start after a final message', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
message(3, usage()),
event(4, 'llm/retry-started', { turn: 1, step: 1, retry: 1 }),
]],
['retry start for the wrong step', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
attempt(3, [
{ type: 'usage', usage: usage() },
{ type: 'finish', reason: { kind: 'error', failure: { code: 'HTTP', message: 'failed' } } },
]),
event(4, 'llm/retry', { turn: 1, step: 1 }),
event(5, 'llm/retry-started', { turn: 1, step: 2, retry: 1 }),
]],
['attempt outside a step', [
event(1, 'turn/start', { turn: 1 }),
attempt(2, [{ type: 'usage', usage: usage() }]),
]],
['attempt for the wrong step', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
attempt(3, [{ type: 'usage', usage: usage() }], 2),
]],
['error finish without usage', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
attempt(3, [{
type: 'finish', reason: { kind: 'error', failure: { code: 'HTTP', message: 'failed' } },
}]),
]],
['retry outside an attempt', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'llm/retry', { turn: 1, step: 1 }),
]],
['retry for the wrong step', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
attempt(3, [
{ type: 'usage', usage: usage() },
{ type: 'finish', reason: { kind: 'error', failure: { code: 'HTTP', message: 'failed' } } },
]),
event(4, 'llm/retry', { turn: 1, step: 2 }),
]],
['retry after a final message', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
message(3, usage()),
event(4, 'llm/retry', { turn: 1, step: 1 }),
]],
['retry before any usage', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
event(3, 'llm/retry', { turn: 1, step: 1 }),
]],
['step end outside an attempt', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/end', { turn: 1, step: 1 }),
]],
['step end for the wrong step', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
event(3, 'step/end', { turn: 1, step: 2 }),
]],
['step end before any usage', [
event(1, 'turn/start', { turn: 1 }),
event(2, 'step/start', { turn: 1, step: 1 }),
event(3, 'step/end', { turn: 1, step: 1 }),
]],
])('fails closed for invalid lifecycle: %s', (_label, events) => {
expect(deriveTurnTokenUsage(events)).toBeUndefined()
})
it('requires the complete turn window', () => {
expect(deriveTurnTokenUsage(completeAttempt(message(3, usage())).slice(1))).toBeUndefined()
expect(deriveTurnTokenUsage(completeAttempt(message(3, usage())).slice(0, -1))).toBeUndefined()
})
})
Reference in New Issue View Git Blame Copy Permalink