|
|
@@ -0,0 +1,222 @@
|
|
|
+import { describe, expect, it } from 'vitest'
|
|
|
+import { Context } from 'cordis'
|
|
|
+import { createMessage, createUserMessage } from '@deepseek-ai/dsh-llm'
|
|
|
+import type { TokenUsage } from '@deepseek-ai/dsh-llm'
|
|
|
+import SessionStore from '@deepseek-ai/dsh-session'
|
|
|
+import type { Session } from '@deepseek-ai/dsh-session'
|
|
|
+import SessionProjectionRegistry from '@deepseek-ai/dsh-session-projection'
|
|
|
+import TokenMeterService from '@deepseek-ai/dsh-token-meter'
|
|
|
+import type { TokenUsageProjection } from '@deepseek-ai/dsh-token-meter/client'
|
|
|
+
|
|
|
+const ZERO: TokenUsageProjection = {
|
|
|
+ uncachedInputTokens: 0,
|
|
|
+ outputTokens: 0,
|
|
|
+ cacheReadTokens: 0,
|
|
|
+ cacheWriteTokens: 0,
|
|
|
+}
|
|
|
+
|
|
|
+async function harness(): Promise<{
|
|
|
+ ctx: Context
|
|
|
+ session: Session
|
|
|
+ meterFiber: Awaited<ReturnType<Context['plugin']>>
|
|
|
+}> {
|
|
|
+ const ctx = new Context()
|
|
|
+ await ctx.plugin(SessionStore)
|
|
|
+ await ctx.plugin(SessionProjectionRegistry)
|
|
|
+ const meterFiber = await ctx.plugin(TokenMeterService)
|
|
|
+ return { ctx, session: ctx.sessions.create(), meterFiber }
|
|
|
+}
|
|
|
+
|
|
|
+function startStep(session: Session, turn: number, step: number): void {
|
|
|
+ session.append('step/start', { turn, step })
|
|
|
+}
|
|
|
+
|
|
|
+function usageChunk(
|
|
|
+ session: Session,
|
|
|
+ usage: TokenUsage,
|
|
|
+ turn: number,
|
|
|
+ step: number,
|
|
|
+): number {
|
|
|
+ return session.append('assistant/chunk', {
|
|
|
+ turn,
|
|
|
+ step,
|
|
|
+ chunk: { type: 'usage', usage },
|
|
|
+ }).seq
|
|
|
+}
|
|
|
+
|
|
|
+function finalUsage(
|
|
|
+ session: Session,
|
|
|
+ usage: TokenUsage,
|
|
|
+ turn: number,
|
|
|
+ step: number,
|
|
|
+ sourceSeqs: number[],
|
|
|
+): void {
|
|
|
+ session.append('assistant/message', {
|
|
|
+ turn,
|
|
|
+ step,
|
|
|
+ message: createMessage({
|
|
|
+ role: 'assistant',
|
|
|
+ content: [],
|
|
|
+ source: { kind: 'model', provider: 'mock', model: 'mock' },
|
|
|
+ }),
|
|
|
+ usage,
|
|
|
+ }, { surfaceOp: 'append', sourceEventSeqs: sourceSeqs })
|
|
|
+ session.append('step/end', { turn, step })
|
|
|
+}
|
|
|
+
|
|
|
+const projected = (ctx: Context, session: Session): TokenUsageProjection => {
|
|
|
+ const value = ctx.sessionProjections.snapshot(session).values.tokenUsage
|
|
|
+ if (value === undefined) throw new Error('tokenUsage projection is not registered')
|
|
|
+ return value
|
|
|
+}
|
|
|
+
|
|
|
+describe('tokenUsage session projection', () => {
|
|
|
+ it('serves zero buckets for an empty log', async () => {
|
|
|
+ const { ctx, session } = await harness()
|
|
|
+ expect(projected(ctx, session)).toEqual(ZERO)
|
|
|
+ })
|
|
|
+
|
|
|
+ it('does not count a usage chunk and identical final usage twice', async () => {
|
|
|
+ const { ctx, session } = await harness()
|
|
|
+ const changes: unknown[] = []
|
|
|
+ ctx.sessionProjections.onChanged((_session, key, value) => {
|
|
|
+ if (key === 'tokenUsage') changes.push(value)
|
|
|
+ })
|
|
|
+ const usage = {
|
|
|
+ inputTokens: 10,
|
|
|
+ outputTokens: 4,
|
|
|
+ cacheReadTokens: 7,
|
|
|
+ cacheWriteTokens: 2,
|
|
|
+ reasoningTokens: 3,
|
|
|
+ }
|
|
|
+ startStep(session, 1, 1)
|
|
|
+ const source = usageChunk(session, usage, 1, 1)
|
|
|
+ finalUsage(session, usage, 1, 1, [source])
|
|
|
+
|
|
|
+ expect(projected(ctx, session)).toEqual({
|
|
|
+ uncachedInputTokens: 10,
|
|
|
+ outputTokens: 4,
|
|
|
+ cacheReadTokens: 7,
|
|
|
+ cacheWriteTokens: 2,
|
|
|
+ })
|
|
|
+ expect(changes).toHaveLength(1)
|
|
|
+ })
|
|
|
+
|
|
|
+ it('replaces an earlier same-step chunk sample with the final usage', async () => {
|
|
|
+ const { ctx, session } = await harness()
|
|
|
+ startStep(session, 1, 1)
|
|
|
+ const source = usageChunk(session, {
|
|
|
+ inputTokens: 10,
|
|
|
+ outputTokens: 2,
|
|
|
+ cacheReadTokens: 3,
|
|
|
+ }, 1, 1)
|
|
|
+ finalUsage(session, {
|
|
|
+ inputTokens: 14,
|
|
|
+ outputTokens: 5,
|
|
|
+ cacheReadTokens: 8,
|
|
|
+ cacheWriteTokens: 1,
|
|
|
+ }, 1, 1, [source])
|
|
|
+
|
|
|
+ expect(projected(ctx, session)).toEqual({
|
|
|
+ uncachedInputTokens: 14,
|
|
|
+ outputTokens: 5,
|
|
|
+ cacheReadTokens: 8,
|
|
|
+ cacheWriteTokens: 1,
|
|
|
+ })
|
|
|
+ })
|
|
|
+
|
|
|
+ it('accumulates disjoint buckets across steps without adding reasoning twice', async () => {
|
|
|
+ const { ctx, session } = await harness()
|
|
|
+ startStep(session, 1, 1)
|
|
|
+ const first = usageChunk(session, {
|
|
|
+ inputTokens: 10,
|
|
|
+ outputTokens: 6,
|
|
|
+ reasoningTokens: 5,
|
|
|
+ cacheReadTokens: 2,
|
|
|
+ }, 1, 1)
|
|
|
+ finalUsage(session, {
|
|
|
+ inputTokens: 10,
|
|
|
+ outputTokens: 6,
|
|
|
+ reasoningTokens: 5,
|
|
|
+ cacheReadTokens: 2,
|
|
|
+ }, 1, 1, [first])
|
|
|
+ startStep(session, 1, 2)
|
|
|
+ const second = usageChunk(session, {
|
|
|
+ inputTokens: 20,
|
|
|
+ outputTokens: 9,
|
|
|
+ reasoningTokens: 7,
|
|
|
+ cacheWriteTokens: 4,
|
|
|
+ }, 1, 2)
|
|
|
+ finalUsage(session, {
|
|
|
+ inputTokens: 20,
|
|
|
+ outputTokens: 9,
|
|
|
+ reasoningTokens: 7,
|
|
|
+ cacheWriteTokens: 4,
|
|
|
+ }, 1, 2, [second])
|
|
|
+
|
|
|
+ expect(projected(ctx, session)).toEqual({
|
|
|
+ uncachedInputTokens: 30,
|
|
|
+ outputTokens: 15,
|
|
|
+ cacheReadTokens: 2,
|
|
|
+ cacheWriteTokens: 4,
|
|
|
+ })
|
|
|
+ })
|
|
|
+
|
|
|
+ it('retains a usage chunk when the request produces no final assistant message', async () => {
|
|
|
+ const { ctx, session } = await harness()
|
|
|
+ startStep(session, 1, 1)
|
|
|
+ usageChunk(session, { inputTokens: 9, outputTokens: 1 }, 1, 1)
|
|
|
+ session.append('step/end', { turn: 1, step: 1 })
|
|
|
+ expect(projected(ctx, session)).toEqual({
|
|
|
+ uncachedInputTokens: 9,
|
|
|
+ outputTokens: 1,
|
|
|
+ cacheReadTokens: 0,
|
|
|
+ cacheWriteTokens: 0,
|
|
|
+ })
|
|
|
+ })
|
|
|
+
|
|
|
+ it('does not erase historical billing when the visible surface is replaced', async () => {
|
|
|
+ const { ctx, session } = await harness()
|
|
|
+ startStep(session, 1, 1)
|
|
|
+ const source = usageChunk(session, { inputTokens: 12, outputTokens: 3 }, 1, 1)
|
|
|
+ finalUsage(session, { inputTokens: 12, outputTokens: 3 }, 1, 1, [source])
|
|
|
+ const before = session.append('user/message', createUserMessage({
|
|
|
+ content: [{ type: 'text', text: 'before compaction' }],
|
|
|
+ source: { kind: 'user' },
|
|
|
+ }), { surfaceOp: 'append' })
|
|
|
+ session.append('user/message', createUserMessage({
|
|
|
+ content: [{ type: 'text', text: 'compacted' }],
|
|
|
+ source: { kind: 'plugin', plugin: 'test' },
|
|
|
+ }), {
|
|
|
+ surfaceOp: { op: 'replace', start: before.seq, end: before.seq },
|
|
|
+ sourceEventSeqs: [before.seq],
|
|
|
+ })
|
|
|
+
|
|
|
+ expect(projected(ctx, session)).toEqual({
|
|
|
+ uncachedInputTokens: 12,
|
|
|
+ outputTokens: 3,
|
|
|
+ cacheReadTokens: 0,
|
|
|
+ cacheWriteTokens: 0,
|
|
|
+ })
|
|
|
+ })
|
|
|
+
|
|
|
+ it('unregisters with the token-meter fiber and restores from a JSON checkpoint', async () => {
|
|
|
+ const { ctx, session, meterFiber } = await harness()
|
|
|
+ startStep(session, 1, 1)
|
|
|
+ usageChunk(session, { inputTokens: 8, outputTokens: 2, cacheReadTokens: 5 }, 1, 1)
|
|
|
+ const checkpoint = JSON.parse(JSON.stringify(
|
|
|
+ ctx.sessionProjections.checkpoint(session),
|
|
|
+ )) as ReturnType<typeof ctx.sessionProjections.checkpoint>
|
|
|
+
|
|
|
+ await meterFiber.dispose()
|
|
|
+ expect(ctx.sessionProjections.snapshot(session).values).not.toHaveProperty('tokenUsage')
|
|
|
+
|
|
|
+ await ctx.plugin(TokenMeterService)
|
|
|
+ expect(ctx.sessionProjections.viewCheckpoint(checkpoint).tokenUsage).toEqual({
|
|
|
+ uncachedInputTokens: 8,
|
|
|
+ outputTokens: 2,
|
|
|
+ cacheReadTokens: 5,
|
|
|
+ cacheWriteTokens: 0,
|
|
|
+ })
|
|
|
+ })
|
|
|
+})
|