237 lines
8.3 KiB
TypeScript
237 lines
8.3 KiB
TypeScript
import { describe, it, expect, beforeEach, afterEach } from 'vitest'
|
|
import { mkdtemp, rm } from 'fs/promises'
|
|
import { tmpdir } from 'os'
|
|
import { join } from 'path'
|
|
import { createRequire } from 'node:module'
|
|
import zlib from 'zlib'
|
|
|
|
import { createZedProvider } from '../../src/providers/zed.js'
|
|
import { isSqliteAvailable } from '../../src/sqlite.js'
|
|
import type { ParsedProviderCall } from '../../src/providers/types.js'
|
|
|
|
const requireForTest = createRequire(import.meta.url)
|
|
|
|
const zstd = (zlib as { zstdCompressSync?: (buf: Buffer) => Buffer }).zstdCompressSync
|
|
|
|
const skipReason = !isSqliteAvailable()
|
|
? 'node:sqlite not available — needs Node 22+; skipping'
|
|
: !zstd
|
|
? 'zlib zstd not available — needs Node 22.15+; skipping'
|
|
: null
|
|
|
|
let tmpDir: string
|
|
|
|
beforeEach(async () => {
|
|
tmpDir = await mkdtemp(join(tmpdir(), 'zed-test-'))
|
|
})
|
|
|
|
afterEach(async () => {
|
|
await rm(tmpDir, { recursive: true, force: true })
|
|
})
|
|
|
|
function buildDb(fn: (db: {
|
|
exec(sql: string): void
|
|
prepare(sql: string): { run(...params: unknown[]): void }
|
|
close(): void
|
|
}) => void): string {
|
|
const dbPath = join(tmpDir, 'threads.db')
|
|
const { DatabaseSync: Database } = requireForTest('node:sqlite')
|
|
const db = new Database(dbPath)
|
|
db.exec(`CREATE TABLE threads (
|
|
id TEXT PRIMARY KEY,
|
|
summary TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
data_type TEXT NOT NULL,
|
|
data BLOB NOT NULL,
|
|
parent_id TEXT, folder_paths TEXT, folder_paths_order TEXT, created_at TEXT
|
|
)`)
|
|
fn(db)
|
|
db.close()
|
|
return dbPath
|
|
}
|
|
|
|
function insertThread(db: {
|
|
prepare(sql: string): { run(...params: unknown[]): void }
|
|
}, opts: {
|
|
id: string
|
|
summary?: string
|
|
updatedAt?: string
|
|
dataType?: string
|
|
thread?: unknown
|
|
rawData?: Buffer
|
|
}): void {
|
|
const data = opts.rawData ?? zstd!(Buffer.from(JSON.stringify(opts.thread ?? {})))
|
|
db.prepare('INSERT INTO threads (id, summary, updated_at, data_type, data) VALUES (?, ?, ?, ?, ?)').run(
|
|
opts.id,
|
|
opts.summary ?? 'a thread',
|
|
opts.updatedAt ?? '2026-06-20T10:00:00Z',
|
|
opts.dataType ?? 'zstd',
|
|
data,
|
|
)
|
|
}
|
|
|
|
async function collectCalls(dbPath: string, seenKeys = new Set<string>()): Promise<ParsedProviderCall[]> {
|
|
const provider = createZedProvider(dbPath)
|
|
const sources = await provider.discoverSessions()
|
|
const calls: ParsedProviderCall[] = []
|
|
for (const source of sources) {
|
|
for await (const call of provider.createSessionParser(source, seenKeys).parse()) {
|
|
calls.push(call)
|
|
}
|
|
}
|
|
return calls
|
|
}
|
|
|
|
describe.skipIf(skipReason !== null)('zed provider (#480)', () => {
|
|
it('emits one call per request with exact Anthropic-shaped token fields', async () => {
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, {
|
|
id: 'thread-1',
|
|
summary: 'refactor the parser',
|
|
updatedAt: '2026-06-21T09:30:00Z',
|
|
thread: {
|
|
model: { provider: 'anthropic', model: 'claude-opus-4-8' },
|
|
request_token_usage: {
|
|
'req-1': { input_tokens: 1200, output_tokens: 300, cache_creation_input_tokens: 5000, cache_read_input_tokens: 90000 },
|
|
'req-2': { input_tokens: 800, output_tokens: 150, cache_creation_input_tokens: 0, cache_read_input_tokens: 95000 },
|
|
},
|
|
cumulative_token_usage: { input_tokens: 2000, output_tokens: 450 },
|
|
},
|
|
})
|
|
})
|
|
|
|
const calls = await collectCalls(dbPath)
|
|
expect(calls.length).toBe(2)
|
|
const first = calls.find(c => c.deduplicationKey === 'zed:thread-1:req-1')
|
|
expect(first).toBeDefined()
|
|
expect(first!.inputTokens).toBe(1200)
|
|
expect(first!.outputTokens).toBe(300)
|
|
expect(first!.cacheCreationInputTokens).toBe(5000)
|
|
expect(first!.cacheReadInputTokens).toBe(90000)
|
|
expect(first!.model).toBe('claude-opus-4-8')
|
|
expect(first!.costUSD).toBeGreaterThan(0)
|
|
expect(first!.sessionId).toBe('thread-1')
|
|
expect(first!.userMessage).toBe('refactor the parser')
|
|
expect(first!.timestamp).toBe('2026-06-21T09:30:00.000Z')
|
|
// Cumulative must not be counted on top of per-request entries.
|
|
expect(calls.reduce((s, c) => s + c.inputTokens, 0)).toBe(2000)
|
|
})
|
|
|
|
it('falls back to cumulative_token_usage when the per-request map is empty', async () => {
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, {
|
|
id: 'thread-2',
|
|
thread: {
|
|
model: { provider: 'anthropic', model: 'claude-sonnet-4-6' },
|
|
request_token_usage: {},
|
|
cumulative_token_usage: { input_tokens: 5400, output_tokens: 900, cache_read_input_tokens: 20000 },
|
|
},
|
|
})
|
|
})
|
|
|
|
const calls = await collectCalls(dbPath)
|
|
expect(calls.length).toBe(1)
|
|
expect(calls[0]!.deduplicationKey).toBe('zed:thread-2:cumulative-remainder')
|
|
expect(calls[0]!.inputTokens).toBe(5400)
|
|
expect(calls[0]!.cacheReadInputTokens).toBe(20000)
|
|
})
|
|
|
|
it('tops threads up to the cumulative counter when the per-request map undercounts', async () => {
|
|
// Mirrors a real thread: the map (keyed by user message) covered only some
|
|
// requests while cumulative carried ~3x the tokens.
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, {
|
|
id: 'thread-real',
|
|
thread: {
|
|
model: { provider: 'zed.dev', model: 'gpt-5.4' },
|
|
request_token_usage: {
|
|
'msg-1': { input_tokens: 9514, output_tokens: 19 },
|
|
'msg-2': { input_tokens: 2757, output_tokens: 310, cache_read_input_tokens: 33408 },
|
|
},
|
|
cumulative_token_usage: { input_tokens: 35464, output_tokens: 1868, cache_read_input_tokens: 140288 },
|
|
},
|
|
})
|
|
})
|
|
|
|
const calls = await collectCalls(dbPath)
|
|
expect(calls.length).toBe(3)
|
|
const remainder = calls.find(c => c.deduplicationKey === 'zed:thread-real:cumulative-remainder')
|
|
expect(remainder).toBeDefined()
|
|
expect(calls.reduce((s, c) => s + c.inputTokens, 0)).toBe(35464)
|
|
expect(calls.reduce((s, c) => s + c.outputTokens, 0)).toBe(1868)
|
|
expect(calls.reduce((s, c) => s + c.cacheReadInputTokens, 0)).toBe(140288)
|
|
})
|
|
|
|
it('skips non-zstd rows and malformed blobs without dropping healthy threads', async () => {
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, { id: 'bad-type', dataType: 'protobuf', rawData: Buffer.from('{}') })
|
|
insertThread(db, { id: 'bad-blob', rawData: Buffer.from('not zstd at all') })
|
|
insertThread(db, {
|
|
id: 'good',
|
|
thread: {
|
|
model: { model: 'claude-opus-4-8' },
|
|
request_token_usage: { 'req-1': { input_tokens: 10, output_tokens: 5 } },
|
|
},
|
|
})
|
|
})
|
|
|
|
const calls = await collectCalls(dbPath)
|
|
expect(calls.length).toBe(1)
|
|
expect(calls[0]!.sessionId).toBe('good')
|
|
})
|
|
|
|
it('reads legacy uncompressed json rows alongside zstd rows', async () => {
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, {
|
|
id: 'legacy',
|
|
dataType: 'json',
|
|
rawData: Buffer.from(JSON.stringify({
|
|
model: { model: 'claude-sonnet-4-6' },
|
|
request_token_usage: { 'req-1': { input_tokens: 40, output_tokens: 8 } },
|
|
})),
|
|
})
|
|
})
|
|
|
|
const calls = await collectCalls(dbPath)
|
|
expect(calls.length).toBe(1)
|
|
expect(calls[0]!.inputTokens).toBe(40)
|
|
expect(calls[0]!.model).toBe('claude-sonnet-4-6')
|
|
})
|
|
|
|
it('dedupes across repeat parses via the shared seenKeys set', async () => {
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, {
|
|
id: 'thread-3',
|
|
thread: {
|
|
model: { model: 'claude-opus-4-8' },
|
|
request_token_usage: { 'req-1': { input_tokens: 100, output_tokens: 50 } },
|
|
},
|
|
})
|
|
})
|
|
|
|
const seen = new Set<string>()
|
|
expect((await collectCalls(dbPath, seen)).length).toBe(1)
|
|
expect((await collectCalls(dbPath, seen)).length).toBe(0)
|
|
})
|
|
|
|
it('skips threads whose usage is entirely zero instead of emitting empty calls', async () => {
|
|
const dbPath = buildDb((db) => {
|
|
insertThread(db, {
|
|
id: 'thread-4',
|
|
thread: {
|
|
model: { model: 'claude-opus-4-8' },
|
|
request_token_usage: { 'req-1': { input_tokens: 0, output_tokens: 0 } },
|
|
cumulative_token_usage: { input_tokens: 0, output_tokens: 0 },
|
|
},
|
|
})
|
|
})
|
|
|
|
expect((await collectCalls(dbPath)).length).toBe(0)
|
|
})
|
|
|
|
it('discovers nothing when the database does not exist', async () => {
|
|
const provider = createZedProvider(join(tmpDir, 'missing.db'))
|
|
expect(await provider.discoverSessions()).toEqual([])
|
|
})
|
|
})
|