113 lines
5.9 KiB
TypeScript
113 lines
5.9 KiB
TypeScript
import assert from 'node:assert/strict'
|
|
import { createContextUsageService, contextSnapshotMatches, type ContextUsageDependencies } from '../src/main/context-usage-service'
|
|
import { contextBreakdownMethodLabel } from '../src/shared/context-usage'
|
|
import type { ChatMessage, ContextUsageBreakdown } from '../src/shared/types'
|
|
|
|
const sample = {
|
|
usage: { inputTokens: 1200, outputTokens: 100, totalTokens: 1300 },
|
|
contextWindow: 10_000,
|
|
source: 'upstream' as const,
|
|
sampledAt: 100,
|
|
requestId: 'request-a',
|
|
modelId: 'model-a'
|
|
}
|
|
const assistant = (overrides: Partial<ChatMessage> = {}): ChatMessage => ({
|
|
id: 'assistant-a',
|
|
role: 'assistant',
|
|
content: 'answer',
|
|
modelId: 'model-a',
|
|
requestId: 'request-a',
|
|
contextUsage: sample,
|
|
createdAt: 1,
|
|
status: 'done',
|
|
...overrides
|
|
})
|
|
const breakdown: ContextUsageBreakdown = {
|
|
requestId: 'request-a',
|
|
sequence: 3,
|
|
sampledAt: 101,
|
|
modelId: 'model-a',
|
|
contextWindow: 10_000,
|
|
method: 'native-summary',
|
|
scope: 'provider-request',
|
|
counts: { system: 100, tools: 100, conversation: 900, mcp: 50, skills: 50, other: 0 },
|
|
inputTokens: 1200,
|
|
partial: false
|
|
}
|
|
|
|
function service(overrides: Partial<ContextUsageDependencies> = {}) {
|
|
const messages = [assistant()]
|
|
const dependencies: ContextUsageDependencies = {
|
|
messages: async () => messages,
|
|
modelId: async () => 'model-a',
|
|
running: () => false,
|
|
summary: async () => undefined,
|
|
persistIfCurrent: async (_id, guard, summary) => {
|
|
const current = await dependencies.messages(_id)
|
|
if(contextSnapshotMatches(current, guard, summary))current.findLast(message=>message.role==='assistant')!.contextBreakdown=summary
|
|
},
|
|
...overrides
|
|
}
|
|
return createContextUsageService(dependencies)
|
|
}
|
|
|
|
assert.deepEqual(await service({ messages: async () => [] }).get('empty'),
|
|
{ status: 'unavailable', reason: '发送消息后可查看用量' })
|
|
|
|
const latest = await service().get('latest')
|
|
assert.equal(latest.status, 'latest')
|
|
assert.equal(latest.sample?.usage.inputTokens, 1200)
|
|
|
|
const previous = await service({ modelId: async () => 'model-b' }).get('previous')
|
|
assert.equal(previous.status, 'previous')
|
|
assert.equal(previous.reason, '来自切换模型前的最近请求')
|
|
|
|
const running = await service({ modelId: async () => 'model-b', running: () => true }).get('running')
|
|
assert.equal(running.status, 'latest', 'the active run must not be relabelled when the next-send model changes')
|
|
|
|
const summary = await service({ summary: async () => breakdown }).get('summary')
|
|
assert.equal(summary.breakdown?.sequence, 3)
|
|
assert.equal(summary.status, 'latest', 'a matching native summary augments the current request without replacing its upstream total')
|
|
|
|
const mismatched = await service({ summary: async () => ({ ...breakdown, requestId: 'older-request' }) }).get('mismatch')
|
|
assert.equal(mismatched.breakdown, undefined, 'a late summary for another request must not attach to the latest assistant message')
|
|
|
|
const summaryFailure = await service({ summary: async () => { throw new Error('query failed') } }).get('failure')
|
|
assert.equal(summaryFailure.sample?.usage.inputTokens, 1200, 'summary failures retain the real streamed usage sample')
|
|
|
|
console.log('context-usage-regression: latest, previous, active-run, summary identity and failure cases passed')
|
|
|
|
const cachedMessages = [assistant()]
|
|
const originalUsage = JSON.stringify(cachedMessages[0].contextUsage)
|
|
await service({messages:async()=>cachedMessages,summary:async()=>breakdown}).get('save')
|
|
assert.deepEqual(cachedMessages[0].contextBreakdown,breakdown)
|
|
assert.equal(JSON.stringify(cachedMessages[0].contextUsage),originalUsage,'saving categories never overwrites actual upstream usage')
|
|
const coldMessages=JSON.parse(JSON.stringify(cachedMessages)) as ChatMessage[]
|
|
for(const summary of [async()=>undefined,async()=>{throw new Error('kernel shut down')}]) {
|
|
const offline=await service({messages:async()=>coldMessages,summary}).get('offline')
|
|
assert.deepEqual(offline.breakdown,breakdown,'persisted categories survive cold/offline/no-instance queries')
|
|
assert.equal(offline.sample?.contextWindow,10_000)
|
|
}
|
|
const oldModel=await service({messages:async()=>coldMessages,modelId:async()=>'model-b'}).get('switch')
|
|
assert.equal(oldModel.status,'previous')
|
|
assert.equal(oldModel.sample?.contextWindow,10_000,'never combine old sample with the new selected model window')
|
|
assert.equal((await service({summary:async()=>({...breakdown,modelId:'wrong'})}).get('wrong-model')).breakdown,undefined)
|
|
for(const change of ['assistant','rewind','pending-user'] as const) {
|
|
let current=[assistant()]
|
|
const result=await service({messages:async()=>current,summary:async()=>{
|
|
if(change==='assistant')current.push(assistant({id:'new',requestId:'new-request',contextUsage:undefined}))
|
|
if(change==='rewind')current=[{id:'different-user',role:'user',content:'fork',createdAt:0},assistant()]
|
|
if(change==='pending-user')current.push({id:'pending',role:'user',content:'next',createdAt:2})
|
|
return breakdown
|
|
}}).get(change)
|
|
assert.equal(result.breakdown,undefined,`${change} while summary awaits cannot pollute the active branch`)
|
|
}
|
|
const pending=[assistant({contextBreakdown:breakdown}),{id:'pending',role:'user',content:'next',createdAt:2} as ChatMessage]
|
|
assert.equal((await service({messages:async()=>pending,summary:async()=>breakdown}).get('pending')).status,'previous','refresh must not erase pending-user previous status')
|
|
const categoriesOnly=await service({messages:async()=>[assistant({contextUsage:undefined,contextBreakdown:breakdown})]}).get('categories-only')
|
|
assert.equal(categoriesOnly.status,'latest','categories without upstream usage remain available, not fake zero usage')
|
|
assert.equal(categoriesOnly.sample,undefined)
|
|
assert.equal(contextBreakdownMethodLabel('native-summary'),'内核 summary 本地估计')
|
|
assert.equal(contextBreakdownMethodLabel('utf8-estimate'),'UTF-8 字节估计')
|
|
console.log('context-usage persistence: offline, shutdown, identity races, usage preservation, category provenance passed')
|