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{