Skip to content

Commit fa7e8a7

Browse files
abhisekabhisek
authored
feat: Token usage collection and cost estimation (#25)
* feat: Add support for cost collection * feat: Add cost tab in stats * fix: TUI stats component * fix: code review fixes * fix: Linter fixes * fix: Linter fixes * fix: Linter fixes * fix: Linter fixes
1 parent 0a73534 commit fa7e8a7

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

59 files changed

+21867
-111
lines changed

‎Makefile‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ generate:
2525
generate-schema:
2626
$(GO) run ./cmd/jsonschema-gen
2727

28+
# Update bundled pricing data from models.dev
29+
update-pricing:
30+
$(GO) run pricing/scripts/update-pricing.go
31+
2832
# Verify Event JSON Schema is up-to-date
2933
verify-schema: generate-schema
3034
@git diff --exit-code schema/event.schema.json || (echo "ERROR: schema/event.schema.json is out of date. Run 'make generate-schema' and commit the result." && exit 1)

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,7 @@ Gryph installs lightweight hooks into your AI agents. When the agent performs an
324324
| Use Case | Description |
325325
| --- | --- |
326326
| [AI Coding Observability](examples/ai-coding-observability/) | Centralized observability for AI coding agents across developer endpoints using Gryph + OpenSearch. Includes SOC dashboards, threat detection alerts, and synthetic data generation. |
327+
| Cost & Token Tracking | Track per-session token usage and estimated costs across models and agents. Group by model, agent, or day. See [docs/cost.md](docs/cost.md). |
327328

328329
## Community
329330

‎agent/claudecode/parser.go‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -102,26 +102,36 @@ func (a *Adapter) parseHookEvent(hookType string, rawData []byte) (*events.Event
102102

103103
agentSessionID := baseInput.SessionID
104104

105+
var event *events.Event
106+
var parseErr error
107+
105108
switch eventName {
106109
case "PreToolUse":
107-
return a.parsePreToolUse(sessionID, agentSessionID, baseInput, rawData)
110+
event, parseErr = a.parsePreToolUse(sessionID, agentSessionID, baseInput, rawData)
108111
case "PostToolUse":
109-
return a.parsePostToolUse(sessionID, agentSessionID, baseInput, rawData, false)
112+
event, parseErr = a.parsePostToolUse(sessionID, agentSessionID, baseInput, rawData, false)
110113
case "PostToolUseFailure":
111-
return a.parsePostToolUse(sessionID, agentSessionID, baseInput, rawData, true)
114+
event, parseErr = a.parsePostToolUse(sessionID, agentSessionID, baseInput, rawData, true)
112115
case "SessionStart":
113-
return parseSessionStart(sessionID, agentSessionID, baseInput, rawData)
116+
event, parseErr = parseSessionStart(sessionID, agentSessionID, baseInput, rawData)
114117
case "SessionEnd":
115-
return parseSessionEnd(sessionID, agentSessionID, baseInput, rawData)
118+
event, parseErr = parseSessionEnd(sessionID, agentSessionID, baseInput, rawData)
116119
case "Notification":
117-
return parseNotification(sessionID, agentSessionID, baseInput, rawData)
120+
event, parseErr = parseNotification(sessionID, agentSessionID, baseInput, rawData)
118121
default:
119-
event := events.NewEvent(sessionID, AgentName, events.ActionUnknown)
122+
event = events.NewEvent(sessionID, AgentName, events.ActionUnknown)
120123
event.AgentSessionID = agentSessionID
121124
event.WorkingDirectory = baseInput.Cwd
122125
event.RawEvent = rawData
123-
return event, nil
124126
}
127+
128+
if parseErr != nil {
129+
return nil, parseErr
130+
}
131+
if event != nil {
132+
event.TranscriptPath = baseInput.TranscriptPath
133+
}
134+
return event, nil
125135
}
126136

127137
func (a *Adapter) parsePreToolUse(sessionID uuid.UUID, agentSessionID string, base HookInput, rawData []byte) (*events.Event, error) {

‎agent/claudecode/testdata/transcript_empty.jsonl‎

Whitespace-only changes.

‎agent/claudecode/testdata/transcript_malformed.jsonl‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
{"type":"assistant","message":{"model":"claude-sonnet-4","usage":{"input_tokens":100,"output_tokens":50,"cache_read_input_tokens":10,"cache_creation_input_tokens":5}}}
2+
this is not valid json
3+
{"type":"assistant","message":{"model":"claude-sonnet-4","usage":{"input_tokens":200,"output_tokens":100,"cache_read_input_tokens":20,"cache_creation_input_tokens":10}}}

‎agent/claudecode/testdata/transcript_multi_model.jsonl‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{"type":"human","message":"hello"}
2+
{"type":"assistant","message":{"model":"claude-sonnet-4-20250514","usage":{"input_tokens":1000,"output_tokens":500,"cache_read_input_tokens":200,"cache_creation_input_tokens":100}}}
3+
{"type":"human","message":"do something"}
4+
{"type":"assistant","message":{"model":"claude-sonnet-4-20250514","usage":{"input_tokens":2000,"output_tokens":800,"cache_read_input_tokens":300,"cache_creation_input_tokens":150}}}
5+
{"type":"assistant","message":{"model":"claude-opus-4-6","usage":{"input_tokens":5000,"output_tokens":3000,"cache_read_input_tokens":1000,"cache_creation_input_tokens":500}}}

‎agent/claudecode/testdata/transcript_single_model.jsonl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
{"type":"assistant","message":{"model":"claude-sonnet-4-20250514","usage":{"input_tokens":1500,"output_tokens":700,"cache_read_input_tokens":400,"cache_creation_input_tokens":200}}}

‎agent/claudecode/transcript.go‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package claudecode
2+
3+
import (
4+
"bufio"
5+
"context"
6+
"encoding/json"
7+
"errors"
8+
"io/fs"
9+
"os"
10+
11+
"github.com/safedep/dry/log"
12+
"github.com/safedep/gryph/core/cost"
13+
)
14+
15+
type transcriptMessage struct {
16+
Type string `json:"type"`
17+
Message *messageEnvelope `json:"message"`
18+
}
19+
20+
type messageEnvelope struct {
21+
Model string `json:"model"`
22+
Usage *messageUsage `json:"usage"`
23+
}
24+
25+
type messageUsage struct {
26+
InputTokens int64 `json:"input_tokens"`
27+
OutputTokens int64 `json:"output_tokens"`
28+
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
29+
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
30+
}
31+
32+
// TranscriptCollector implements cost.TokenCollector by parsing Claude Code transcript files.
33+
type TranscriptCollector struct{}
34+
35+
func NewTranscriptCollector() *TranscriptCollector {
36+
return &TranscriptCollector{}
37+
}
38+
39+
func (c *TranscriptCollector) Source() cost.CostSource {
40+
return cost.CostSourceTranscript
41+
}
42+
43+
func (c *TranscriptCollector) Collect(_ context.Context, transcriptPath string) (*cost.SessionUsage, error) {
44+
if transcriptPath == "" {
45+
return nil, nil
46+
}
47+
48+
f, err := os.Open(transcriptPath)
49+
if err != nil {
50+
if errors.Is(err, fs.ErrNotExist) {
51+
return nil, nil
52+
}
53+
return nil, err
54+
}
55+
defer func() { _ = f.Close() }()
56+
57+
usageByModel := make(map[string]*cost.ModelUsage)
58+
59+
scanner := bufio.NewScanner(f)
60+
scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024)
61+
62+
for scanner.Scan() {
63+
line := scanner.Bytes()
64+
if len(line) == 0 {
65+
continue
66+
}
67+
68+
var msg transcriptMessage
69+
if err := json.Unmarshal(line, &msg); err != nil {
70+
log.Debugf("skipping malformed transcript line: %v", err)
71+
continue
72+
}
73+
74+
if msg.Type != "assistant" || msg.Message == nil || msg.Message.Usage == nil {
75+
continue
76+
}
77+
78+
model := msg.Message.Model
79+
if model == "" {
80+
model = "unknown"
81+
}
82+
83+
mu, ok := usageByModel[model]
84+
if !ok {
85+
mu = &cost.ModelUsage{Model: model}
86+
usageByModel[model] = mu
87+
}
88+
89+
mu.InputTokens += msg.Message.Usage.InputTokens
90+
mu.OutputTokens += msg.Message.Usage.OutputTokens
91+
mu.CacheReadTokens += msg.Message.Usage.CacheReadInputTokens
92+
mu.CacheWriteTokens += msg.Message.Usage.CacheCreationInputTokens
93+
}
94+
95+
if err := scanner.Err(); err != nil {
96+
return nil, err
97+
}
98+
99+
if len(usageByModel) == 0 {
100+
return nil, nil
101+
}
102+
103+
usage := &cost.SessionUsage{}
104+
for _, mu := range usageByModel {
105+
usage.Models = append(usage.Models, *mu)
106+
}
107+
usage.Aggregate()
108+
109+
return usage, nil
110+
}

‎agent/claudecode/transcript_test.go‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
package claudecode
2+
3+
import (
4+
"context"
5+
"path/filepath"
6+
"testing"
7+
)
8+
9+
func TestTranscriptCollector_MultiModel(t *testing.T) {
10+
c := NewTranscriptCollector()
11+
usage, err := c.Collect(context.Background(), filepath.Join("testdata", "transcript_multi_model.jsonl"))
12+
if err != nil {
13+
t.Fatalf("unexpected error: %v", err)
14+
}
15+
if usage == nil {
16+
t.Fatal("expected usage, got nil")
17+
}
18+
19+
if len(usage.Models) != 2 {
20+
t.Fatalf("expected 2 models, got %d", len(usage.Models))
21+
}
22+
23+
// Check aggregate totals
24+
// sonnet: input=3000, output=1300, cache_read=500, cache_write=250
25+
// opus: input=5000, output=3000, cache_read=1000, cache_write=500
26+
if usage.InputTokens != 8000 {
27+
t.Errorf("expected input_tokens=8000, got %d", usage.InputTokens)
28+
}
29+
if usage.OutputTokens != 4300 {
30+
t.Errorf("expected output_tokens=4300, got %d", usage.OutputTokens)
31+
}
32+
if usage.CacheReadTokens != 1500 {
33+
t.Errorf("expected cache_read_tokens=1500, got %d", usage.CacheReadTokens)
34+
}
35+
if usage.CacheWriteTokens != 750 {
36+
t.Errorf("expected cache_write_tokens=750, got %d", usage.CacheWriteTokens)
37+
}
38+
}
39+
40+
func TestTranscriptCollector_SingleModel(t *testing.T) {
41+
c := NewTranscriptCollector()
42+
usage, err := c.Collect(context.Background(), filepath.Join("testdata", "transcript_single_model.jsonl"))
43+
if err != nil {
44+
t.Fatalf("unexpected error: %v", err)
45+
}
46+
if usage == nil {
47+
t.Fatal("expected usage, got nil")
48+
}
49+
if len(usage.Models) != 1 {
50+
t.Fatalf("expected 1 model, got %d", len(usage.Models))
51+
}
52+
if usage.Models[0].Model != "claude-sonnet-4-20250514" {
53+
t.Errorf("expected model claude-sonnet-4-20250514, got %s", usage.Models[0].Model)
54+
}
55+
if usage.InputTokens != 1500 {
56+
t.Errorf("expected input_tokens=1500, got %d", usage.InputTokens)
57+
}
58+
}
59+
60+
func TestTranscriptCollector_EmptyFile(t *testing.T) {
61+
c := NewTranscriptCollector()
62+
usage, err := c.Collect(context.Background(), filepath.Join("testdata", "transcript_empty.jsonl"))
63+
if err != nil {
64+
t.Fatalf("unexpected error: %v", err)
65+
}
66+
if usage != nil {
67+
t.Errorf("expected nil usage for empty file, got %+v", usage)
68+
}
69+
}
70+
71+
func TestTranscriptCollector_MissingFile(t *testing.T) {
72+
c := NewTranscriptCollector()
73+
usage, err := c.Collect(context.Background(), "/nonexistent/path.jsonl")
74+
if err != nil {
75+
t.Fatalf("unexpected error for missing file: %v", err)
76+
}
77+
if usage != nil {
78+
t.Errorf("expected nil usage for missing file, got %+v", usage)
79+
}
80+
}
81+
82+
func TestTranscriptCollector_EmptyPath(t *testing.T) {
83+
c := NewTranscriptCollector()
84+
usage, err := c.Collect(context.Background(), "")
85+
if err != nil {
86+
t.Fatalf("unexpected error: %v", err)
87+
}
88+
if usage != nil {
89+
t.Errorf("expected nil usage for empty path, got %+v", usage)
90+
}
91+
}
92+
93+
func TestTranscriptCollector_MalformedLines(t *testing.T) {
94+
c := NewTranscriptCollector()
95+
usage, err := c.Collect(context.Background(), filepath.Join("testdata", "transcript_malformed.jsonl"))
96+
if err != nil {
97+
t.Fatalf("unexpected error: %v", err)
98+
}
99+
if usage == nil {
100+
t.Fatal("expected usage (valid lines should be parsed)")
101+
}
102+
// Should have parsed 2 valid assistant messages, skipping the malformed line
103+
if usage.InputTokens != 300 {
104+
t.Errorf("expected input_tokens=300, got %d", usage.InputTokens)
105+
}
106+
}
107+
108+
func TestTranscriptCollector_Source(t *testing.T) {
109+
c := NewTranscriptCollector()
110+
if c.Source() != "transcript" {
111+
t.Errorf("expected source 'transcript', got '%s'", c.Source())
112+
}
113+
}

‎agent/cursor/parser.go‎

Lines changed: 30 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -220,57 +220,67 @@ func (a *Adapter) parseHookEvent(hookType string, rawData []byte) (*events.Event
220220

221221
agentSessionID := baseInput.ConversationID
222222

223+
var event *events.Event
224+
var parseErr error
225+
223226
switch hookType {
224227
case "preToolUse":
225-
return a.parsePreToolUse(sessionID, agentSessionID, baseInput, rawData)
228+
event, parseErr = a.parsePreToolUse(sessionID, agentSessionID, baseInput, rawData)
226229
case "postToolUse":
227-
return a.parsePostToolUse(sessionID, agentSessionID, baseInput, rawData)
230+
event, parseErr = a.parsePostToolUse(sessionID, agentSessionID, baseInput, rawData)
228231
case "postToolUseFailure":
229-
return a.parsePostToolUseFailure(sessionID, agentSessionID, baseInput, rawData)
232+
event, parseErr = a.parsePostToolUseFailure(sessionID, agentSessionID, baseInput, rawData)
230233
case "beforeShellExecution":
231-
return parseBeforeShellExecution(sessionID, agentSessionID, baseInput, rawData)
234+
event, parseErr = parseBeforeShellExecution(sessionID, agentSessionID, baseInput, rawData)
232235
case "beforeReadFile":
233-
return a.parseBeforeReadFile(sessionID, agentSessionID, baseInput, rawData)
236+
event, parseErr = a.parseBeforeReadFile(sessionID, agentSessionID, baseInput, rawData)
234237
case "afterFileEdit":
235-
return a.parseAfterFileEdit(sessionID, agentSessionID, baseInput, rawData)
238+
event, parseErr = a.parseAfterFileEdit(sessionID, agentSessionID, baseInput, rawData)
236239
case "beforeSubmitPrompt":
237-
return parseBeforeSubmitPrompt(sessionID, agentSessionID, baseInput, rawData)
240+
event, parseErr = parseBeforeSubmitPrompt(sessionID, agentSessionID, baseInput, rawData)
238241
case "sessionStart":
239-
return parseSessionStart(sessionID, agentSessionID, baseInput, rawData)
242+
event, parseErr = parseSessionStart(sessionID, agentSessionID, baseInput, rawData)
240243
case "sessionEnd":
241-
return parseSessionEnd(sessionID, agentSessionID, baseInput, rawData)
244+
event, parseErr = parseSessionEnd(sessionID, agentSessionID, baseInput, rawData)
242245
case "stop":
243-
return parseStop(sessionID, agentSessionID, baseInput, rawData)
246+
event, parseErr = parseStop(sessionID, agentSessionID, baseInput, rawData)
244247
case "beforeTabFileRead":
245-
return a.parseBeforeReadFile(sessionID, agentSessionID, baseInput, rawData)
248+
event, parseErr = a.parseBeforeReadFile(sessionID, agentSessionID, baseInput, rawData)
246249
case "afterTabFileEdit":
247-
return a.parseAfterFileEdit(sessionID, agentSessionID, baseInput, rawData)
250+
event, parseErr = a.parseAfterFileEdit(sessionID, agentSessionID, baseInput, rawData)
248251
case "beforeMCPExecution":
249-
return parseBeforeMCPExecution(sessionID, agentSessionID, baseInput, rawData)
252+
event, parseErr = parseBeforeMCPExecution(sessionID, agentSessionID, baseInput, rawData)
250253
case "afterShellExecution":
251-
return parseAfterShellExecution(sessionID, agentSessionID, baseInput, rawData)
254+
event, parseErr = parseAfterShellExecution(sessionID, agentSessionID, baseInput, rawData)
252255
case "afterMCPExecution":
253-
return parseAfterMCPExecution(sessionID, agentSessionID, baseInput, rawData)
256+
event, parseErr = parseAfterMCPExecution(sessionID, agentSessionID, baseInput, rawData)
254257
case "subagentStart":
255-
return parseSubagentStart(sessionID, agentSessionID, baseInput, rawData)
258+
event, parseErr = parseSubagentStart(sessionID, agentSessionID, baseInput, rawData)
256259
case "subagentStop":
257-
return parseSubagentStop(sessionID, agentSessionID, baseInput, rawData)
260+
event, parseErr = parseSubagentStop(sessionID, agentSessionID, baseInput, rawData)
258261
case "afterAgentThought":
259-
return parseAfterAgentThought(sessionID, agentSessionID, baseInput, rawData)
262+
event, parseErr = parseAfterAgentThought(sessionID, agentSessionID, baseInput, rawData)
260263
default:
261264
actionType := events.ActionUnknown
262265
if at, ok := HookTypeMapping[hookType]; ok {
263266
actionType = at
264267
}
265-
event := events.NewEvent(sessionID, AgentName, actionType)
268+
event = events.NewEvent(sessionID, AgentName, actionType)
266269
event.AgentSessionID = agentSessionID
267270
event.ToolName = hookType
268271
event.RawEvent = rawData
269272
if len(baseInput.WorkspaceRoots) > 0 {
270273
event.WorkingDirectory = baseInput.WorkspaceRoots[0]
271274
}
272-
return event, nil
273275
}
276+
277+
if parseErr != nil {
278+
return nil, parseErr
279+
}
280+
if event != nil {
281+
event.TranscriptPath = baseInput.TranscriptPath
282+
}
283+
return event, nil
274284
}
275285

276286
func (a *Adapter) parsePreToolUse(sessionID uuid.UUID, agentSessionID string, base HookInput, rawData []byte) (*events.Event, error) {

0 commit comments

Comments
 (0)