diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b1f1dc..bd13088 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -52,6 +52,13 @@ workspace-based project names from session data. Closes #248. ### Fixed (CLI) +- **One-shot rate detection for non-Claude providers.** Gemini and Mistral Vibe + now emit per-assistant-message calls grouped by user turn, so retry detection + sees multi-message `Edit -> Bash -> Edit` flows instead of counting each + message as an independent one-shot turn. Kiro and Goose record per-message + tool ordering via `toolSequence` for the same effect on aggregated sessions. + Vibe prefers `meta.json.stats.session_cost` over price-derived estimates when + available. Session cache bumped to v2. Closes #351. - **Reduced Claude parser OOM risk.** Large Claude JSONL sessions retained full entry objects (text, thinking blocks, tool results) in memory during parsing, causing V8 heap exhaustion on heavy usage months. Entries are now diff --git a/docs/providers/mistral-vibe.md b/docs/providers/mistral-vibe.md index c7005f7..244ebae 100644 --- a/docs/providers/mistral-vibe.md +++ b/docs/providers/mistral-vibe.md @@ -21,7 +21,11 @@ Subagent traces are stored under a parent session's `agents/` folder with the sa ## Caching -None. +Current Vibe local logs do not expose cache-read/cache-write token fields, so +CodeBurn reports cache token counts as `0`. When `meta.json.stats.session_cost` +is present, CodeBurn uses that session total instead of re-estimating from +prompt/completion token prices because it is the best cache-aware cost signal +available in the local log shape. ## Deduplication @@ -29,8 +33,8 @@ Per `mistral-vibe:`. ## Quirks -- **Usage is cumulative per session.** Vibe does not write per-assistant-message token usage into `messages.jsonl`; token counts come from `meta.json.stats.session_prompt_tokens` and `session_completion_tokens`. CodeBurn emits one usage record per Vibe session. -- **Cost prefers Vibe's own model prices.** `meta.json.stats.input_price_per_million` and `output_price_per_million` are used first, with the active model config as a fallback. LiteLLM pricing is only used when Vibe provides no price data. +- **Usage is cumulative per session.** Vibe does not write per-assistant-message token usage into `messages.jsonl`; token counts come from `meta.json.stats.session_prompt_tokens` and `session_completion_tokens`. CodeBurn splits assistant-message tools into their user turns for classification and distributes the cumulative token/cost totals across those assistant calls so session totals remain unchanged. +- **Cost prefers Vibe's own session total.** `meta.json.stats.session_cost` is used first. If it is missing, `meta.json.stats.input_price_per_million` and `output_price_per_million` are used with the active model config as a fallback. LiteLLM pricing is only used when Vibe provides no price data. - **Project names come from metadata.** Discovery uses `meta.json.environment.working_directory` and falls back to the session directory name if that field is missing. - **Tool calls come from messages.** Assistant `tool_calls[*].function.name` is normalized to the standard CodeBurn names (`bash` to `Bash`, `search_replace` to `Edit`, etc.). Bash commands are extracted from `function.arguments.command`. diff --git a/mac/Sources/CodeBurnMenubar/Views/HeatmapSection.swift b/mac/Sources/CodeBurnMenubar/Views/HeatmapSection.swift index e1255a2..b782be3 100644 --- a/mac/Sources/CodeBurnMenubar/Views/HeatmapSection.swift +++ b/mac/Sources/CodeBurnMenubar/Views/HeatmapSection.swift @@ -95,22 +95,25 @@ private struct InsightPillSwitcher: View { let visibleModes: [InsightMode] var body: some View { - HStack(spacing: 4) { - ForEach(visibleModes) { mode in - Button { - selected = mode - } label: { - Text(mode.rawValue) - .font(.system(size: 11, weight: .medium)) - .foregroundStyle(selected == mode ? AnyShapeStyle(.white) : AnyShapeStyle(.secondary)) - .padding(.horizontal, 10) - .padding(.vertical, 4) - .background( - RoundedRectangle(cornerRadius: 6) - .fill(selected == mode ? AnyShapeStyle(Theme.brandAccent) : AnyShapeStyle(Color.secondary.opacity(0.10))) - ) + ScrollView(.horizontal, showsIndicators: false) { + HStack(spacing: 4) { + ForEach(visibleModes) { mode in + Button { + selected = mode + } label: { + Text(mode.rawValue) + .font(.system(size: 11, weight: .medium)) + .fixedSize() + .foregroundStyle(selected == mode ? AnyShapeStyle(.white) : AnyShapeStyle(.secondary)) + .padding(.horizontal, 10) + .padding(.vertical, 4) + .background( + RoundedRectangle(cornerRadius: 6) + .fill(selected == mode ? AnyShapeStyle(Theme.brandAccent) : AnyShapeStyle(Color.secondary.opacity(0.10))) + ) + } + .buttonStyle(.plain) } - .buttonStyle(.plain) } } } diff --git a/src/classifier.ts b/src/classifier.ts index 9a5de49..b0d97bd 100644 --- a/src/classifier.ts +++ b/src/classifier.ts @@ -154,13 +154,22 @@ function classifyConversation(userMessage: string): TaskCategory { } function countRetries(turn: ParsedTurn): number { + const steps: string[][] = [] + for (const call of turn.assistantCalls) { + if (call.toolSequence && call.toolSequence.length > 0) { + steps.push(...call.toolSequence) + } else if (call.tools.length > 0) { + steps.push(call.tools) + } + } + let sawEditBeforeBash = false let sawBashAfterEdit = false let retries = 0 - for (const call of turn.assistantCalls) { - const hasEdit = call.tools.some(t => EDIT_TOOLS.has(t)) - const hasBash = call.tools.some(t => BASH_TOOLS.has(t)) + for (const tools of steps) { + const hasEdit = tools.some(t => EDIT_TOOLS.has(t)) + const hasBash = tools.some(t => BASH_TOOLS.has(t)) if (hasEdit) { if (sawBashAfterEdit) retries++ diff --git a/src/parser.ts b/src/parser.ts index 2700e46..c5a91e1 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -1517,6 +1517,33 @@ function providerCallToTurn(call: ParsedProviderCall): ParsedTurn { // ── Cache Conversion ─────────────────────────────────────────────────── +function providerCallToCachedCall(call: ParsedProviderCall): CachedCall { + return { + provider: call.provider, + model: call.model, + usage: { + inputTokens: call.inputTokens, + outputTokens: call.outputTokens, + cacheCreationInputTokens: call.cacheCreationInputTokens, + cacheReadInputTokens: call.cacheReadInputTokens, + cachedInputTokens: call.cachedInputTokens, + reasoningTokens: call.reasoningTokens, + webSearchRequests: call.webSearchRequests, + cacheCreationOneHourTokens: 0, + }, + costUSD: call.provider === 'mistral-vibe' ? call.costUSD : undefined, + speed: call.speed, + timestamp: call.timestamp, + tools: call.tools, + bashCommands: call.bashCommands, + skills: [], + deduplicationKey: call.deduplicationKey, + project: call.project, + projectPath: call.projectPath, + toolSequence: call.toolSequence, + } +} + function apiCallToCachedCall(call: ParsedApiCall): CachedCall { return { provider: call.provider, @@ -1545,31 +1572,38 @@ function providerCallToCachedTurn(call: ParsedProviderCall): CachedTurn { timestamp: call.timestamp, sessionId: call.sessionId, userMessage: call.userMessage.slice(0, 2000), - calls: [{ - provider: call.provider, - model: call.model, - usage: { - inputTokens: call.inputTokens, - outputTokens: call.outputTokens, - cacheCreationInputTokens: call.cacheCreationInputTokens, - cacheReadInputTokens: call.cacheReadInputTokens, - cachedInputTokens: call.cachedInputTokens, - reasoningTokens: call.reasoningTokens, - webSearchRequests: call.webSearchRequests, - cacheCreationOneHourTokens: 0, - }, - speed: call.speed, - timestamp: call.timestamp, - tools: call.tools, - bashCommands: call.bashCommands, - skills: [], - deduplicationKey: call.deduplicationKey, - project: call.project, - projectPath: call.projectPath, - }], + calls: [providerCallToCachedCall(call)], } } +function providerCallsToCachedTurns(calls: ParsedProviderCall[]): CachedTurn[] { + const turns: CachedTurn[] = [] + const grouped = new Map() + + for (const call of calls) { + if (!call.turnId) { + turns.push(providerCallToCachedTurn(call)) + continue + } + + const key = `${call.sessionId}\0${call.turnId}` + let turn = grouped.get(key) + if (!turn) { + turn = { + timestamp: call.timestamp, + sessionId: call.sessionId, + userMessage: call.userMessage.slice(0, 2000), + calls: [], + } + grouped.set(key, turn) + turns.push(turn) + } + turn.calls.push(providerCallToCachedCall(call)) + } + + return turns +} + function cachedCallToApiCall(call: CachedCall): ParsedApiCall { const u = call.usage const outputForCost = call.provider === 'claude' @@ -1592,7 +1626,7 @@ function cachedCallToApiCall(call: CachedCall): ParsedApiCall { reasoningTokens: u.reasoningTokens, webSearchRequests: u.webSearchRequests, }, - costUSD, + costUSD: call.costUSD ?? costUSD, tools: call.tools, mcpTools: extractMcpTools(call.tools), skills: call.skills, @@ -1603,6 +1637,7 @@ function cachedCallToApiCall(call: CachedCall): ParsedApiCall { bashCommands: call.bashCommands, deduplicationKey: call.deduplicationKey, cacheCreationOneHourTokens: u.cacheCreationOneHourTokens || undefined, + toolSequence: call.toolSequence, } } @@ -1725,10 +1760,11 @@ async function parseProviderSources( ) try { - const turns: CachedTurn[] = [] + const providerCalls: ParsedProviderCall[] = [] for await (const call of parser.parse()) { - turns.push(providerCallToCachedTurn(call)) + providerCalls.push(call) } + const turns = providerCallsToCachedTurns(providerCalls) section.files[source.path] = { fingerprint: fp, mcpInventory: [], turns } didParse = true ;(diskCache as { _dirty?: boolean })._dirty = true diff --git a/src/providers/gemini.ts b/src/providers/gemini.ts index 3f4b590..2ea71d4 100644 --- a/src/providers/gemini.ts +++ b/src/providers/gemini.ts @@ -67,6 +67,8 @@ function parseSession(data: GeminiSession, seenKeys: Set): ParsedProvide const results: ParsedProviderCall[] = [] let lastUserMessage = '' + let turnOrdinal = 0 + let currentTurnId = `${data.sessionId}:prelude` let geminiOrdinal = 0 for (const msg of data.messages) { @@ -76,6 +78,7 @@ function parseSession(data: GeminiSession, seenKeys: Set): ParsedProvide } else if (typeof msg.content === 'string') { lastUserMessage = msg.content.slice(0, 500) } + currentTurnId = `${data.sessionId}:turn-${turnOrdinal++}` continue } @@ -136,6 +139,7 @@ function parseSession(data: GeminiSession, seenKeys: Set): ParsedProvide timestamp: tsDate.toISOString(), speed: 'standard', deduplicationKey: dedupKey, + turnId: currentTurnId, userMessage: lastUserMessage, sessionId: data.sessionId, }) diff --git a/src/providers/goose.ts b/src/providers/goose.ts index 9f4abe5..6a10d1d 100644 --- a/src/providers/goose.ts +++ b/src/providers/goose.ts @@ -80,14 +80,15 @@ function parseModelConfig(raw: string | null): ModelConfig { } } -function extractToolsFromMessages(db: SqliteDatabase, sessionId: string): { tools: string[]; bashCommands: string[] } { +function extractToolsFromMessages(db: SqliteDatabase, sessionId: string): { tools: string[]; bashCommands: string[]; toolSequence: string[][] } { const tools: string[] = [] const bashCommands: string[] = [] const seen = new Set() + const toolSequence: string[][] = [] try { const rows = db.query<{ content_json: Uint8Array | string }>( - "SELECT CAST(content_json AS BLOB) AS content_json FROM messages WHERE session_id = ? AND role = 'assistant' AND content_json LIKE '%toolRequest%'", + "SELECT CAST(content_json AS BLOB) AS content_json FROM messages WHERE session_id = ? AND role = 'assistant' AND content_json LIKE '%toolRequest%' ORDER BY created_timestamp ASC", [sessionId], ) @@ -98,6 +99,7 @@ function extractToolsFromMessages(db: SqliteDatabase, sessionId: string): { tool } catch { continue } + const msgTools: string[] = [] for (const item of items) { if (item.type !== 'toolRequest') continue const rawName = item.toolCall?.value?.name ?? '' @@ -107,6 +109,7 @@ function extractToolsFromMessages(db: SqliteDatabase, sessionId: string): { tool seen.add(mapped) tools.push(mapped) } + msgTools.push(mapped) if (mapped === 'Bash') { const cmd = item.toolCall?.value?.arguments?.command if (typeof cmd === 'string') { @@ -116,10 +119,11 @@ function extractToolsFromMessages(db: SqliteDatabase, sessionId: string): { tool } } } + if (msgTools.length > 0) toolSequence.push(msgTools) } } catch { /* best-effort */ } - return { tools, bashCommands } + return { tools, bashCommands, toolSequence } } function getFirstUserMessage(db: SqliteDatabase, sessionId: string): string { @@ -179,7 +183,7 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars const model = config.model_name ?? 'unknown' const costUSD = calculateCost(model, inputTokens, outputTokens, 0, 0, 0) - const { tools, bashCommands } = extractToolsFromMessages(db, sessionId) + const { tools, bashCommands, toolSequence } = extractToolsFromMessages(db, sessionId) const userMessage = getFirstUserMessage(db, sessionId) const raw = session.updated_at || session.created_at || '' @@ -200,6 +204,7 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars costUSD, tools, bashCommands, + toolSequence: toolSequence.length > 1 ? toolSequence : undefined, timestamp: ts.toISOString(), speed: 'standard', deduplicationKey: dedupKey, diff --git a/src/providers/kiro.ts b/src/providers/kiro.ts index 118bd06..7d616d7 100644 --- a/src/providers/kiro.ts +++ b/src/providers/kiro.ts @@ -86,6 +86,7 @@ function parseChatFile(data: KiroChatFile, sessionId: string, project: string, s let pendingUserMessage = '' const allTools: string[] = [] + const toolSequence: string[][] = [] for (const msg of chat) { if (msg.role === 'human') { @@ -93,7 +94,9 @@ function parseChatFile(data: KiroChatFile, sessionId: string, project: string, s pendingUserMessage = msg.content.slice(0, 500) } if (msg.role === 'bot') { - allTools.push(...extractToolNames(msg.content)) + const msgTools = extractToolNames(msg.content) + allTools.push(...msgTools) + if (msgTools.length > 0) toolSequence.push(msgTools) } } @@ -125,6 +128,7 @@ function parseChatFile(data: KiroChatFile, sessionId: string, project: string, s costUSD, tools: [...new Set(allTools)], bashCommands: [], + toolSequence: toolSequence.length > 1 ? toolSequence : undefined, timestamp, speed: 'standard', deduplicationKey: dedupKey, diff --git a/src/providers/mistral-vibe.ts b/src/providers/mistral-vibe.ts index 7feb988..48e7203 100644 --- a/src/providers/mistral-vibe.ts +++ b/src/providers/mistral-vibe.ts @@ -38,6 +38,7 @@ const toolNameMap: Record = { type VibeStats = { session_prompt_tokens?: number session_completion_tokens?: number + session_cost?: number input_price_per_million?: number output_price_per_million?: number tokens_per_second?: number @@ -75,6 +76,8 @@ type VibeToolCall = { type VibeMessage = { role?: string content?: unknown + message_id?: string + timestamp?: string tool_calls?: VibeToolCall[] | null } @@ -179,6 +182,9 @@ function safeNumber(value: unknown): number { function calculateSessionCost(metadata: VibeMetadata, model: string, inputTokens: number, outputTokens: number): number { const stats = metadata.stats ?? {} + const sessionCost = safeNumber(stats.session_cost) + if (sessionCost > 0) return sessionCost + const configured = activeModelConfig(metadata) const inputPrice = safeNumber(stats.input_price_per_million) || safeNumber(configured?.input_price) const outputPrice = safeNumber(stats.output_price_per_million) || safeNumber(configured?.output_price) @@ -216,26 +222,41 @@ function parseToolArguments(raw: string | Record | null | undef } } +function extractMessageTools(message: VibeMessage): { tools: string[]; bashCommands: string[] } { + const tools: string[] = [] + const bashCommands: string[] = [] + + if (message.role !== 'assistant') return { tools, bashCommands } + + for (const toolCall of message.tool_calls ?? []) { + const rawName = toolCall.function?.name + if (!rawName) continue + + const mappedName = toolNameMap[rawName] ?? rawName + tools.push(mappedName) + + if (mappedName !== 'Bash') continue + const args = parseToolArguments(toolCall.function?.arguments) + const command = args['command'] + if (typeof command === 'string') { + bashCommands.push(...extractBashCommands(command)) + } + } + + return { + tools: [...new Set(tools)], + bashCommands: [...new Set(bashCommands)], + } +} + function extractTools(messages: VibeMessage[]): { tools: string[]; bashCommands: string[] } { const tools: string[] = [] const bashCommands: string[] = [] for (const message of messages) { - if (message.role !== 'assistant') continue - for (const toolCall of message.tool_calls ?? []) { - const rawName = toolCall.function?.name - if (!rawName) continue - - const mappedName = toolNameMap[rawName] ?? rawName - tools.push(mappedName) - - if (mappedName !== 'Bash') continue - const args = parseToolArguments(toolCall.function?.arguments) - const command = args['command'] - if (typeof command === 'string') { - bashCommands.push(...extractBashCommands(command)) - } - } + const extracted = extractMessageTools(message) + tools.push(...extracted.tools) + bashCommands.push(...extracted.bashCommands) } return { @@ -267,6 +288,17 @@ function firstUserMessage(messages: VibeMessage[], fallback?: string | null): st return (fallback ?? '').slice(0, 500) } +function allocateInteger(total: number, index: number, count: number): number { + if (count <= 1) return total + const base = Math.floor(total / count) + const remainder = total % count + return base + (index < remainder ? 1 : 0) +} + +function allocateCost(total: number, count: number): number { + return count <= 1 ? total : total / count +} + function createParser(source: SessionSource, seenKeys: Set): SessionParser { return { async *parse(): AsyncGenerator { @@ -281,33 +313,85 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars if (inputTokens === 0 && outputTokens === 0) return const sessionId = metadata.session_id || basename(source.path) - const deduplicationKey = `mistral-vibe:${sessionId}` - if (seenKeys.has(deduplicationKey)) return - seenKeys.add(deduplicationKey) - const messages = await readMessages(messagesPath) const model = resolveModel(metadata) - const { tools, bashCommands } = extractTools(messages) const costUSD = calculateSessionCost(metadata, model, inputTokens, outputTokens) + const assistantMessages = messages.filter(m => m.role === 'assistant') + const fallbackTimestamp = metadata.end_time ?? metadata.start_time ?? '' - yield { - provider: 'mistral-vibe', - model, - inputTokens, - outputTokens, - cacheCreationInputTokens: 0, - cacheReadInputTokens: 0, - cachedInputTokens: 0, - reasoningTokens: 0, - webSearchRequests: 0, - costUSD, - tools, - bashCommands, - timestamp: metadata.end_time ?? metadata.start_time ?? '', - speed: 'standard', - deduplicationKey, - userMessage: firstUserMessage(messages, metadata.title), - sessionId, + if (assistantMessages.length === 0) { + const deduplicationKey = `mistral-vibe:${sessionId}` + if (seenKeys.has(deduplicationKey)) return + seenKeys.add(deduplicationKey) + const { tools, bashCommands } = extractTools(messages) + + yield { + provider: 'mistral-vibe', + model, + inputTokens, + outputTokens, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + cachedInputTokens: 0, + reasoningTokens: 0, + webSearchRequests: 0, + costUSD, + tools, + bashCommands, + timestamp: fallbackTimestamp, + speed: 'standard', + deduplicationKey, + userMessage: firstUserMessage(messages, metadata.title), + sessionId, + } + return + } + + let currentUserMessage = (metadata.title ?? '').slice(0, 500) + let turnOrdinal = 0 + let currentTurnId = `${sessionId}:prelude` + let assistantOrdinal = 0 + + for (const message of messages) { + if (message.role === 'user') { + const text = normalizeContent(message.content).trim() + if (text) currentUserMessage = text.slice(0, 500) + currentTurnId = `${sessionId}:turn-${turnOrdinal++}` + continue + } + + if (message.role !== 'assistant') continue + + const messageKey = message.message_id || `idx-${assistantOrdinal}` + const deduplicationKey = `mistral-vibe:${sessionId}:${messageKey}` + const allocationIndex = assistantOrdinal + assistantOrdinal++ + + if (seenKeys.has(deduplicationKey)) continue + seenKeys.add(deduplicationKey) + + const { tools, bashCommands } = extractMessageTools(message) + + yield { + provider: 'mistral-vibe', + model, + inputTokens: allocateInteger(inputTokens, allocationIndex, assistantMessages.length), + outputTokens: allocateInteger(outputTokens, allocationIndex, assistantMessages.length), + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + cachedInputTokens: 0, + reasoningTokens: 0, + webSearchRequests: 0, + costUSD: allocateCost(costUSD, assistantMessages.length), + tools, + bashCommands, + timestamp: message.timestamp ?? fallbackTimestamp, + speed: 'standard', + deduplicationKey, + turnId: currentTurnId, + userMessage: currentUserMessage, + sessionId, + } } }, } diff --git a/src/providers/types.ts b/src/providers/types.ts index 90d5e1c..8e0c937 100644 --- a/src/providers/types.ts +++ b/src/providers/types.ts @@ -25,6 +25,8 @@ export type ParsedProviderCall = { timestamp: string speed: 'standard' | 'fast' deduplicationKey: string + turnId?: string + toolSequence?: string[][] userMessage: string sessionId: string project?: string diff --git a/src/session-cache.ts b/src/session-cache.ts index 2537a7d..5534ad6 100644 --- a/src/session-cache.ts +++ b/src/session-cache.ts @@ -21,6 +21,7 @@ export type CachedCall = { provider: string model: string usage: CachedUsage + costUSD?: number speed: 'standard' | 'fast' timestamp: string tools: string[] @@ -29,6 +30,7 @@ export type CachedCall = { deduplicationKey: string project?: string projectPath?: string + toolSequence?: string[][] } export type CachedTurn = { @@ -65,7 +67,7 @@ export type SessionCache = { // ── Constants ────────────────────────────────────────────────────────── -export const CACHE_VERSION = 1 +export const CACHE_VERSION = 2 const CACHE_FILE = 'session-cache.json' const TEMP_FILE_MAX_AGE_MS = 5 * 60 * 1000 @@ -147,11 +149,13 @@ function validateCall(c: unknown): c is CachedCall { && typeof o['deduplicationKey'] === 'string' && typeof o['timestamp'] === 'string' && (o['speed'] === 'standard' || o['speed'] === 'fast') + && isOptionalNum(o['costUSD']) && isStringArray(o['tools']) && isStringArray(o['bashCommands']) && isStringArray(o['skills']) && isOptionalString(o['project']) && isOptionalString(o['projectPath']) + && (o['toolSequence'] === undefined || (Array.isArray(o['toolSequence']) && (o['toolSequence'] as unknown[]).every(s => isStringArray(s)))) && validateUsage(o['usage']) } diff --git a/src/types.ts b/src/types.ts index 312906d..7e362e7 100644 --- a/src/types.ts +++ b/src/types.ts @@ -84,6 +84,7 @@ export type ParsedApiCall = { bashCommands: string[] deduplicationKey: string cacheCreationOneHourTokens?: number + toolSequence?: string[][] } export type TaskCategory = diff --git a/tests/classifier.test.ts b/tests/classifier.test.ts index 6a02d64..42c3441 100644 --- a/tests/classifier.test.ts +++ b/tests/classifier.test.ts @@ -151,3 +151,46 @@ describe('classifyTurn — feature vs debugging precedence (#196)', () => { expect(c.category).toBe('debugging') }) }) + +describe('classifyTurn — retry detection via toolSequence', () => { + it('detects retries from multi-call turns (Claude-style)', () => { + const turn = makeTurn([ + makeCall({ tools: ['Edit'] }), + makeCall({ tools: ['Bash'] }), + makeCall({ tools: ['Edit'] }), + ], 'fix the build') + const c = classifyTurn(turn) + expect(c.retries).toBe(1) + }) + + it('detects retries from toolSequence on a single call (Kiro/Goose-style)', () => { + const call = makeCall({ tools: ['Edit', 'Bash'] }) + call.toolSequence = [['Edit'], ['Bash'], ['Edit']] + const turn = makeTurn([call], 'fix the build') + const c = classifyTurn(turn) + expect(c.retries).toBe(1) + }) + + it('returns 0 retries for single call without toolSequence', () => { + const call = makeCall({ tools: ['Edit', 'Bash'] }) + const turn = makeTurn([call], 'fix the build') + const c = classifyTurn(turn) + expect(c.retries).toBe(0) + }) + + it('counts multiple retries from toolSequence', () => { + const call = makeCall({ tools: ['Edit', 'Bash'] }) + call.toolSequence = [['Edit'], ['Bash'], ['Edit'], ['Bash'], ['Edit']] + const turn = makeTurn([call], 'fix the build') + const c = classifyTurn(turn) + expect(c.retries).toBe(2) + }) + + it('ignores toolSequence with only one step', () => { + const call = makeCall({ tools: ['Edit', 'Bash'] }) + call.toolSequence = [['Edit', 'Bash']] + const turn = makeTurn([call], 'fix the build') + const c = classifyTurn(turn) + expect(c.retries).toBe(0) + }) +}) diff --git a/tests/provider-turn-grouping.test.ts b/tests/provider-turn-grouping.test.ts new file mode 100644 index 0000000..27ff5a7 --- /dev/null +++ b/tests/provider-turn-grouping.test.ts @@ -0,0 +1,170 @@ +import { mkdir, mkdtemp, rm, writeFile } from 'fs/promises' +import { join } from 'path' +import { tmpdir } from 'os' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import type { DateRange } from '../src/types.js' + +let home: string +let cacheDir: string +let vibeHome: string +let originalHome: string | undefined +let originalCacheDir: string | undefined +let originalVibeHome: string | undefined +let clearParserCache: (() => void) | undefined + +beforeEach(async () => { + home = await mkdtemp(join(tmpdir(), 'codeburn-turn-group-home-')) + cacheDir = await mkdtemp(join(tmpdir(), 'codeburn-turn-group-cache-')) + vibeHome = await mkdtemp(join(tmpdir(), 'codeburn-turn-group-vibe-')) + originalHome = process.env['HOME'] + originalCacheDir = process.env['CODEBURN_CACHE_DIR'] + originalVibeHome = process.env['VIBE_HOME'] + process.env['HOME'] = home + process.env['CODEBURN_CACHE_DIR'] = cacheDir + process.env['VIBE_HOME'] = vibeHome +}) + +afterEach(async () => { + clearParserCache?.() + clearParserCache = undefined + vi.resetModules() + if (originalHome === undefined) delete process.env['HOME'] + else process.env['HOME'] = originalHome + if (originalCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = originalCacheDir + if (originalVibeHome === undefined) delete process.env['VIBE_HOME'] + else process.env['VIBE_HOME'] = originalVibeHome + await rm(home, { recursive: true, force: true }) + await rm(cacheDir, { recursive: true, force: true }) + await rm(vibeHome, { recursive: true, force: true }) +}) + +function dayRange(): DateRange { + return { + start: new Date('2026-05-16T00:00:00.000Z'), + end: new Date('2026-05-16T23:59:59.999Z'), + } +} + +async function loadParser() { + vi.resetModules() + const parser = await import('../src/parser.js') + clearParserCache = parser.clearSessionCache + return parser.parseAllSessions +} + +describe('provider turn grouping', () => { + it('groups Gemini assistant messages under their user turn so retries are counted', async () => { + const chatsDir = join(home, '.gemini', 'tmp', 'project-a', 'chats') + await mkdir(chatsDir, { recursive: true }) + await writeFile(join(chatsDir, 'session-gemini.json'), JSON.stringify({ + sessionId: 'gemini-session-1', + startTime: '2026-05-16T10:00:00.000Z', + messages: [ + { id: 'u1', timestamp: '2026-05-16T10:00:00.000Z', type: 'user', content: 'implement parser update in src/parser.ts' }, + { + id: 'g1', + timestamp: '2026-05-16T10:00:05.000Z', + type: 'gemini', + content: 'editing', + model: 'gemini-3.1-pro-preview', + tokens: { input: 100, output: 30 }, + toolCalls: [{ id: 't1', name: 'edit_file', args: { path: 'src/parser.ts' } }], + }, + { + id: 'g2', + timestamp: '2026-05-16T10:00:10.000Z', + type: 'gemini', + content: 'testing', + model: 'gemini-3.1-pro-preview', + tokens: { input: 80, output: 20 }, + toolCalls: [{ id: 't2', name: 'run_command', args: { command: 'npm test' } }], + }, + { + id: 'g3', + timestamp: '2026-05-16T10:00:15.000Z', + type: 'gemini', + content: 'fixing after test', + model: 'gemini-3.1-pro-preview', + tokens: { input: 90, output: 25 }, + toolCalls: [{ id: 't3', name: 'edit_file', args: { path: 'src/parser.ts' } }], + }, + ], + })) + + const parseAllSessions = await loadParser() + const projects = await parseAllSessions(dayRange(), 'gemini') + const session = projects[0]!.sessions[0]! + const turn = session.turns[0]! + + expect(session.turns).toHaveLength(1) + expect(turn.assistantCalls.map(call => call.deduplicationKey)).toEqual([ + 'gemini:gemini-session-1:g1', + 'gemini:gemini-session-1:g2', + 'gemini:gemini-session-1:g3', + ]) + expect(turn.hasEdits).toBe(true) + expect(turn.retries).toBe(1) + expect(session.categoryBreakdown[turn.category].editTurns).toBe(1) + expect(session.categoryBreakdown[turn.category].oneShotTurns).toBe(0) + }) + + it('groups Mistral Vibe assistant messages and uses Vibe session_cost when present', async () => { + const sessionDir = join(vibeHome, 'logs', 'session', 'session_20260516_100000_vibe') + await mkdir(sessionDir, { recursive: true }) + await writeFile(join(sessionDir, 'meta.json'), JSON.stringify({ + session_id: 'vibe-session-1', + start_time: '2026-05-16T10:00:00.000Z', + end_time: '2026-05-16T10:01:00.000Z', + environment: { working_directory: '/Users/test/project-a' }, + stats: { + session_prompt_tokens: 300, + session_completion_tokens: 90, + session_cost: 0.123456, + input_price_per_million: 100, + output_price_per_million: 100, + }, + config: { active_model: 'mistral-medium-3.5', models: [] }, + title: 'vibe parser update', + })) + await writeFile(join(sessionDir, 'messages.jsonl'), [ + { role: 'user', content: 'implement parser update in src/providers/mistral-vibe.ts', message_id: 'u1' }, + { + role: 'assistant', + content: 'editing', + message_id: 'a1', + tool_calls: [{ id: 't1', type: 'function', function: { name: 'search_replace', arguments: '{"file_path":"src/providers/mistral-vibe.ts"}' } }], + }, + { + role: 'assistant', + content: 'testing', + message_id: 'a2', + tool_calls: [{ id: 't2', type: 'function', function: { name: 'bash', arguments: '{"command":"npm test"}' } }], + }, + { + role: 'assistant', + content: 'fixing after test', + message_id: 'a3', + tool_calls: [{ id: 't3', type: 'function', function: { name: 'write_file', arguments: '{"path":"src/providers/mistral-vibe.ts"}' } }], + }, + ].map(message => JSON.stringify(message)).join('\n') + '\n') + + const parseAllSessions = await loadParser() + const projects = await parseAllSessions(dayRange(), 'mistral-vibe') + const session = projects[0]!.sessions[0]! + const turn = session.turns[0]! + + expect(session.turns).toHaveLength(1) + expect(turn.assistantCalls.map(call => call.deduplicationKey)).toEqual([ + 'mistral-vibe:vibe-session-1:a1', + 'mistral-vibe:vibe-session-1:a2', + 'mistral-vibe:vibe-session-1:a3', + ]) + expect(turn.retries).toBe(1) + expect(session.totalCostUSD).toBeCloseTo(0.123456, 8) + expect(session.totalInputTokens).toBe(300) + expect(session.totalOutputTokens).toBe(90) + expect(session.categoryBreakdown[turn.category].oneShotTurns).toBe(0) + }) +}) diff --git a/tests/providers/mistral-vibe.test.ts b/tests/providers/mistral-vibe.test.ts index 51bc03c..b8d49be 100644 --- a/tests/providers/mistral-vibe.test.ts +++ b/tests/providers/mistral-vibe.test.ts @@ -29,6 +29,7 @@ function metadata(opts: { cwd?: string input?: number output?: number + sessionCost?: number inputPrice?: number outputPrice?: number activeModel?: string @@ -49,6 +50,7 @@ function metadata(opts: { stats: { session_prompt_tokens: opts.input ?? 2000, session_completion_tokens: opts.output ?? 3000, + session_cost: opts.sessionCost, input_price_per_million: opts.inputPrice ?? 1.5, output_price_per_million: opts.outputPrice ?? 7.5, tokens_per_second: 42, @@ -202,7 +204,22 @@ describe('mistral-vibe provider - parsing', () => { expect(call.timestamp).toBe('2026-05-11T10:05:00+00:00') expect(call.userMessage).toBe('track Mistral Vibe usage') expect(call.sessionId).toBe('session-abc123') - expect(call.deduplicationKey).toBe('mistral-vibe:session-abc123') + expect(call.deduplicationKey).toBe('mistral-vibe:session-abc123:msg-assistant-1') + }) + + it('prefers Vibe session_cost over price-derived estimates when present', async () => { + const sessionDir = await writeSession('session_20260511_100000_sessiona', metadata({ + input: 364147, + output: 1731, + sessionCost: 0.381681, + inputPrice: 100, + outputPrice: 100, + })) + + const calls = await collect(sessionDir, createMistralVibeProvider(tmpDir)) + + expect(calls).toHaveLength(1) + expect(calls[0]!.costUSD).toBe(0.381681) }) it('uses configured model prices when stats omit prices', async () => {