Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 2013ecb

Browse files
ozymandiashhozymandiashh
authored
Track agent calls across providers (getagentseal#340)
1 parent 303c945 commit 2013ecb

7 files changed

Lines changed: 503 additions & 72 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,15 @@
33
## Unreleased
44

55
### Added (CLI)
6+
- **Agent and subagent tracking coverage.** Gemini sessions now emit one
7+
provider call per assistant message with token usage instead of one aggregate
8+
call per session, preserving per-message tools, bash commands, timestamps,
9+
and nearest user prompts. Existing cached aggregate Gemini entries are
10+
reparsed so the new per-message shape takes effect, and per-tool counts may
11+
increase because repeated tools are now attributed to the specific Gemini
12+
message that used them. Claude discovery also scans direct project-level
13+
`subagents/*.jsonl` files, and Codex agent tool normalization is covered by
14+
regression tests. Addresses #336.
615
- **Multiple subscription plans can be tracked at the same time.**
716
`codeburn plan set` now stores plans in a provider-keyed `plans` map, so
817
setting a Codex custom plan no longer overwrites an existing Claude plan.

‎src/parser.ts‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1321,18 +1321,24 @@ async function parseSessionFile(
13211321

13221322
async function collectJsonlFiles(dirPath: string): Promise<string[]> {
13231323
const files = await readdir(dirPath).catch(() => [])
1324-
const jsonlFiles = files.filter(f => f.endsWith('.jsonl')).map(f => join(dirPath, f))
1324+
const jsonlFiles = new Set(files.filter(f => f.endsWith('.jsonl')).map(f => join(dirPath, f)))
1325+
1326+
const directSubagentsPath = join(dirPath, 'subagents')
1327+
const directSubFiles = await readdir(directSubagentsPath).catch(() => [])
1328+
for (const sf of directSubFiles) {
1329+
if (sf.endsWith('.jsonl')) jsonlFiles.add(join(directSubagentsPath, sf))
1330+
}
13251331

13261332
for (const entry of files) {
13271333
if (entry.endsWith('.jsonl')) continue
13281334
const subagentsPath = join(dirPath, entry, 'subagents')
13291335
const subFiles = await readdir(subagentsPath).catch(() => [])
13301336
for (const sf of subFiles) {
1331-
if (sf.endsWith('.jsonl')) jsonlFiles.push(join(subagentsPath, sf))
1337+
if (sf.endsWith('.jsonl')) jsonlFiles.add(join(subagentsPath, sf))
13321338
}
13331339
}
13341340

1335-
return jsonlFiles
1341+
return [...jsonlFiles]
13361342
}
13371343

13381344
async function scanProjectDirs(
@@ -1639,6 +1645,14 @@ function getOrCreateProviderSection(cache: SessionCache, provider: string): Prov
16391645
return section
16401646
}
16411647

1648+
function cachedFileNeedsProviderReparse(providerName: string, cached: CachedFile): boolean {
1649+
if (providerName !== 'gemini') return false
1650+
1651+
return cached.turns.some(turn =>
1652+
turn.calls.some(call => call.deduplicationKey === `gemini:${turn.sessionId}`),
1653+
)
1654+
}
1655+
16421656
const warnedProviderReadFailures = new Set<string>()
16431657

16441658
function warnProviderReadFailureOnce(providerName: string, err: unknown): void {
@@ -1674,9 +1688,10 @@ async function parseProviderSources(
16741688
const fp = await fingerprintFile(source.path)
16751689
if (!fp) continue
16761690

1677-
const action = reconcileFile(fp, section.files[source.path])
1678-
if (action.action === 'unchanged') {
1679-
unchangedSources.push({ source, cached: section.files[source.path]! })
1691+
const cached = section.files[source.path]
1692+
const action = reconcileFile(fp, cached)
1693+
if (action.action === 'unchanged' && cached && !cachedFileNeedsProviderReparse(providerName, cached)) {
1694+
unchangedSources.push({ source, cached })
16801695
} else {
16811696
changedSources.push({ source, fp })
16821697
}

‎src/providers/gemini.ts‎

Lines changed: 63 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -66,84 +66,81 @@ type GeminiSession = {
6666
function parseSession(data: GeminiSession, seenKeys: Set<string>): ParsedProviderCall[] {
6767
const results: ParsedProviderCall[] = []
6868

69-
const geminiMessages = data.messages.filter(m => m.type === 'gemini' && m.tokens && m.model)
70-
if (geminiMessages.length === 0) return results
71-
72-
const dedupKey = `gemini:${data.sessionId}`
73-
if (seenKeys.has(dedupKey)) return results
74-
seenKeys.add(dedupKey)
75-
76-
let totalInput = 0
77-
let totalOutput = 0
78-
let totalCached = 0
79-
let totalThoughts = 0
80-
const allTools: string[] = []
81-
const bashCommands: string[] = []
82-
let model = ''
83-
84-
for (const msg of geminiMessages) {
85-
const t = msg.tokens!
86-
totalInput += t.input ?? 0
87-
totalOutput += t.output ?? 0
88-
totalCached += t.cached ?? 0
89-
totalThoughts += t.thoughts ?? 0
90-
if (msg.model && !model) model = msg.model
69+
let lastUserMessage = ''
70+
let geminiOrdinal = 0
71+
72+
for (const msg of data.messages) {
73+
if (msg.type === 'user') {
74+
if (Array.isArray(msg.content)) {
75+
lastUserMessage = msg.content.map(c => c.text).join(' ').slice(0, 500)
76+
} else if (typeof msg.content === 'string') {
77+
lastUserMessage = msg.content.slice(0, 500)
78+
}
79+
continue
80+
}
81+
82+
if (msg.type !== 'gemini' || !msg.tokens || !msg.model) continue
83+
84+
const t = msg.tokens
85+
const totalInput = t.input ?? 0
86+
const totalOutput = t.output ?? 0
87+
const totalCached = t.cached ?? 0
88+
const totalThoughts = t.thoughts ?? 0
89+
if (totalInput === 0 && totalOutput === 0 && totalCached === 0 && totalThoughts === 0) continue
90+
91+
const messageKey = msg.id || `idx-${geminiOrdinal}`
92+
geminiOrdinal++
93+
const dedupKey = `gemini:${data.sessionId}:${messageKey}`
94+
if (seenKeys.has(dedupKey)) continue
95+
96+
const tools: string[] = []
97+
const bashCommands: string[] = []
9198

9299
if (msg.toolCalls) {
93100
for (const tc of msg.toolCalls) {
94101
const mapped = toolNameMap[tc.displayName ?? ''] ?? toolNameMap[tc.name] ?? tc.displayName ?? tc.name
95-
allTools.push(mapped)
102+
tools.push(mapped)
96103
if (mapped === 'Bash' && tc.args && typeof tc.args.command === 'string') {
97104
bashCommands.push(...extractBashCommands(tc.args.command))
98105
}
99106
}
100107
}
101-
}
102-
103-
if (totalInput === 0 && totalOutput === 0) return results
104108

105-
// Gemini's `input` count includes `cached` tokens as a subset, so fresh input
106-
// must subtract cached to avoid double-charging at both rates.
107-
const freshInput = totalInput - totalCached
108-
109-
let userMessage = ''
110-
const firstUser = data.messages.find(m => m.type === 'user')
111-
if (firstUser) {
112-
if (Array.isArray(firstUser.content)) {
113-
userMessage = firstUser.content.map(c => c.text).join(' ').slice(0, 500)
114-
} else if (typeof firstUser.content === 'string') {
115-
userMessage = firstUser.content.slice(0, 500)
116-
}
109+
// Gemini's `input` count includes `cached` tokens as a subset, so fresh
110+
// input must subtract cached to avoid double-charging at both rates.
111+
const freshInput = Math.max(0, totalInput - totalCached)
112+
113+
const tsDate = new Date(msg.timestamp || data.startTime)
114+
if (isNaN(tsDate.getTime()) || tsDate.getTime() < 1_000_000_000_000) continue
115+
116+
seenKeys.add(dedupKey)
117+
118+
// Gemini bills thoughts at the output token rate; calculateCost does not
119+
// accept a reasoning parameter, so fold thoughts into the output count for
120+
// pricing while keeping outputTokens / reasoningTokens reported separately.
121+
const costUSD = calculateCost(msg.model, freshInput, totalOutput + totalThoughts, 0, totalCached, 0)
122+
123+
results.push({
124+
provider: 'gemini',
125+
model: msg.model,
126+
inputTokens: freshInput,
127+
outputTokens: totalOutput,
128+
cacheCreationInputTokens: 0,
129+
cacheReadInputTokens: totalCached,
130+
cachedInputTokens: totalCached,
131+
reasoningTokens: totalThoughts,
132+
webSearchRequests: 0,
133+
costUSD,
134+
tools: [...new Set(tools)],
135+
bashCommands: [...new Set(bashCommands)],
136+
timestamp: tsDate.toISOString(),
137+
speed: 'standard',
138+
deduplicationKey: dedupKey,
139+
userMessage: lastUserMessage,
140+
sessionId: data.sessionId,
141+
})
117142
}
118143

119-
const tsDate = new Date(data.startTime)
120-
if (isNaN(tsDate.getTime()) || tsDate.getTime() < 1_000_000_000_000) return results
121-
122-
// Gemini bills thoughts at the output token rate; calculateCost does not
123-
// accept a reasoning parameter, so fold thoughts into the output count for
124-
// pricing while keeping outputTokens / reasoningTokens reported separately.
125-
const costUSD = calculateCost(model, freshInput, totalOutput + totalThoughts, 0, totalCached, 0)
126-
127-
results.push({
128-
provider: 'gemini',
129-
model,
130-
inputTokens: freshInput,
131-
outputTokens: totalOutput,
132-
cacheCreationInputTokens: 0,
133-
cacheReadInputTokens: totalCached,
134-
cachedInputTokens: totalCached,
135-
reasoningTokens: totalThoughts,
136-
webSearchRequests: 0,
137-
costUSD,
138-
tools: [...new Set(allTools)],
139-
bashCommands: [...new Set(bashCommands)],
140-
timestamp: tsDate.toISOString(),
141-
speed: 'standard',
142-
deduplicationKey: dedupKey,
143-
userMessage,
144-
sessionId: data.sessionId,
145-
})
146-
147144
return results
148145
}
149146

‎tests/parser-gemini-cache.test.ts‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import { mkdir, mkdtemp, readFile, rm, stat, writeFile } from 'fs/promises'
2+
import { tmpdir } from 'os'
3+
import { join } from 'path'
4+
5+
import { afterEach, beforeEach, describe, expect, it } from 'vitest'
6+
7+
import { clearSessionCache, parseAllSessions } from '../src/parser.js'
8+
import { CACHE_VERSION, computeEnvFingerprint } from '../src/session-cache.js'
9+
import type { DateRange } from '../src/types.js'
10+
11+
let home: string
12+
let cacheDir: string
13+
let previousHome: string | undefined
14+
let previousCacheDir: string | undefined
15+
16+
beforeEach(async () => {
17+
home = await mkdtemp(join(tmpdir(), 'codeburn-gemini-home-'))
18+
cacheDir = await mkdtemp(join(tmpdir(), 'codeburn-gemini-cache-'))
19+
previousHome = process.env['HOME']
20+
previousCacheDir = process.env['CODEBURN_CACHE_DIR']
21+
process.env['HOME'] = home
22+
process.env['CODEBURN_CACHE_DIR'] = cacheDir
23+
})
24+
25+
afterEach(async () => {
26+
clearSessionCache()
27+
if (previousHome === undefined) delete process.env['HOME']
28+
else process.env['HOME'] = previousHome
29+
if (previousCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR']
30+
else process.env['CODEBURN_CACHE_DIR'] = previousCacheDir
31+
await rm(home, { recursive: true, force: true })
32+
await rm(cacheDir, { recursive: true, force: true })
33+
})
34+
35+
describe('Gemini session cache migration', () => {
36+
it('reparses cached legacy aggregate Gemini entries into granular calls', async () => {
37+
const chatsDir = join(home, '.gemini', 'tmp', 'project-a', 'chats')
38+
await mkdir(chatsDir, { recursive: true })
39+
const sessionPath = join(chatsDir, 'session-2026-05-16.json')
40+
await writeFile(sessionPath, JSON.stringify({
41+
sessionId: 'gemini-session-1',
42+
startTime: '2026-05-16T10:00:00.000Z',
43+
messages: [
44+
{ id: 'u1', timestamp: '2026-05-16T10:00:00.000Z', type: 'user', content: 'work' },
45+
{
46+
id: 'g1',
47+
timestamp: '2026-05-16T10:00:05.000Z',
48+
type: 'gemini',
49+
content: 'first',
50+
model: 'gemini-3.1-pro-preview',
51+
tokens: { input: 10, output: 5 },
52+
},
53+
{
54+
id: 'g2',
55+
timestamp: '2026-05-16T10:00:10.000Z',
56+
type: 'gemini',
57+
content: 'second',
58+
model: 'gemini-3.1-pro-preview',
59+
tokens: { input: 12, output: 6 },
60+
},
61+
],
62+
}))
63+
64+
const fileStat = await stat(sessionPath)
65+
await writeFile(join(cacheDir, 'session-cache.json'), JSON.stringify({
66+
version: CACHE_VERSION,
67+
providers: {
68+
gemini: {
69+
envFingerprint: computeEnvFingerprint('gemini'),
70+
files: {
71+
[sessionPath]: {
72+
fingerprint: {
73+
dev: fileStat.dev,
74+
ino: fileStat.ino,
75+
mtimeMs: fileStat.mtimeMs,
76+
sizeBytes: fileStat.size,
77+
},
78+
mcpInventory: [],
79+
turns: [{
80+
timestamp: '2026-05-16T10:00:00.000Z',
81+
sessionId: 'gemini-session-1',
82+
userMessage: 'work',
83+
calls: [{
84+
provider: 'gemini',
85+
model: 'gemini-3.1-pro-preview',
86+
usage: {
87+
inputTokens: 22,
88+
outputTokens: 11,
89+
cacheCreationInputTokens: 0,
90+
cacheReadInputTokens: 0,
91+
cachedInputTokens: 0,
92+
reasoningTokens: 0,
93+
webSearchRequests: 0,
94+
cacheCreationOneHourTokens: 0,
95+
},
96+
speed: 'standard',
97+
timestamp: '2026-05-16T10:00:00.000Z',
98+
tools: [],
99+
bashCommands: [],
100+
skills: [],
101+
deduplicationKey: 'gemini:gemini-session-1',
102+
}],
103+
}],
104+
},
105+
},
106+
},
107+
},
108+
}))
109+
110+
const range: DateRange = {
111+
start: new Date('2026-05-16T00:00:00.000Z'),
112+
end: new Date('2026-05-16T23:59:59.999Z'),
113+
}
114+
115+
const projects = await parseAllSessions(range, 'gemini')
116+
const keys = projects.flatMap(project =>
117+
project.sessions.flatMap(session =>
118+
session.turns.flatMap(turn => turn.assistantCalls.map(call => call.deduplicationKey)),
119+
),
120+
)
121+
122+
expect(projects[0]!.totalApiCalls).toBe(2)
123+
expect(keys).toEqual([
124+
'gemini:gemini-session-1:g1',
125+
'gemini:gemini-session-1:g2',
126+
])
127+
128+
const savedCache = JSON.parse(await readFile(join(cacheDir, 'session-cache.json'), 'utf-8'))
129+
const savedKeys = savedCache.providers.gemini.files[sessionPath].turns.flatMap((turn: { calls: Array<{ deduplicationKey: string }> }) =>
130+
turn.calls.map(call => call.deduplicationKey),
131+
)
132+
expect(savedKeys).toEqual(keys)
133+
})
134+
})

0 commit comments

Comments
 (0)