From 54520de9577ea4fe1bbcc4cfcc3755cae643ae93 Mon Sep 17 00:00:00 2001 From: Multica Eve Date: Thu, 24 Sep 2026 13:44:51 +0800 Subject: [PATCH] MUL-7646: fix(claude): report per-run usage for resumed sessions (#8792) * fix(claude): report per-run usage for resumed sessions Co-authored-by: multica-agent * fix(claude): preserve usage when counters reset Co-authored-by: multica-agent --------- Co-authored-by: Sol-Boy Co-authored-by: multica-agent --- server/pkg/agent/claude.go | 317 +++++++++++++++++++++++-- server/pkg/agent/claude_usage_test.go | 318 +++++++++++++++++++++++++- 2 files changed, 617 insertions(+), 18 deletions(-) diff --git a/server/pkg/agent/claude.go b/server/pkg/agent/claude.go index 59be6a925..347501b2c 100644 --- a/server/pkg/agent/claude.go +++ b/server/pkg/agent/claude.go @@ -1,6 +1,7 @@ package agent import ( + "bytes" "context" "encoding/json" "errors" @@ -93,6 +94,16 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt return nil, err } + var usageSnapshot *claudeUsageSnapshot + if opts.ResumeSessionID != "" { + snapshot, snapshotErr := captureClaudeUsageSnapshot(cmd.Env, cmd.Dir, opts.ResumeSessionID) + if snapshotErr != nil { + b.cfg.Logger.Warn("claude usage baseline unavailable; falling back to reported totals", "error", snapshotErr) + } else { + usageSnapshot = snapshot + } + } + stdout, err := cmd.StdoutPipe() if err != nil { cancel() @@ -188,6 +199,7 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt resultIsError := false terminalReasonError := "" var sessionID string + var lastUsageResult *claudeSDKMessage sawAsyncLaunch := false usage := make(map[string]TokenUsage) seenUsage := make(map[string]struct{}) @@ -270,8 +282,14 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt resultIsError = msg.IsError terminalReasonError = claudeTerminalReasonFailure(msg.TerminalReason, msg.ResultText) sessionID = msg.SessionID - if resultUsage := claudeResultUsage(msg, opts.Model); len(resultUsage) > 0 { + var baseline map[string]TokenUsage + if usageSnapshot != nil { + baseline = usageSnapshot.baseline + } + if resultUsage, authoritative := claudeResultUsageSince(msg, opts.Model, baseline); authoritative { usage = resultUsage + usageMsg := msg + lastUsageResult = &usageMsg } closeStdin() case "log": @@ -343,6 +361,16 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt terminalReasonError = fmt.Sprintf("claude input/control protocol failed: %v", writeErr) } + if !sawResult && usageSnapshot != nil { + appendedUsage, found, snapshotErr := usageSnapshot.appendedCostStateUsage() + switch { + case snapshotErr != nil: + b.cfg.Logger.Warn("claude final cost state unavailable; keeping stream usage fallback", "error", snapshotErr) + case found: + usage = subtractClaudeUsage(appendedUsage, usageSnapshot.baseline) + } + } + completionGuardError := "" if sawAsyncLaunch { completionGuardError = "claude launched an async background task; Multica-managed runs require foreground execution" @@ -399,6 +427,13 @@ func (b *claudeBackend) Execute(ctx context.Context, prompt string, opts ExecOpt resumeRejected := resumeWasRejected(opts.ResumeSessionID, sessionID, finalStatus == "failed", finalError, stderrTail) reportedSessionID := resolveSessionID(opts.ResumeSessionID, sessionID, finalStatus == "failed", finalError, stderrTail) if resumeRejected { + // A rejected resume may emit usage for a newly-created session. Its + // totals do not share the requested session's baseline. + if lastUsageResult != nil { + if resultUsage, authoritative := claudeResultUsageSince(*lastUsageResult, opts.Model, nil); authoritative { + usage = resultUsage + } + } b.cfg.Logger.Info("claude resume was rejected; dropping session id and signalling fresh-session retry", "requested_resume", opts.ResumeSessionID, "emitted_session", sessionID, @@ -697,21 +732,15 @@ func claudeTerminalReasonFailure(terminalReason, resultText string) string { } func claudeResultUsage(msg claudeSDKMessage, fallbackModel string) map[string]TokenUsage { + usage, _ := claudeResultUsageSince(msg, fallbackModel, nil) + return usage +} + +func claudeResultUsageSince(msg claudeSDKMessage, fallbackModel string, baseline map[string]TokenUsage) (map[string]TokenUsage, bool) { if len(msg.ModelUsage) > 0 { - usage := make(map[string]TokenUsage, len(msg.ModelUsage)) - for model, u := range msg.ModelUsage { - if model == "" || !claudeUsageHasTokens(u.InputTokens, u.OutputTokens, u.CacheReadInputTokens, u.CacheCreationInputTokens) { - continue - } - usage[model] = TokenUsage{ - InputTokens: u.InputTokens, - OutputTokens: u.OutputTokens, - CacheReadTokens: u.CacheReadInputTokens, - CacheWriteTokens: u.CacheCreationInputTokens, - } - } + usage := claudeModelUsage(msg.ModelUsage) if len(usage) > 0 { - return usage + return subtractClaudeUsage(usage, baseline), true } } @@ -725,7 +754,7 @@ func claudeResultUsage(msg claudeSDKMessage, fallbackModel string) map[string]To msg.Usage.CacheReadInputTokens, msg.Usage.CacheCreationInputTokens, ) { - return nil + return nil, false } return map[string]TokenUsage{ model: { @@ -734,7 +763,265 @@ func claudeResultUsage(msg claudeSDKMessage, fallbackModel string) map[string]To CacheReadTokens: msg.Usage.CacheReadInputTokens, CacheWriteTokens: msg.Usage.CacheCreationInputTokens, }, + }, true +} + +func claudeModelUsage(modelUsage map[string]claudeResultModelUsage) map[string]TokenUsage { + usage := make(map[string]TokenUsage, len(modelUsage)) + for model, u := range modelUsage { + if model == "" || !claudeUsageHasTokens(u.InputTokens, u.OutputTokens, u.CacheReadInputTokens, u.CacheCreationInputTokens) { + continue + } + usage[model] = TokenUsage{ + InputTokens: u.InputTokens, + OutputTokens: u.OutputTokens, + CacheReadTokens: u.CacheReadInputTokens, + CacheWriteTokens: u.CacheCreationInputTokens, + } } + return usage +} + +func subtractClaudeUsage(current, baseline map[string]TokenUsage) map[string]TokenUsage { + // A baseline is usable only when Claude restored all of it. Older versions, + // downgrades, and future counter resets report per-run totals instead; in + // those cases subtracting even one historical field would silently undercount. + if !claudeUsageIncludesBaseline(current, baseline) { + usage := make(map[string]TokenUsage, len(current)) + for model, modelUsage := range current { + usage[model] = modelUsage + } + return usage + } + usage := make(map[string]TokenUsage, len(current)) + for model, currentUsage := range current { + base := baseline[model] + delta := TokenUsage{ + InputTokens: currentUsage.InputTokens - base.InputTokens, + OutputTokens: currentUsage.OutputTokens - base.OutputTokens, + CacheReadTokens: currentUsage.CacheReadTokens - base.CacheReadTokens, + CacheWriteTokens: currentUsage.CacheWriteTokens - base.CacheWriteTokens, + } + if claudeUsageHasTokens(delta.InputTokens, delta.OutputTokens, delta.CacheReadTokens, delta.CacheWriteTokens) { + usage[model] = delta + } + } + return usage +} + +func claudeUsageIncludesBaseline(current, baseline map[string]TokenUsage) bool { + for model, base := range baseline { + modelUsage, ok := current[model] + if !ok || modelUsage.InputTokens < base.InputTokens || + modelUsage.OutputTokens < base.OutputTokens || + modelUsage.CacheReadTokens < base.CacheReadTokens || + modelUsage.CacheWriteTokens < base.CacheWriteTokens { + return false + } + } + return true +} + +const claudeSessionReadChunkSize = 64 * 1024 + +type claudeUsageSnapshot struct { + path string + fileInfo os.FileInfo + offset int64 + baseline map[string]TokenUsage +} + +func captureClaudeUsageSnapshot(env []string, cwd, sessionID string) (*claudeUsageSnapshot, error) { + configDir, err := claudeConfigDir(env, cwd) + if err != nil { + return nil, err + } + path, err := findClaudeSessionFile(configDir, sessionID) + if err != nil { + return nil, err + } + f, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("open Claude session: %w", err) + } + defer f.Close() + info, err := f.Stat() + if err != nil { + return nil, fmt.Errorf("stat Claude session: %w", err) + } + baseline, _, err := readLastClaudeCostStateUsage(f, 0, info.Size()) + if err != nil { + return nil, fmt.Errorf("read Claude usage baseline: %w", err) + } + return &claudeUsageSnapshot{ + path: path, + fileInfo: info, + offset: info.Size(), + baseline: baseline, + }, nil +} + +func (s *claudeUsageSnapshot) appendedCostStateUsage() (map[string]TokenUsage, bool, error) { + f, err := os.Open(s.path) + if err != nil { + return nil, false, fmt.Errorf("reopen Claude session: %w", err) + } + defer f.Close() + info, err := f.Stat() + if err != nil { + return nil, false, fmt.Errorf("stat Claude session after exit: %w", err) + } + if !os.SameFile(s.fileInfo, info) { + return nil, false, errors.New("Claude session file was replaced during execution") + } + if info.Size() < s.offset { + return nil, false, errors.New("Claude session file shrank during execution") + } + return readLastClaudeCostStateUsage(f, s.offset, info.Size()) +} + +func claudeConfigDir(env []string, cwd string) (string, error) { + if value, ok := lastClaudeEnvValue(env, "CLAUDE_CONFIG_DIR"); ok && value != "" { + return resolveClaudeConfigPath(value, cwd) + } + if runtime.GOOS == "windows" { + if value, ok := lastClaudeEnvValue(env, "USERPROFILE"); ok && value != "" { + return resolveClaudeConfigPath(filepath.Join(value, ".claude"), cwd) + } + } else if value, ok := lastClaudeEnvValue(env, "HOME"); ok && value != "" { + return resolveClaudeConfigPath(filepath.Join(value, ".claude"), cwd) + } + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("resolve home directory for Claude config: %w", err) + } + return resolveClaudeConfigPath(filepath.Join(home, ".claude"), cwd) +} + +func resolveClaudeConfigPath(path, cwd string) (string, error) { + if filepath.IsAbs(path) { + return filepath.Clean(path), nil + } + if cwd != "" { + if !filepath.IsAbs(cwd) { + absoluteCwd, err := filepath.Abs(cwd) + if err != nil { + return "", fmt.Errorf("resolve Claude working directory: %w", err) + } + cwd = absoluteCwd + } + return filepath.Clean(filepath.Join(cwd, path)), nil + } + absolutePath, err := filepath.Abs(path) + if err != nil { + return "", fmt.Errorf("resolve Claude config directory: %w", err) + } + return absolutePath, nil +} + +func lastClaudeEnvValue(env []string, key string) (string, bool) { + for i := len(env) - 1; i >= 0; i-- { + name, value, ok := strings.Cut(env[i], "=") + matches := name == key + if runtime.GOOS == "windows" { + matches = strings.EqualFold(name, key) + } + if ok && matches { + return value, true + } + } + return "", false +} + +func findClaudeSessionFile(configDir, sessionID string) (string, error) { + if sessionID == "" || sessionID == "." || sessionID == ".." || strings.ContainsAny(sessionID, "/\\:\x00") { + return "", fmt.Errorf("invalid Claude session id %q", sessionID) + } + projectsDir := filepath.Join(configDir, "projects") + projects, err := os.ReadDir(projectsDir) + if err != nil { + return "", fmt.Errorf("read Claude projects directory: %w", err) + } + var bestPath string + var bestInfo os.FileInfo + var firstErr error + for _, project := range projects { + path := filepath.Join(projectsDir, project.Name(), sessionID+".jsonl") + info, statErr := os.Stat(path) + if statErr != nil { + if !errors.Is(statErr, os.ErrNotExist) && firstErr == nil { + firstErr = statErr + } + continue + } + if !info.Mode().IsRegular() { + continue + } + if bestInfo == nil || info.ModTime().After(bestInfo.ModTime()) { + bestPath = path + bestInfo = info + } + } + if bestPath != "" { + return bestPath, nil + } + if firstErr != nil { + return "", fmt.Errorf("locate Claude session: %w", firstErr) + } + return "", fmt.Errorf("locate Claude session %q: %w", sessionID, os.ErrNotExist) +} + +func readLastClaudeCostStateUsage(f *os.File, start, end int64) (map[string]TokenUsage, bool, error) { + if start < 0 || end < start { + return nil, false, fmt.Errorf("invalid Claude session range [%d,%d)", start, end) + } + var suffix []byte + for position := end; position > start; { + chunkStart := max(position-claudeSessionReadChunkSize, start) + chunk := make([]byte, position-chunkStart) + if _, err := f.ReadAt(chunk, chunkStart); err != nil && !errors.Is(err, io.EOF) { + return nil, false, fmt.Errorf("read Claude session: %w", err) + } + data := make([]byte, 0, len(chunk)+len(suffix)) + data = append(data, chunk...) + data = append(data, suffix...) + lines := bytes.Split(data, []byte{'\n'}) + firstComplete := chunkStart == start + firstLine := 0 + if !firstComplete { + suffix = append(suffix[:0], lines[0]...) + firstLine = 1 + } + for i := len(lines) - 1; i >= firstLine; i-- { + line := bytes.TrimSpace(lines[i]) + if len(line) == 0 { + continue + } + var envelope struct { + Type string `json:"type"` + } + if err := json.Unmarshal(line, &envelope); err != nil { + continue + } + if envelope.Type == "cost-state" { + var state struct { + ModelUsage map[string]claudeResultModelUsage `json:"modelUsage"` + } + if err := json.Unmarshal(line, &state); err != nil { + continue + } + usage := claudeModelUsage(state.ModelUsage) + if len(usage) == 0 { + // An empty or newly-incompatible shape is not authoritative. + // Keeping the stream fallback is safer than dropping usage. + return nil, false, nil + } + return usage, true, nil + } + } + position = chunkStart + } + return nil, false, nil } func claudeUsageHasTokens(input, output, cacheRead, cacheWrite int64) bool { diff --git a/server/pkg/agent/claude_usage_test.go b/server/pkg/agent/claude_usage_test.go index fe91d06e5..bb9977272 100644 --- a/server/pkg/agent/claude_usage_test.go +++ b/server/pkg/agent/claude_usage_test.go @@ -28,6 +28,27 @@ func runFakeClaudeUsageFixture() { if _, err := os.Stdout.Write(data); err != nil { os.Exit(2) } + if appendPath := os.Getenv("CLAUDE_USAGE_APPEND_FILE"); appendPath != "" { + appendData, err := os.ReadFile(os.Getenv("CLAUDE_USAGE_APPEND_FIXTURE")) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(2) + } + f, err := os.OpenFile(appendPath, os.O_APPEND|os.O_WRONLY, 0) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(2) + } + if _, err := f.Write(appendData); err != nil { + _ = f.Close() + fmt.Fprintln(os.Stderr, err) + os.Exit(2) + } + if err := f.Close(); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(2) + } + } if os.Getenv("CLAUDE_USAGE_FIXTURE_FAIL") == "1" { os.Exit(1) } @@ -131,7 +152,290 @@ func TestClaudeFallbackUsageIsScopedToExecution(t *testing.T) { } } +func TestClaudeExecuteSubtractsResumedCostStateBaseline(t *testing.T) { + const sessionID = "session-cumulative" + const model = "claude-sonnet-4-6" + const otherModel = "claude-haiku-4-5" + configDir := t.TempDir() + writeClaudeUsageSession(t, configDir, sessionID, + claudeCostStateFixture(t, map[string]claudeResultModelUsage{ + model: {InputTokens: 100, OutputTokens: 20, CacheReadInputTokens: 300, CacheCreationInputTokens: 40}, + }), + []byte(`{"type":"user","message":{"role":"user","content":"after baseline"}}`+"\n"), + ) + + resultEvent := mustMarshal(t, map[string]any{ + "type": "result", "subtype": "success", "session_id": sessionID, "result": "done", + "modelUsage": map[string]claudeResultModelUsage{ + model: {InputTokens: 115, OutputTokens: 27, CacheReadInputTokens: 360, CacheCreationInputTokens: 42}, + otherModel: {InputTokens: 3, OutputTokens: 5, CacheReadInputTokens: 11, CacheCreationInputTokens: 2}, + }, + }) + backend := claudeUsageFixtureBackendWithEnv(t, nil, resultEvent, false, map[string]string{ + "CLAUDE_CONFIG_DIR": configDir, + }) + result := executeClaudeUsageFixtureWithOptions(t, backend, 0, ExecOptions{ + Timeout: 5 * time.Second, + ResumeSessionID: sessionID, + }) + want := map[string]TokenUsage{ + model: {InputTokens: 15, OutputTokens: 7, CacheReadTokens: 60, CacheWriteTokens: 2}, + otherModel: {InputTokens: 3, OutputTokens: 5, CacheReadTokens: 11, CacheWriteTokens: 2}, + } + if !reflect.DeepEqual(result.Usage, want) { + t.Fatalf("usage = %#v, want %#v", result.Usage, want) + } +} + +func TestClaudeExecuteKeepsPerRunUsageWithoutCostState(t *testing.T) { + const sessionID = "session-pre-cost-state" + const model = "claude-sonnet-4-6" + configDir := t.TempDir() + writeClaudeUsageSession(t, configDir, sessionID, + []byte(`{"type":"user","message":{"role":"user","content":"old Claude Code session"}}`+"\n"), + ) + resultEvent := mustMarshal(t, map[string]any{ + "type": "result", "subtype": "success", "session_id": sessionID, "result": "done", + "modelUsage": map[string]claudeResultModelUsage{ + model: {InputTokens: 15, OutputTokens: 7, CacheReadInputTokens: 60, CacheCreationInputTokens: 2}, + }, + }) + backend := claudeUsageFixtureBackendWithEnv(t, nil, resultEvent, false, map[string]string{ + "CLAUDE_CONFIG_DIR": configDir, + }) + result := executeClaudeUsageFixtureWithOptions(t, backend, 0, ExecOptions{ + Timeout: 5 * time.Second, + ResumeSessionID: sessionID, + }) + want := map[string]TokenUsage{ + model: {InputTokens: 15, OutputTokens: 7, CacheReadTokens: 60, CacheWriteTokens: 2}, + } + if !reflect.DeepEqual(result.Usage, want) { + t.Fatalf("usage = %#v, want %#v", result.Usage, want) + } +} + +func TestClaudeExecuteFallsBackWhenSessionFileIsMalformed(t *testing.T) { + const sessionID = "session-malformed" + const model = "claude-sonnet-4-6" + configDir := t.TempDir() + writeClaudeUsageSession(t, configDir, sessionID, []byte("not-json\n")) + resultEvent := mustMarshal(t, map[string]any{ + "type": "result", "subtype": "success", "session_id": sessionID, "result": "done", + "modelUsage": map[string]claudeResultModelUsage{ + model: {InputTokens: 15, OutputTokens: 7, CacheReadInputTokens: 60, CacheCreationInputTokens: 2}, + }, + }) + backend := claudeUsageFixtureBackendWithEnv(t, nil, resultEvent, false, map[string]string{ + "CLAUDE_CONFIG_DIR": configDir, + }) + result := executeClaudeUsageFixtureWithOptions(t, backend, 0, ExecOptions{ + Timeout: 5 * time.Second, + ResumeSessionID: sessionID, + }) + want := map[string]TokenUsage{ + model: {InputTokens: 15, OutputTokens: 7, CacheReadTokens: 60, CacheWriteTokens: 2}, + } + if !reflect.DeepEqual(result.Usage, want) { + t.Fatalf("usage = %#v, want %#v", result.Usage, want) + } +} + +func TestClaudeExecuteUsesAppendedCostStateWithoutResult(t *testing.T) { + const sessionID = "session-cancelled" + const model = "claude-sonnet-4-6" + configDir := t.TempDir() + sessionPath := writeClaudeUsageSession(t, configDir, sessionID, + claudeCostStateFixture(t, map[string]claudeResultModelUsage{ + model: {InputTokens: 100, OutputTokens: 20, CacheReadInputTokens: 300, CacheCreationInputTokens: 40}, + }), + ) + appended := filepath.Join(t.TempDir(), "appended.jsonl") + if err := os.WriteFile(appended, claudeCostStateFixture(t, map[string]claudeResultModelUsage{ + model: {InputTokens: 115, OutputTokens: 27, CacheReadInputTokens: 360, CacheCreationInputTokens: 42}, + }), 0o600); err != nil { + t.Fatal(err) + } + assistant := json.RawMessage(`{"type":"assistant","message":{"id":"cancelled","role":"assistant","model":"claude-sonnet-4-6","usage":{"input_tokens":2,"output_tokens":1,"cache_read_input_tokens":3,"cache_creation_input_tokens":1},"content":[{"type":"text","text":"partial"},{"type":"tool_use","id":"tool_cancelled","name":"Read","input":{}}]}}`) + backend := claudeUsageFixtureBackendWithEnv(t, []json.RawMessage{assistant}, nil, true, map[string]string{ + "CLAUDE_CONFIG_DIR": configDir, + "CLAUDE_USAGE_APPEND_FILE": sessionPath, + "CLAUDE_USAGE_APPEND_FIXTURE": appended, + }) + result := executeClaudeUsageFixtureWithOptions(t, backend, 1, ExecOptions{ + Timeout: 5 * time.Second, + ResumeSessionID: sessionID, + }) + want := map[string]TokenUsage{ + model: {InputTokens: 15, OutputTokens: 7, CacheReadTokens: 60, CacheWriteTokens: 2}, + } + if !reflect.DeepEqual(result.Usage, want) { + t.Fatalf("usage = %#v, want %#v", result.Usage, want) + } +} + +func TestClaudeExecuteKeepsStreamFallbackWithoutAppendedCostState(t *testing.T) { + const sessionID = "session-killed" + const model = "claude-sonnet-4-6" + configDir := t.TempDir() + writeClaudeUsageSession(t, configDir, sessionID, + claudeCostStateFixture(t, map[string]claudeResultModelUsage{ + model: {InputTokens: 100, OutputTokens: 20, CacheReadInputTokens: 300, CacheCreationInputTokens: 40}, + }), + ) + assistant := json.RawMessage(`{"type":"assistant","message":{"id":"killed","role":"assistant","model":"claude-sonnet-4-6","usage":{"input_tokens":2,"output_tokens":1,"cache_read_input_tokens":3,"cache_creation_input_tokens":1},"content":[{"type":"text","text":"partial"},{"type":"tool_use","id":"tool_killed","name":"Read","input":{}}]}}`) + backend := claudeUsageFixtureBackendWithEnv(t, []json.RawMessage{assistant}, nil, true, map[string]string{ + "CLAUDE_CONFIG_DIR": configDir, + }) + result := executeClaudeUsageFixtureWithOptions(t, backend, 1, ExecOptions{ + Timeout: 5 * time.Second, + ResumeSessionID: sessionID, + }) + want := map[string]TokenUsage{ + model: {InputTokens: 2, CacheReadTokens: 3, CacheWriteTokens: 1}, + } + if !reflect.DeepEqual(result.Usage, want) { + t.Fatalf("usage = %#v, want %#v", result.Usage, want) + } +} + +func TestClaudeResultUsageSinceCumulativeTurns(t *testing.T) { + t.Parallel() + const model = "claude-sonnet-4-6" + totals := []claudeResultModelUsage{ + {InputTokens: 10, OutputTokens: 4, CacheReadInputTokens: 100, CacheCreationInputTokens: 20}, + {InputTokens: 17, OutputTokens: 9, CacheReadInputTokens: 160, CacheCreationInputTokens: 23}, + {InputTokens: 20, OutputTokens: 15, CacheReadInputTokens: 205, CacheCreationInputTokens: 31}, + } + wants := []TokenUsage{ + {InputTokens: 10, OutputTokens: 4, CacheReadTokens: 100, CacheWriteTokens: 20}, + {InputTokens: 7, OutputTokens: 5, CacheReadTokens: 60, CacheWriteTokens: 3}, + {InputTokens: 3, OutputTokens: 6, CacheReadTokens: 45, CacheWriteTokens: 8}, + } + var baseline map[string]TokenUsage + for turn, total := range totals { + msg := claudeSDKMessage{ModelUsage: map[string]claudeResultModelUsage{model: total}} + usage, authoritative := claudeResultUsageSince(msg, model, baseline) + if !authoritative { + t.Fatalf("turn %d usage was not authoritative", turn+1) + } + if got := usage[model]; got != wants[turn] { + t.Fatalf("turn %d usage = %+v, want %+v", turn+1, got, wants[turn]) + } + baseline = claudeModelUsage(msg.ModelUsage) + } + + perRun := claudeSDKMessage{ + Model: model, + Usage: &claudeUsage{InputTokens: 2, OutputTokens: 3, CacheReadInputTokens: 4, CacheCreationInputTokens: 5}, + } + usage, authoritative := claudeResultUsageSince(perRun, model, baseline) + if !authoritative { + t.Fatal("top-level per-run usage was not authoritative") + } + if got, want := usage[model], (TokenUsage{InputTokens: 2, OutputTokens: 3, CacheReadTokens: 4, CacheWriteTokens: 5}); got != want { + t.Fatalf("top-level per-run usage = %+v, want %+v", got, want) + } + + reset := claudeSDKMessage{ModelUsage: map[string]claudeResultModelUsage{ + model: {InputTokens: 30, OutputTokens: 14, CacheReadInputTokens: 300, CacheCreationInputTokens: 40}, + }} + usage, authoritative = claudeResultUsageSince(reset, model, baseline) + if want := claudeModelUsage(reset.ModelUsage); !authoritative || !reflect.DeepEqual(usage, want) { + t.Fatalf("reset usage = %#v, authoritative=%v; want raw usage %#v", usage, authoritative, want) + } + + baselineWithAnotherModel := map[string]TokenUsage{ + model: baseline[model], + "claude-haiku-4-5": {InputTokens: 1}, + } + currentWithoutModel := claudeSDKMessage{ModelUsage: map[string]claudeResultModelUsage{model: totals[2]}} + usage, authoritative = claudeResultUsageSince(currentWithoutModel, model, baselineWithAnotherModel) + if want := claudeModelUsage(currentWithoutModel.ModelUsage); !authoritative || !reflect.DeepEqual(usage, want) { + t.Fatalf("missing-model usage = %#v, authoritative=%v; want raw usage %#v", usage, authoritative, want) + } +} + +func TestReadLastClaudeCostStateUsageAcrossChunks(t *testing.T) { + t.Parallel() + const model = "claude-sonnet-4-6" + baseline := claudeCostStateFixture(t, map[string]claudeResultModelUsage{ + model: {InputTokens: 10, OutputTokens: 4, CacheReadInputTokens: 100, CacheCreationInputTokens: 20}, + }) + largeTrailingLine := mustMarshal(t, map[string]any{ + "type": "user", + "message": map[string]any{"content": strings.Repeat("x", claudeSessionReadChunkSize*2)}, + }) + path := filepath.Join(t.TempDir(), "session.jsonl") + data := append(append([]byte(nil), baseline...), largeTrailingLine...) + data = append(data, '\n') + data = append(data, []byte(`{"type":"cost-state"`)...) + if err := os.WriteFile(path, data, 0o600); err != nil { + t.Fatal(err) + } + f, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + usage, found, err := readLastClaudeCostStateUsage(f, 0, int64(len(data))) + if err != nil { + t.Fatal(err) + } + if !found { + t.Fatal("cost-state not found") + } + want := TokenUsage{InputTokens: 10, OutputTokens: 4, CacheReadTokens: 100, CacheWriteTokens: 20} + if got := usage[model]; got != want { + t.Fatalf("usage = %+v, want %+v", got, want) + } + if _, found, err := readLastClaudeCostStateUsage(f, int64(len(baseline)), int64(len(data))); err != nil || found { + t.Fatalf("appended range found=%v err=%v, want no cost-state", found, err) + } + + emptyPath := filepath.Join(t.TempDir(), "empty-cost-state.jsonl") + emptyCostState := claudeCostStateFixture(t, map[string]claudeResultModelUsage{}) + if err := os.WriteFile(emptyPath, emptyCostState, 0o600); err != nil { + t.Fatal(err) + } + emptyFile, err := os.Open(emptyPath) + if err != nil { + t.Fatal(err) + } + defer emptyFile.Close() + if usage, found, err := readLastClaudeCostStateUsage(emptyFile, 0, int64(len(emptyCostState))); err != nil || found || len(usage) != 0 { + t.Fatalf("empty cost-state usage=%#v found=%v err=%v, want non-authoritative", usage, found, err) + } +} + +func claudeCostStateFixture(t *testing.T, usage map[string]claudeResultModelUsage) []byte { + t.Helper() + data := mustMarshal(t, map[string]any{"type": "cost-state", "modelUsage": usage}) + return append(data, '\n') +} + +func writeClaudeUsageSession(t *testing.T, configDir, sessionID string, lines ...[]byte) string { + t.Helper() + dir := filepath.Join(configDir, "projects", "fixture-project") + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatal(err) + } + path := filepath.Join(dir, sessionID+".jsonl") + var data []byte + for _, line := range lines { + data = append(data, line...) + } + if err := os.WriteFile(path, data, 0o600); err != nil { + t.Fatal(err) + } + return path +} + func claudeUsageFixtureBackend(t *testing.T, events []json.RawMessage, result json.RawMessage, fail bool) Backend { + return claudeUsageFixtureBackendWithEnv(t, events, result, fail, nil) +} + +func claudeUsageFixtureBackendWithEnv(t *testing.T, events []json.RawMessage, result json.RawMessage, fail bool, extraEnv map[string]string) Backend { t.Helper() var stream strings.Builder for _, event := range events { @@ -154,9 +458,13 @@ func claudeUsageFixtureBackend(t *testing.T, events []json.RawMessage, result js if err != nil { t.Fatal(err) } - backend, err := New("claude", Config{ExecutablePath: executable, Logger: slog.Default(), Env: map[string]string{ + env := map[string]string{ "IS_SANDBOX": "1", "CLAUDE_FAKE_MODE": "usage_fixture", "CLAUDE_USAGE_FIXTURE": fixture, "CLAUDE_USAGE_FIXTURE_FAIL": failEnv, - }}) + } + for key, value := range extraEnv { + env[key] = value + } + backend, err := New("claude", Config{ExecutablePath: executable, Logger: slog.Default(), Env: env}) if err != nil { t.Fatal(err) } @@ -164,10 +472,14 @@ func claudeUsageFixtureBackend(t *testing.T, events []json.RawMessage, result js } func executeClaudeUsageFixture(t *testing.T, backend Backend, eventCount int) Result { + return executeClaudeUsageFixtureWithOptions(t, backend, eventCount, ExecOptions{Timeout: 5 * time.Second}) +} + +func executeClaudeUsageFixtureWithOptions(t *testing.T, backend Backend, eventCount int, opts ExecOptions) Result { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - session, err := backend.Execute(ctx, "fixture", ExecOptions{Timeout: 5 * time.Second}) + session, err := backend.Execute(ctx, "fixture", opts) if err != nil { t.Fatal(err) }