From fb0793767d8413223e85c459eae117e33b5b28e1 Mon Sep 17 00:00:00 2001 From: Jordan Taylor Date: Wed, 18 Mar 2026 22:03:16 -0500 Subject: [PATCH] feat: extract token usage from Claude stream-json events The Claude agent runner was not populating the Usage field on CodexEvents, so token counts were always zero in the dashboard and API. Add extractClaudeUsage() to parse token data from both shapes Claude emits: - result events: top-level "usage" field - assistant events: nested "message.usage" field Wire the extracted usage into the CodexEvent forwarded to the orchestrator, matching the existing behaviour in the Codex runner. Also add tests for extractClaudeUsage and dashboard token count rendering. Co-Authored-By: Claude Sonnet 4.6 --- internal/agent/claude.go | 44 +++++++++++++++++++ internal/agent/claude_test.go | 79 ++++++++++++++++++++++++++++++++++ internal/server/server_test.go | 56 ++++++++++++++++++++++++ 3 files changed, 179 insertions(+) create mode 100644 internal/agent/claude_test.go diff --git a/internal/agent/claude.go b/internal/agent/claude.go index 689037a..d7fe56c 100644 --- a/internal/agent/claude.go +++ b/internal/agent/claude.go @@ -214,6 +214,7 @@ func (s *ClaudeSession) RunTurn(ctx context.Context, prompt string, issue model. Timestamp: time.Now().UTC(), Message: extractClaudeMessage(event), SessionID: sessionID, + Usage: extractClaudeUsage([]byte(line)), } onEvent(ce) } @@ -328,3 +329,46 @@ func shellescape(s string) string { // Wrap in single quotes and escape any single quotes within return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'" } + +// extractClaudeUsage parses token usage from a raw Claude stream-json line. +// Claude emits usage in two shapes: +// +// result events: {"type":"result", "usage":{"input_tokens":N,"output_tokens":N,...}} +// assistant events: {"type":"assistant","message":{"usage":{"input_tokens":N,"output_tokens":N,...}}} +func extractClaudeUsage(raw []byte) *model.TokenUsage { + // Use a flexible struct that captures both shapes without conflicting tags. + var envelope struct { + Usage *struct { + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + } `json:"usage"` + Message *struct { + Usage *struct { + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + } `json:"usage"` + } `json:"message"` + } + if err := json.Unmarshal(raw, &envelope); err != nil { + return nil + } + + var input, output int64 + + if envelope.Usage != nil { + input = envelope.Usage.InputTokens + output = envelope.Usage.OutputTokens + } else if envelope.Message != nil && envelope.Message.Usage != nil { + input = envelope.Message.Usage.InputTokens + output = envelope.Message.Usage.OutputTokens + } + + if input == 0 && output == 0 { + return nil + } + return &model.TokenUsage{ + InputTokens: input, + OutputTokens: output, + TotalTokens: input + output, + } +} diff --git a/internal/agent/claude_test.go b/internal/agent/claude_test.go new file mode 100644 index 0000000..d5724f9 --- /dev/null +++ b/internal/agent/claude_test.go @@ -0,0 +1,79 @@ +package agent + +import ( + "testing" +) + +func TestExtractClaudeUsage_ResultEvent(t *testing.T) { + raw := []byte(`{"type":"result","subtype":"success","is_error":false,"result":"done","session_id":"abc","usage":{"input_tokens":1000,"output_tokens":500}}`) + u := extractClaudeUsage(raw) + if u == nil { + t.Fatal("expected usage, got nil") + } + if u.InputTokens != 1000 { + t.Errorf("InputTokens: want 1000, got %d", u.InputTokens) + } + if u.OutputTokens != 500 { + t.Errorf("OutputTokens: want 500, got %d", u.OutputTokens) + } + if u.TotalTokens != 1500 { + t.Errorf("TotalTokens: want 1500, got %d", u.TotalTokens) + } +} + +func TestExtractClaudeUsage_AssistantEvent(t *testing.T) { + raw := []byte(`{"type":"assistant","message":{"id":"msg_01","role":"assistant","content":[],"usage":{"input_tokens":800,"output_tokens":200}}}`) + u := extractClaudeUsage(raw) + if u == nil { + t.Fatal("expected usage, got nil") + } + if u.InputTokens != 800 { + t.Errorf("InputTokens: want 800, got %d", u.InputTokens) + } + if u.OutputTokens != 200 { + t.Errorf("OutputTokens: want 200, got %d", u.OutputTokens) + } + if u.TotalTokens != 1000 { + t.Errorf("TotalTokens: want 1000, got %d", u.TotalTokens) + } +} + +func TestExtractClaudeUsage_NoUsage(t *testing.T) { + raw := []byte(`{"type":"system","subtype":"init","session_id":"abc"}`) + u := extractClaudeUsage(raw) + if u != nil { + t.Errorf("expected nil usage for system event, got %+v", u) + } +} + +func TestExtractClaudeUsage_ZeroTokens(t *testing.T) { + raw := []byte(`{"type":"result","usage":{"input_tokens":0,"output_tokens":0}}`) + u := extractClaudeUsage(raw) + if u != nil { + t.Errorf("expected nil for zero-token usage, got %+v", u) + } +} + +func TestExtractClaudeUsage_InvalidJSON(t *testing.T) { + u := extractClaudeUsage([]byte(`not json`)) + if u != nil { + t.Errorf("expected nil for invalid JSON, got %+v", u) + } +} + +func TestExtractClaudeUsage_WithCacheTokens(t *testing.T) { + raw := []byte(`{"type":"result","usage":{"input_tokens":600,"output_tokens":300,"cache_creation_input_tokens":100,"cache_read_input_tokens":50}}`) + u := extractClaudeUsage(raw) + if u == nil { + t.Fatal("expected usage, got nil") + } + if u.InputTokens != 600 { + t.Errorf("InputTokens: want 600, got %d", u.InputTokens) + } + if u.OutputTokens != 300 { + t.Errorf("OutputTokens: want 300, got %d", u.OutputTokens) + } + if u.TotalTokens != 900 { + t.Errorf("TotalTokens: want 900, got %d", u.TotalTokens) + } +} diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 6e1ffe2..862bbf7 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -9,6 +9,7 @@ import ( "testing" "time" + "github.com/jordan/go-symphony/internal/model" "github.com/jordan/go-symphony/internal/orchestrator" ) @@ -26,6 +27,61 @@ func newTestServer(snap orchestrator.StateSnapshot) *httptest.Server { return httptest.NewServer(mux) } +func TestDashboardShowsTokenCounts(t *testing.T) { + now := time.Now() + snap := orchestrator.StateSnapshot{ + GeneratedAt: now, + Running: []orchestrator.RunningSnapshot{ + { + IssueID: "id-1", + IssueIdentifier: "ZYX-42", + IssueTitle: "Some issue", + State: "In Progress", + StartedAt: now, + Tokens: orchestrator.TokenSnapshot{ + InputTokens: 1234, + OutputTokens: 567, + TotalTokens: 1801, + }, + }, + }, + CodexTotals: model.CodexTotals{ + InputTokens: 5000, + OutputTokens: 2000, + TotalTokens: 7000, + }, + } + + ts := newTestServer(snap) + defer ts.Close() + + resp, err := http.Get(ts.URL + "/") //nolint:noctx // test-only HTTP call + if err != nil { + t.Fatalf("GET /: %v", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + html := string(body) + + // Aggregate totals in header + for _, want := range []string{"5000", "2000", "7000"} { + if !strings.Contains(html, want) { + t.Errorf("dashboard missing aggregate token count %q", want) + } + } + + // Per-session tokens in running table + for _, want := range []string{"1234", "567", "1801"} { + if !strings.Contains(html, want) { + t.Errorf("dashboard missing per-session token count %q", want) + } + } +} + func TestDashboardShowsIssueTitle(t *testing.T) { now := time.Now() snap := orchestrator.StateSnapshot{