Files
deepseek-harness/packages/llm/openrouter-usage/tests/projection.spec.ts
Coder 99b63abde6 feat: show per-step OpenRouter cost on the Trajectory tab
Extend the openRouterCost projection with a per-step priced-cost map
keyed by turn:step so surfaces can render spend without re-pricing, and
have the Trajectory view read it through the framework useProjection seat:
each priced assistant step shows its USD cost on the request boundary chip
and in the inspector summary, while unpriced steps stay blank. Also fix a
pre-existing exactOptionalPropertyTypes error in extractCacheRates that the
host typecheck surfaced.
2026-08-20 22:14:09 +07:00

256 lines
11 KiB
TypeScript

import { describe, expect, it } from 'vitest'
import { Context } from '@deepseek-ai/cordis'
import { createMessage } 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 type { ModelPricing } from '@deepseek-ai/dsh-openrouter-usage/client'
import { createOpenRouterCostProjection } from '../src/projection.ts'
const PRICING = new Map<string, ModelPricing>([
['deepseek/deepseek-chat', {
promptUsd: 1.4e-6,
completionUsd: 2.8e-6,
requestUsd: 0,
cacheReadUsd: 1.4e-7,
cacheWriteUsd: 1.4e-6,
}],
// A free model: priced, but at zero USD per bucket (still a priced step).
['deepseek/deepseek-chat:free', { promptUsd: 0, completionUsd: 0, requestUsd: 0 }],
// A model with a flat per-request fee and no disclosed cache rate.
['expensive/request-fee', { promptUsd: 1e-5, completionUsd: 1e-5, requestUsd: 0.5 }],
])
async function harness(): Promise<{ ctx: Context; session: Session }> {
const ctx = new Context()
await ctx.plugin(SessionStore)
await ctx.plugin(SessionProjectionRegistry)
const session = ctx.sessions.create()
ctx.sessionProjections.register(createOpenRouterCostProjection(model => PRICING.get(model)))
return { ctx, session }
}
function startStep(session: Session, turn: number, step: number): void {
session.append('step/start', { turn, step })
}
/** Append a usage chunk and return its seq. */
function usageChunk(session: Session, usage: TokenUsage, turn: number, step: number): number {
return session.append('assistant/chunk', { turn, step, chunk: { type: 'usage', usage } }).seq
}
/** Append the assistant message for a step and close it. */
function finalUsage(
session: Session,
usage: TokenUsage,
turn: number,
step: number,
sourceSeqs: number[],
model = 'deepseek/deepseek-chat',
): void {
session.append('assistant/message', {
turn,
step,
message: createMessage({
role: 'assistant',
content: [],
source: { kind: 'model', provider: 'openrouter', model },
}),
usage,
}, { surfaceOp: 'append', sourceEventSeqs: sourceSeqs })
session.append('step/end', { turn, step })
}
function recordContext(session: Session): void {
session.append('request/context', { provider: 'openrouter', model: 'deepseek/deepseek-chat' })
}
const projected = (ctx: Context, session: Session) => {
const value = ctx.sessionProjections.snapshot(session).values.openRouterCost
if (value === undefined) throw new Error('openRouterCost projection is not registered')
return value
}
describe('openRouterCost session projection', () => {
it('serves an all-zero view on an empty log', async () => {
const { ctx, session } = await harness()
expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 0, steps: {}, currency: 'USD' })
})
it('prices input/output/cache buckets at the model rates', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
const source = usageChunk(session, { inputTokens: 1_000, outputTokens: 500, cacheReadTokens: 400, cacheWriteTokens: 100 }, 1, 1)
finalUsage(session, { inputTokens: 1_000, outputTokens: 500, cacheReadTokens: 400, cacheWriteTokens: 100 }, 1, 1, [source])
const input = 1_000 * 1.4e-6
const cacheRead = 400 * 1.4e-7
const cacheWrite = 100 * 1.4e-6
const output = 500 * 2.8e-6
expect(projected(ctx, session).totalUsd).toBeCloseTo(input + cacheRead + cacheWrite + output, 12)
expect(projected(ctx, session).pricedSteps).toBe(1)
})
it('does not double-count a usage chunk and the identical final usage', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
const source = usageChunk(session, { inputTokens: 10, outputTokens: 4 }, 1, 1)
finalUsage(session, { inputTokens: 10, outputTokens: 4 }, 1, 1, [source])
expect(projected(ctx, session)).toEqual({
totalUsd: 10 * 1.4e-6 + 4 * 2.8e-6,
pricedSteps: 1,
unknownModelSteps: 0,
steps: { '1:1': 10 * 1.4e-6 + 4 * 2.8e-6 },
currency: 'USD',
})
})
it('replaces an earlier same-step chunk sample with the final usage', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
const source = usageChunk(session, { inputTokens: 10, outputTokens: 2 }, 1, 1)
finalUsage(session, { inputTokens: 14, outputTokens: 5 }, 1, 1, [source])
expect(projected(ctx, session)).toEqual({
totalUsd: 14 * 1.4e-6 + 5 * 2.8e-6,
pricedSteps: 1,
unknownModelSteps: 0,
steps: { '1:1': 14 * 1.4e-6 + 5 * 2.8e-6 },
currency: 'USD',
})
})
it('retains a usage chunk when no final assistant message lands (failed step)', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 9, outputTokens: 1 }, 1, 1)
session.append('step/end', { turn: 1, step: 1 })
expect(projected(ctx, session).totalUsd).toBeCloseTo(9 * 1.4e-6 + 1 * 2.8e-6, 12)
expect(projected(ctx, session).pricedSteps).toBe(1)
expect(projected(ctx, session).steps).toEqual({ '1:1': 9 * 1.4e-6 + 1 * 2.8e-6 })
})
it('counts an unknown-priced OpenRouter model as an unpriced step', async () => {
const { ctx, session } = await harness()
session.append('request/context', { provider: 'openrouter', model: 'brand-new/model' })
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 100, outputTokens: 1 }, 1, 1)
session.append('step/end', { turn: 1, step: 1 })
expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 1, steps: {}, currency: 'USD' })
})
it('ignores a step on a non-openrouter provider entirely', async () => {
const { ctx, session } = await harness()
session.append('request/context', { provider: 'deepseek', model: 'deepseek-chat' })
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 100, outputTokens: 1 }, 1, 1)
session.append('step/end', { turn: 1, step: 1 })
expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 0, steps: {}, currency: 'USD' })
})
it('prefers the assistant-message source over the last request/context record', async () => {
const { ctx, session } = await harness()
session.append('request/context', { provider: 'openrouter', model: 'brand-new/model' })
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 100, outputTokens: 10 }, 1, 1)
// The message attributes the step to a priced model, overriding the last
// request/context attribution used for the early chunk.
const source = session.events.length - 1
finalUsage(session, { inputTokens: 100, outputTokens: 10 }, 1, 1, [source])
expect(projected(ctx, session).pricedSteps).toBe(1)
expect(projected(ctx, session).unknownModelSteps).toBe(0)
})
it('charges the flat per-request fee once per step', async () => {
const { ctx, session } = await harness()
session.append('request/context', { provider: 'openrouter', model: 'expensive/request-fee' })
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 100, outputTokens: 1 }, 1, 1)
session.append('step/end', { turn: 1, step: 1 })
expect(projected(ctx, session).totalUsd).toBeCloseTo(100 * 1e-5 + 1 * 1e-5 + 0.5, 12)
expect(projected(ctx, session).pricedSteps).toBe(1)
})
it('prices a zero-rate free model as a priced (not unknown) step', async () => {
const { ctx, session } = await harness()
session.append('request/context', { provider: 'openrouter', model: 'deepseek/deepseek-chat:free' })
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 100, outputTokens: 10 }, 1, 1)
session.append('step/end', { turn: 1, step: 1 })
expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 1, unknownModelSteps: 0, steps: { '1:1': 0 }, currency: 'USD' })
})
it('maps each step to its priced cost under turn:step keys', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 1_000, outputTokens: 500 }, 1, 1)
startStep(session, 1, 2)
usageChunk(session, { inputTokens: 200, outputTokens: 100 }, 1, 2)
session.append('step/end', { turn: 1, step: 2 })
expect(projected(ctx, session)).toEqual({
totalUsd: (1_000 * 1.4e-6 + 500 * 2.8e-6) + (200 * 1.4e-6 + 100 * 2.8e-6),
pricedSteps: 2,
unknownModelSteps: 0,
steps: {
'1:1': 1_000 * 1.4e-6 + 500 * 2.8e-6,
'1:2': 200 * 1.4e-6 + 100 * 2.8e-6,
},
currency: 'USD',
})
})
it('drops a step entry when the final sample re-attributes it to an unknown-priced model', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
const source = usageChunk(session, { inputTokens: 100, outputTokens: 10 }, 1, 1)
session.append('assistant/message', {
turn: 1,
step: 1,
message: createMessage({
role: 'assistant',
content: [],
source: { kind: 'model', provider: 'openrouter', model: 'brand-new/model' },
}),
usage: { inputTokens: 100, outputTokens: 10 },
}, { surfaceOp: 'append', sourceEventSeqs: [source] })
session.append('step/end', { turn: 1, step: 1 })
// The earlier priced chunk's entry is removed with the re-attribution: the
// step is no longer priced, so it has no entry.
expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 1, steps: {}, currency: 'USD' })
})
it('pushes no change for unrelated events', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 10, outputTokens: 1 }, 1, 1)
const changed: string[] = []
ctx.sessionProjections.onChanged((_session, key) => { changed.push(key) })
session.append('todo/write', { todos: [] })
expect(changed).not.toContain('openRouterCost')
})
it('restores from a JSON checkpoint', async () => {
const { ctx, session } = await harness()
recordContext(session)
startStep(session, 1, 1)
usageChunk(session, { inputTokens: 8, outputTokens: 2 }, 1, 1)
session.append('step/end', { turn: 1, step: 1 })
const checkpoint = JSON.parse(JSON.stringify(
ctx.sessionProjections.checkpoint(session),
)) as ReturnType<typeof ctx.sessionProjections.checkpoint>
expect(ctx.sessionProjections.viewCheckpoint(checkpoint).openRouterCost).toEqual({
totalUsd: 8 * 1.4e-6 + 2 * 2.8e-6,
pricedSteps: 1,
unknownModelSteps: 0,
steps: { '1:1': 8 * 1.4e-6 + 2 * 2.8e-6 },
currency: 'USD',
})
})
})