From 1f6bd0d2fc894d5672a955946a0e8b9ac01c01e5 Mon Sep 17 00:00:00 2001 From: Jacques Verre Date: Fri, 12 Jun 2026 12:55:41 +0100 Subject: [PATCH] fix hook flush overhead --- src/count_tokens.go | 4 ++++ src/count_tokens_test.go | 21 +++++++++++++++++++++ src/dryrun_test.go | 9 +++++++++ src/main.go | 23 +++++++++++++++++------ src/metrics.go | 17 ++++++----------- 5 files changed, 57 insertions(+), 17 deletions(-) diff --git a/src/count_tokens.go b/src/count_tokens.go index badb7c8..e691100 100644 --- a/src/count_tokens.go +++ b/src/count_tokens.go @@ -3,6 +3,7 @@ package main import ( "bytes" "encoding/json" + "fmt" "net/http" "os" "os/exec" @@ -238,6 +239,9 @@ var countTokensHTTP = func(payload []byte, headers map[string]string) (int, erro return 0, err } defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return 0, fmt.Errorf("count_tokens: %s", resp.Status) + } var out struct { InputTokens int `json:"input_tokens"` } diff --git a/src/count_tokens_test.go b/src/count_tokens_test.go index 0ebd0f5..467ad29 100644 --- a/src/count_tokens_test.go +++ b/src/count_tokens_test.go @@ -2,6 +2,7 @@ package main import ( "encoding/json" + "errors" "os" "strings" "testing" @@ -110,6 +111,26 @@ func TestRunTokenCountPassBudgetAndBaseline(t *testing.T) { _ = spent } +func TestRunTokenCountPassDoesNotCacheHTTPError(t *testing.T) { + resetTokenCache(t) + t.Setenv("ANTHROPIC_API_KEY", "test-key") + + old := countTokensHTTP + countTokensHTTP = func(payload []byte, headers map[string]string) (int, error) { + return 0, errors.New("count_tokens: 401 Unauthorized") + } + defer func() { countTokensHTTP = old }() + + text := strings.Repeat("uncached prompt body\n", 20) + entries := []TranscriptEntry{userPromptEntry(text)} + if spent := runTokenCountPass(entries, 2); spent != 0 { + t.Fatalf("spent = %d, want 0 after HTTP error", spent) + } + if _, ok := tokenCacheGet(defaultCountModel + "|" + sha256hex(text)); ok { + t.Error("failed count_tokens response should not populate cache") + } +} + func jsonReadFile(path string) ([]byte, error) { return os.ReadFile(path) } diff --git a/src/dryrun_test.go b/src/dryrun_test.go index cd265de..3fac19f 100644 --- a/src/dryrun_test.go +++ b/src/dryrun_test.go @@ -120,6 +120,15 @@ func TestTraceNameResolution(t *testing.T) { } } +func TestSpansHaveUsage(t *testing.T) { + if spansHaveUsage([]Span{{Name: "Read"}, {Name: "Edit"}}) { + t.Fatal("tool-only spans should not require context snapshot work") + } + if !spansHaveUsage([]Span{{Name: "Thinking", Usage: map[string]int{"total_tokens": 12}}}) { + t.Fatal("LLM spans with usage should require context snapshot work") + } +} + // TestToolResultDebug enumerates every tool_use → tool_result pair and // flags any tool_use whose result the extractor isn't seeing. func TestToolResultDebug(t *testing.T) { diff --git a/src/main.go b/src/main.go index da8e7bf..ce65fac 100644 --- a/src/main.go +++ b/src/main.go @@ -622,13 +622,15 @@ func flush(state *State) { // billed tokens to categories with a single-row query, no JOIN back // to trace.cc.context_runtime. See context_snapshot.go for the // accuracy tradeoff. - if snapshot := buildContextSnapshot(state); snapshot != nil { - for i := range spans { - if spans[i].Usage == nil { - continue + if spansHaveUsage(spans) { + if snapshot := buildContextSnapshot(state); snapshot != nil { + for i := range spans { + if spans[i].Usage == nil { + continue + } + cc := ensureCCMap(&spans[i]) + cc["context_snapshot"] = snapshot } - cc := ensureCCMap(&spans[i]) - cc["context_snapshot"] = snapshot } } @@ -638,6 +640,15 @@ func flush(state *State) { } } +func spansHaveUsage(spans []Span) bool { + for _, span := range spans { + if span.Usage != nil { + return true + } + } + return false +} + // findSlug returns the best per-session identifier available on the // transcript. Historic shape: per-entry `slug` (session-stable kebab-case). // Claude Code 2.1.150+ shape: dedicated `type:"ai-title"` events carrying diff --git a/src/metrics.go b/src/metrics.go index ae27012..ee427ff 100644 --- a/src/metrics.go +++ b/src/metrics.go @@ -15,15 +15,10 @@ type EditAggregate struct { LinesOverwritten int } -// aggregateEdits walks transcript entries from state.StartLine forward and -// returns counts for Edit/Write/MultiEdit tool calls Claude made in this trace. -func aggregateEdits(state *State) *EditAggregate { +// aggregateEdits walks transcript entries and returns counts for +// Edit/Write/MultiEdit tool calls Claude made in this trace. +func aggregateEdits(entries []TranscriptEntry) *EditAggregate { agg := &EditAggregate{Files: map[string]struct{}{}} - entries, err := ReadTranscript(state.Transcript, state.StartLine) - if err != nil { - return agg - } - for _, entry := range entries { if entry.Type != "assistant" || entry.Message == nil { continue @@ -184,12 +179,14 @@ func postTraceMetrics(state *State) { } metrics := map[string]interface{}{} + fullEntries, _ := ReadTranscript(state.Transcript, 0) + turnEntries, _ := ReadTranscript(state.Transcript, state.StartLine) var repo, branch string var commits, insC, delC int var agg *EditAggregate if cwd != "" && git(cwd, "rev-parse", "--is-inside-work-tree") == "true" { - agg = aggregateEdits(state) + agg = aggregateEdits(turnEntries) repo = repoName(cwd) branch = git(cwd, "branch", "--show-current") @@ -224,8 +221,6 @@ func postTraceMetrics(state *State) { debugLog("postTraceMetrics: skipping git block (cwd=%q not a git work tree)", cwd) } - fullEntries, _ := ReadTranscript(state.Transcript, 0) - turnEntries, _ := ReadTranscript(state.Transcript, state.StartLine) for domain, snap := range domainSnapshotsFromEntries(fullEntries, turnEntries) { if snap != nil { metrics[domain] = snap