Merge pull request #26 from comet-ml/codex/hook-performance-cleanup

[codex] Fix hook flush overhead

Author: jverre

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"`
 	}

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)
 }

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) {

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

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