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.
256 lines
11 KiB
TypeScript
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',
|
|
})
|
|
})
|
|
})
|