diff --git a/api/checkpoint/metadata.go b/api/checkpoint/metadata.go index 5a39e7ecd7..d7e5163a25 100644 --- a/api/checkpoint/metadata.go +++ b/api/checkpoint/metadata.go @@ -117,6 +117,9 @@ type WriteOptions struct { // TokenUsage contains the token usage for this checkpoint TokenUsage *types.TokenUsage + // ModelUsage is the per-model breakdown of TokenUsage for this checkpoint. + ModelUsage []types.ModelUsage + // SkillEvents records explicit native skill signals observed in this session. SkillEvents []types.SkillEvent @@ -371,6 +374,12 @@ type Metadata struct { // Token usage for this checkpoint TokenUsage *types.TokenUsage `json:"token_usage,omitempty"` + // ModelUsage is the per-model breakdown of TokenUsage for this checkpoint. Each + // entry pairs a model identifier with its token usage. Backends ingest this as + // the canonical per-model shape ("model" + nested "token_usage"). Omitted when + // no per-model data is available. + ModelUsage []types.ModelUsage `json:"model_usage,omitempty"` + // SkillEvents records explicit native skill signals observed in this session. // Consumers use these anchors to collapse skill-related raw transcript events. SkillEventsVersion int `json:"skill_events_version,omitempty"` @@ -433,6 +442,64 @@ func (m Metadata) GetCompactTranscriptStart() (offset int, ok bool) { return *m.CompactTranscriptStart, true } +// WithoutCost returns a copy of the session Metadata with all token-usage cost +// provenance removed (the flat TokenUsage, every per-model ModelUsage bucket, and +// each nested SubagentTokens subtree). Token counts, APICallCount, model ids, +// SessionMetrics, and every other field are preserved. +// +// The CLI no longer persists cost: entire-api prices server-side from the token +// breakdown, so cost is a display-only local estimate (see the token commands). +// Every checkpoint-metadata write funnels through this so no committed blob ever +// carries cost_usd/cost_source, while keeping the four token fields + model id +// per model that the platform's server-side pricing depends on. The returned copy +// never aliases the receiver's TokenUsage/ModelUsage, so callers can keep using +// the originals (e.g. for local diagnostics) after persisting. +func (m Metadata) WithoutCost() Metadata { + m.TokenUsage = tokenUsageWithoutCost(m.TokenUsage) + m.ModelUsage = modelUsageWithoutCost(m.ModelUsage) + return m +} + +// WithoutCost returns a copy of the root CheckpointSummary with all token-usage +// cost provenance removed. See Metadata.WithoutCost. +func (s CheckpointSummary) WithoutCost() CheckpointSummary { + s.TokenUsage = tokenUsageWithoutCost(s.TokenUsage) + s.ModelUsage = modelUsageWithoutCost(s.ModelUsage) + return s +} + +// tokenUsageWithoutCost deep-copies u with its cost fields cleared at every level +// (flat plus the whole SubagentTokens subtree), preserving the token counts and +// APICallCount. Returns nil for nil so an absent usage stays absent. +func tokenUsageWithoutCost(u *types.TokenUsage) *types.TokenUsage { + if u == nil { + return nil + } + c := *u + c.CostUSD = nil + c.CostSource = "" + c.SubagentTokens = tokenUsageWithoutCost(u.SubagentTokens) + return &c +} + +// modelUsageWithoutCost copies the per-model buckets with each bucket's cost +// cleared, preserving the model id and the four token fields the platform prices +// from. Returns nil for a nil/empty slice. +func modelUsageWithoutCost(in []types.ModelUsage) []types.ModelUsage { + if len(in) == 0 { + return in + } + out := make([]types.ModelUsage, len(in)) + for i := range in { + tu := in[i].TokenUsage + tu.CostUSD = nil + tu.CostSource = "" + tu.SubagentTokens = tokenUsageWithoutCost(in[i].TokenUsage.SubagentTokens) + out[i] = types.ModelUsage{Model: in[i].Model, TokenUsage: tu} + } + return out +} + // SessionFilePaths contains the absolute paths to session files from the git tree root. // Paths include the full checkpoint path prefix (e.g., "/a1/b2c3d4e5f6/1/metadata.json"). // Used in CheckpointSummary.Sessions to map session IDs to their file locations. @@ -482,6 +549,7 @@ type CheckpointSummary struct { FilesTouched []string `json:"files_touched"` Sessions []SessionFilePaths `json:"sessions"` TokenUsage *types.TokenUsage `json:"token_usage,omitempty"` + ModelUsage []types.ModelUsage `json:"model_usage,omitempty"` CombinedAttribution *Attribution `json:"combined_attribution,omitempty"` // HasReview is the umbrella "any review happened" flag: true when at least diff --git a/api/checkpoint/metadata_test.go b/api/checkpoint/metadata_test.go index 827374d80b..924f08021b 100644 --- a/api/checkpoint/metadata_test.go +++ b/api/checkpoint/metadata_test.go @@ -4,6 +4,8 @@ import ( "encoding/json" "strings" "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" ) func TestImportedFlagsOnSummaryAndInfo(t *testing.T) { @@ -67,3 +69,216 @@ func TestCompactTranscriptStart_JSONRoundTrip(t *testing.T) { t.Fatalf("round-trip: got (%d, %v), want (0, true)", offset, ok) } } + +// costPtr is a small helper for constructing optional cost values in tests. +func costPtr(v float64) *float64 { return &v } + +const testWithoutCostModel = "claude-opus-4-8" + +// WithoutCost must strip cost provenance at every level (flat, nested subagent, +// and per-model bucket) while preserving the four token fields, APICallCount, +// model ids, subagent token counts, and all non-cost metadata — and it must not +// mutate the receiver (the CLI keeps the costed copy for local diagnostics). +func TestMetadata_WithoutCost(t *testing.T) { + t.Parallel() + + orig := Metadata{ + Model: testWithoutCostModel, + TokenUsage: &types.TokenUsage{ + InputTokens: 1000, + OutputTokens: 200, + CacheReadTokens: 300, + CacheCreationTokens: 40, + APICallCount: 5, + CostUSD: costPtr(1.23), + CostSource: types.CostSourceEstimated, + SubagentTokens: &types.TokenUsage{ + InputTokens: 10, + OutputTokens: 2, + CostUSD: costPtr(0.05), + CostSource: types.CostSourceReported, + }, + }, + ModelUsage: []types.ModelUsage{ + {Model: testWithoutCostModel, TokenUsage: types.TokenUsage{ + InputTokens: 1000, OutputTokens: 200, CacheReadTokens: 300, CacheCreationTokens: 40, + CostUSD: costPtr(1.23), CostSource: types.CostSourceEstimated, + }}, + }, + } + + got := orig.WithoutCost() + + // Model + token counts preserved. + if got.Model != testWithoutCostModel { + t.Errorf("model = %q, want claude-opus-4-8", got.Model) + } + if got.TokenUsage.InputTokens != 1000 || got.TokenUsage.OutputTokens != 200 || + got.TokenUsage.CacheReadTokens != 300 || got.TokenUsage.CacheCreationTokens != 40 || + got.TokenUsage.APICallCount != 5 { + t.Errorf("flat token counts changed: %+v", got.TokenUsage) + } + // Cost stripped at flat, subagent, and per-model levels. + if got.TokenUsage.CostUSD != nil || got.TokenUsage.CostSource != "" { + t.Errorf("flat cost must be cleared, got %v/%q", got.TokenUsage.CostUSD, got.TokenUsage.CostSource) + } + if got.TokenUsage.SubagentTokens == nil { + t.Fatal("subagent tokens must be preserved") + } + if got.TokenUsage.SubagentTokens.InputTokens != 10 || got.TokenUsage.SubagentTokens.OutputTokens != 2 { + t.Errorf("subagent token counts changed: %+v", got.TokenUsage.SubagentTokens) + } + if got.TokenUsage.SubagentTokens.CostUSD != nil || got.TokenUsage.SubagentTokens.CostSource != "" { + t.Errorf("subagent cost must be cleared, got %v/%q", got.TokenUsage.SubagentTokens.CostUSD, got.TokenUsage.SubagentTokens.CostSource) + } + if len(got.ModelUsage) != 1 || got.ModelUsage[0].Model != testWithoutCostModel { + t.Fatalf("per-model breakdown must be preserved with model id, got %+v", got.ModelUsage) + } + mu := got.ModelUsage[0].TokenUsage + if mu.InputTokens != 1000 || mu.OutputTokens != 200 || mu.CacheReadTokens != 300 || mu.CacheCreationTokens != 40 { + t.Errorf("per-model token counts changed: %+v", mu) + } + if mu.CostUSD != nil || mu.CostSource != "" { + t.Errorf("per-model cost must be cleared, got %v/%q", mu.CostUSD, mu.CostSource) + } + + // Receiver not mutated. + if orig.TokenUsage.CostUSD == nil || *orig.TokenUsage.CostUSD != 1.23 { + t.Errorf("WithoutCost must not mutate the receiver's flat cost") + } + if orig.TokenUsage.SubagentTokens.CostUSD == nil { + t.Errorf("WithoutCost must not mutate the receiver's subagent cost") + } + if orig.ModelUsage[0].TokenUsage.CostUSD == nil { + t.Errorf("WithoutCost must not mutate the receiver's per-model cost") + } +} + +// CheckpointSummary.WithoutCost strips cost from the aggregated summary while +// keeping token counts and per-model ids. +func TestCheckpointSummary_WithoutCost(t *testing.T) { + t.Parallel() + + orig := CheckpointSummary{ + TokenUsage: &types.TokenUsage{InputTokens: 1000, OutputTokens: 200, CostUSD: costPtr(1.23), CostSource: types.CostSourceEstimated}, + ModelUsage: []types.ModelUsage{ + {Model: "gpt-5.5", TokenUsage: types.TokenUsage{InputTokens: 1000, OutputTokens: 200, CostUSD: costPtr(1.23), CostSource: types.CostSourceEstimated}}, + }, + } + + got := orig.WithoutCost() + if got.TokenUsage.InputTokens != 1000 || got.TokenUsage.OutputTokens != 200 { + t.Errorf("summary token counts changed: %+v", got.TokenUsage) + } + if got.TokenUsage.CostUSD != nil || got.TokenUsage.CostSource != "" { + t.Errorf("summary flat cost must be cleared, got %v/%q", got.TokenUsage.CostUSD, got.TokenUsage.CostSource) + } + if len(got.ModelUsage) != 1 || got.ModelUsage[0].Model != "gpt-5.5" { + t.Fatalf("summary per-model breakdown must be preserved, got %+v", got.ModelUsage) + } + if got.ModelUsage[0].TokenUsage.CostUSD != nil || got.ModelUsage[0].TokenUsage.CostSource != "" { + t.Errorf("summary per-model cost must be cleared, got %+v", got.ModelUsage[0].TokenUsage) + } + if orig.TokenUsage.CostUSD == nil { + t.Errorf("WithoutCost must not mutate the receiver") + } +} + +// The per-model breakdown must be omitted from the wire form when empty, so +// legacy readers and backends that never asked for it see no field at all. +func TestModelUsage_OmittedWhenEmpty(t *testing.T) { + t.Parallel() + + b, err := json.Marshal(Metadata{}) + if err != nil { + t.Fatalf("marshal metadata: %v", err) + } + if strings.Contains(string(b), "model_usage") { + t.Fatalf("empty ModelUsage should be omitted, got: %s", b) + } + + b, err = json.Marshal(CheckpointSummary{}) + if err != nil { + t.Fatalf("marshal summary: %v", err) + } + if strings.Contains(string(b), "model_usage") { + t.Fatalf("empty summary ModelUsage should be omitted, got: %s", b) + } +} + +// The backend contract is a nested shape: each entry has a "model" string and a +// "token_usage" object. Assert the exact raw JSON keys so a field rename can't +// silently break ingestion. +func TestModelUsage_NestedShapeExact(t *testing.T) { + t.Parallel() + + meta := Metadata{ + ModelUsage: []types.ModelUsage{ + { + Model: testWithoutCostModel, + TokenUsage: types.TokenUsage{ + InputTokens: 1200000, + OutputTokens: 3400, + CostUSD: costPtr(0.42), + CostSource: types.CostSourceEstimated, + }, + }, + }, + } + b, err := json.Marshal(meta) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + // Decode into a generic map to assert the exact nested key names. + var raw map[string]any + if err := json.Unmarshal(b, &raw); err != nil { + t.Fatalf("unmarshal raw: %v", err) + } + list, ok := raw["model_usage"].([]any) + if !ok || len(list) != 1 { + t.Fatalf("model_usage not a 1-element array: %v", raw["model_usage"]) + } + entry, ok := list[0].(map[string]any) + if !ok { + t.Fatalf("model_usage[0] not an object: %v", list[0]) + } + if entry["model"] != testWithoutCostModel { + t.Fatalf(`entry["model"] = %v, want %s`, entry["model"], testWithoutCostModel) + } + tu, ok := entry["token_usage"].(map[string]any) + if !ok { + t.Fatalf(`entry["token_usage"] not a nested object: %v`, entry["token_usage"]) + } + if in, ok := tu["input_tokens"].(float64); !ok || in != 1200000 { + t.Fatalf("token_usage.input_tokens = %v, want 1200000", tu["input_tokens"]) + } + if cost, ok := tu["cost_usd"].(float64); !ok || cost != 0.42 { + t.Fatalf("token_usage.cost_usd = %v, want 0.42", tu["cost_usd"]) + } + + // Full round-trip preserves the typed value. + var got Metadata + if err := json.Unmarshal(b, &got); err != nil { + t.Fatalf("unmarshal metadata: %v", err) + } + if len(got.ModelUsage) != 1 || got.ModelUsage[0].Model != testWithoutCostModel || + got.ModelUsage[0].TokenUsage.InputTokens != 1200000 { + t.Fatalf("round-trip mismatch: %+v", got.ModelUsage) + } +} + +// Legacy metadata written before this field existed has no "model_usage" key; +// decoding must leave ModelUsage nil (not panic, not a zero-length non-nil slice). +func TestModelUsage_LegacyJSONNil(t *testing.T) { + t.Parallel() + + legacy := `{"strategy":"manual-commit","token_usage":{"input_tokens":5}}` + var got Metadata + if err := json.Unmarshal([]byte(legacy), &got); err != nil { + t.Fatalf("unmarshal legacy: %v", err) + } + if got.ModelUsage != nil { + t.Fatalf("legacy ModelUsage = %v, want nil", got.ModelUsage) + } +} diff --git a/cmd/entire/cli/agent/agent.go b/cmd/entire/cli/agent/agent.go index 597291fa7b..d4eefd0fe3 100644 --- a/cmd/entire/cli/agent/agent.go +++ b/cmd/entire/cli/agent/agent.go @@ -211,6 +211,21 @@ type TokenCalculator interface { CalculateTokenUsage(transcriptData []byte, fromOffset int) (*TokenUsage, error) } +// ModelUsageCalculator is an optional capability for agents that can attribute +// token usage to the specific model that produced it. Implementations must cover +// the same usage the flat TokenUsage covers (including subagent usage when the +// agent also implements SubagentAwareExtractor) so per-model buckets sum to the +// flat total. Buckets carry token counts only; CostUSD is left nil unless the +// agent itself reports cost (then CostSource must be CostSourceReported). +type ModelUsageCalculator interface { + Agent + + // CalculateModelUsage attributes token usage to each producing model, + // starting at the given offset. The returned buckets must sum to the flat + // TokenUsage from CalculateTokenUsage over the same slice. + CalculateModelUsage(transcriptData []byte, fromOffset int) ([]types.ModelUsage, error) +} + // ModelExtractor extracts the LLM model identifier from a transcript for agents // that do not report the model through lifecycle hooks. Pi, for example, records // the model on every assistant message (message.model) but its hook events carry diff --git a/cmd/entire/cli/agent/capabilities.go b/cmd/entire/cli/agent/capabilities.go index 77084156bb..d8799fcc80 100644 --- a/cmd/entire/cli/agent/capabilities.go +++ b/cmd/entire/cli/agent/capabilities.go @@ -25,6 +25,7 @@ type DeclaredCaps struct { TranscriptAnalyzer bool `json:"transcript_analyzer"` TranscriptPreparer bool `json:"transcript_preparer"` TokenCalculator bool `json:"token_calculator"` + ModelUsageCalculator bool `json:"model_usage_calculator"` CompactTranscript bool `json:"compact_transcript"` TextGenerator bool `json:"text_generator"` HookResponseWriter bool `json:"hook_response_writer"` @@ -94,6 +95,12 @@ func AsTokenCalculator(ag Agent) (TokenCalculator, bool) { return declaredCapability[TokenCalculator](ag, func(c DeclaredCaps) bool { return c.TokenCalculator }) } +// AsModelUsageCalculator returns the agent as ModelUsageCalculator if it both +// implements the interface and (for CapabilityDeclarer agents) has declared the capability. +func AsModelUsageCalculator(ag Agent) (ModelUsageCalculator, bool) { + return declaredCapability[ModelUsageCalculator](ag, func(c DeclaredCaps) bool { return c.ModelUsageCalculator }) +} + // AsTextGenerator returns the agent as TextGenerator if it both // implements the interface and (for CapabilityDeclarer agents) has declared the capability. func AsTextGenerator(ag Agent) (TextGenerator, bool) { diff --git a/cmd/entire/cli/agent/claudecode/lifecycle.go b/cmd/entire/cli/agent/claudecode/lifecycle.go index cc7e1f2e30..56bff1b917 100644 --- a/cmd/entire/cli/agent/claudecode/lifecycle.go +++ b/cmd/entire/cli/agent/claudecode/lifecycle.go @@ -20,6 +20,7 @@ var ( _ agent.TranscriptAnalyzer = (*ClaudeCodeAgent)(nil) _ agent.TranscriptPreparer = (*ClaudeCodeAgent)(nil) _ agent.TokenCalculator = (*ClaudeCodeAgent)(nil) + _ agent.ModelUsageCalculator = (*ClaudeCodeAgent)(nil) _ agent.SkillEventExtractor = (*ClaudeCodeAgent)(nil) _ agent.SubagentAwareExtractor = (*ClaudeCodeAgent)(nil) _ agent.HookResponseWriter = (*ClaudeCodeAgent)(nil) diff --git a/cmd/entire/cli/agent/claudecode/testdata/stream_session_fast.jsonl b/cmd/entire/cli/agent/claudecode/testdata/stream_session_fast.jsonl new file mode 100644 index 0000000000..40a15f0bbb --- /dev/null +++ b/cmd/entire/cli/agent/claudecode/testdata/stream_session_fast.jsonl @@ -0,0 +1,6 @@ +{"type":"system","subtype":"init","cwd":"/redacted/worktree","session_id":"fa570000-aaaa-aaaa-aaaa-aaaaaaaaaaaa","model":"claude-opus-4-8","fast_mode_state":"on","uuid":"11111111-1111-1111-1111-111111111111"} +{"type":"user","message":{"role":"user","content":"go"},"session_id":"fa570000-aaaa-aaaa-aaaa-aaaaaaaaaaaa","uuid":"22222222-2222-2222-2222-222222222222"} +{"type":"assistant","message":{"model":"claude-opus-4-8","id":"msg_std","type":"message","role":"assistant","content":[{"type":"text","text":"standard turn"}],"usage":{"input_tokens":10,"cache_creation_input_tokens":100,"cache_read_input_tokens":50,"output_tokens":100,"speed":"standard"}},"session_id":"fa570000-aaaa-aaaa-aaaa-aaaaaaaaaaaa","uuid":"33333333-3333-3333-3333-333333333333"} +{"type":"assistant","message":{"model":"claude-opus-4-8","id":"msg_fast","type":"message","role":"assistant","content":[{"type":"text","text":"fast turn (streaming)"}],"usage":{"input_tokens":20,"cache_creation_input_tokens":200,"cache_read_input_tokens":80,"output_tokens":50,"speed":"fast"}},"session_id":"fa570000-aaaa-aaaa-aaaa-aaaaaaaaaaaa","uuid":"44444444-4444-4444-4444-444444444444"} +{"type":"assistant","message":{"model":"claude-opus-4-8","id":"msg_fast","type":"message","role":"assistant","content":[{"type":"text","text":"fast turn (final)"}],"usage":{"input_tokens":20,"cache_creation_input_tokens":200,"cache_read_input_tokens":80,"output_tokens":200,"speed":"fast"}},"session_id":"fa570000-aaaa-aaaa-aaaa-aaaaaaaaaaaa","uuid":"55555555-5555-5555-5555-555555555555"} +{"type":"assistant","message":{"model":"claude-sonnet-5","id":"msg_sonnet","type":"message","role":"assistant","content":[{"type":"text","text":"second model fast turn"}],"usage":{"input_tokens":5,"cache_creation_input_tokens":0,"cache_read_input_tokens":9,"output_tokens":40,"speed":"fast"}},"session_id":"fa570000-aaaa-aaaa-aaaa-aaaaaaaaaaaa","uuid":"66666666-6666-6666-6666-666666666666"} diff --git a/cmd/entire/cli/agent/claudecode/transcript.go b/cmd/entire/cli/agent/claudecode/transcript.go index a7206578ec..801134f71d 100644 --- a/cmd/entire/cli/agent/claudecode/transcript.go +++ b/cmd/entire/cli/agent/claudecode/transcript.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "path/filepath" + "sort" "strings" "github.com/entireio/cli/cmd/entire/cli/agent" @@ -197,6 +198,102 @@ func CalculateTokenUsageFromFile(path string, startLine int) (*agent.TokenUsage, return CalculateTokenUsage(lines), nil } +// modelUsageRow pairs a deduplicated assistant message's usage with the model +// that produced it, so per-model attribution keys off the SAME kept row as the +// flat token loop (dedup by message.id, keeping the highest-output row). +type modelUsageRow struct { + model string + usage messageUsage +} + +// modelKeyWithSpeed returns the pricing/attribution key for a turn's model, +// appending a "-fast" suffix when the turn was billed at the fast-mode premium +// (speed == "fast", case-insensitive). It leaves the model unchanged for +// standard/empty speed, an empty model id, or an id that already ends in +// "-fast", so buckets never double-suffix. +func modelKeyWithSpeed(model, speed string) string { + if !strings.EqualFold(speed, "fast") || model == "" || strings.HasSuffix(model, "-fast") { + return model + } + return model + "-fast" +} + +// CalculateModelUsage attributes per-message token usage to the model that +// produced it (message.model), over the MAIN transcript only. The +// ModelUsageCalculator signature carries no subagentsDir, so unlike +// CalculateTotalTokenUsage this covers only the main-agent transcript; the +// dispatcher's remainder bucket attributes any subagent shortfall under the +// session model. It reuses the flat token loop's dedup semantics — one row per +// message.id, keeping the highest output_tokens (final streaming state) — and +// buckets each kept row by that row's model. Buckets are returned sorted by +// model for deterministic output and sum to the flat CalculateTokenUsage total +// over the same main-transcript slice. Cost fields are left nil; the framework +// prices each bucket. +func (c *ClaudeCodeAgent) CalculateModelUsage(transcriptData []byte, fromOffset int) ([]agent.ModelUsage, error) { + if len(transcriptData) == 0 { + return nil, nil + } + + sliced := transcript.SliceFromLine(transcriptData, fromOffset) + parsed, err := transcript.ParseFromBytes(sliced) + if err != nil { + return nil, fmt.Errorf("failed to parse transcript: %w", err) + } + + // Deduplicate by message.id, keeping the highest-output row; the model rides + // the kept row so attribution matches the flat token loop exactly. + rowByID := make(map[string]modelUsageRow) + for _, line := range parsed { + if line.Type != envelopeTypeAssistant { + continue + } + var msg messageWithUsage + if err := json.Unmarshal(line.Message, &msg); err != nil { + continue + } + if msg.ID == "" { + continue + } + existing, exists := rowByID[msg.ID] + if !exists || msg.Usage.OutputTokens > existing.usage.OutputTokens { + rowByID[msg.ID] = modelUsageRow{model: msg.Model, usage: msg.Usage} + } + } + + // Aggregate the deduplicated rows into per-model buckets. A fast-mode turn + // (usage.speed == "fast" on the kept row) buckets under the model's "-fast" + // variant so the dispatcher prices it at the fast premium; the speed rides + // the SAME dedup-kept row as the model, keeping attribution consistent with + // the flat token loop. + byModel := make(map[string]*agent.TokenUsage) + for _, row := range rowByID { + key := modelKeyWithSpeed(row.model, row.usage.Speed) + u := byModel[key] + if u == nil { + u = &agent.TokenUsage{} + byModel[key] = u + } + u.InputTokens += row.usage.InputTokens + u.CacheCreationTokens += row.usage.CacheCreationInputTokens + u.CacheReadTokens += row.usage.CacheReadInputTokens + u.OutputTokens += row.usage.OutputTokens + u.APICallCount++ + } + + // Emit in a deterministic order (map iteration is randomized). + models := make([]string, 0, len(byModel)) + for m := range byModel { + models = append(models, m) + } + sort.Strings(models) + + buckets := make([]agent.ModelUsage, 0, len(models)) + for _, m := range models { + buckets = append(buckets, agent.ModelUsage{Model: m, TokenUsage: *byModel[m]}) + } + return buckets, nil +} + // ExtractSpawnedAgentIDs extracts agent IDs from Task tool results in a transcript. // When a Task tool completes, the tool_result contains "agentId: " in its content. // Returns a map of agentID -> toolUseID for all spawned agents. diff --git a/cmd/entire/cli/agent/claudecode/transcript_test.go b/cmd/entire/cli/agent/claudecode/transcript_test.go index 746dc836e9..d22fe3e420 100644 --- a/cmd/entire/cli/agent/claudecode/transcript_test.go +++ b/cmd/entire/cli/agent/claudecode/transcript_test.go @@ -340,6 +340,200 @@ func TestCalculateTokenUsage_StreamingDeduplication(t *testing.T) { } } +func TestCalculateModelUsage_TwoModelsWithStreamingDedup(t *testing.T) { + t.Parallel() + + // Two models interleaved. msg_001 (opus) streams three rows with increasing + // output_tokens; the model stays constant on the id. msg_002 (opus) and + // msg_003 (fable) are single rows. Dedup keeps the max-output row per id, and + // the model rides that kept row. + data := buildJSONL( + `{"type":"user","uuid":"u1","message":{"content":"go"}}`, + `{"type":"assistant","uuid":"a1a","message":{"id":"msg_001","model":"claude-opus-4-8","usage":{"input_tokens":10,"cache_creation_input_tokens":100,"cache_read_input_tokens":50,"output_tokens":1}}}`, + `{"type":"assistant","uuid":"a1b","message":{"id":"msg_001","model":"claude-opus-4-8","usage":{"input_tokens":10,"cache_creation_input_tokens":100,"cache_read_input_tokens":50,"output_tokens":5}}}`, + `{"type":"assistant","uuid":"a1c","message":{"id":"msg_001","model":"claude-opus-4-8","usage":{"input_tokens":10,"cache_creation_input_tokens":100,"cache_read_input_tokens":50,"output_tokens":20}}}`, + `{"type":"assistant","uuid":"a2","message":{"id":"msg_002","model":"claude-opus-4-8","usage":{"input_tokens":7,"cache_creation_input_tokens":30,"cache_read_input_tokens":0,"output_tokens":40}}}`, + `{"type":"assistant","uuid":"a3","message":{"id":"msg_003","model":"claude-fable-5","usage":{"input_tokens":3,"cache_creation_input_tokens":0,"cache_read_input_tokens":9,"output_tokens":12}}}`, + ) + + c := &ClaudeCodeAgent{} + buckets, err := c.CalculateModelUsage(data, 0) + if err != nil { + t.Fatalf("CalculateModelUsage error: %v", err) + } + if len(buckets) != 2 { + t.Fatalf("buckets = %d, want 2", len(buckets)) + } + + // Deterministic order: sorted by model string. "claude-fable-5" < "claude-opus-4-8". + if buckets[0].Model != "claude-fable-5" || buckets[1].Model != "claude-opus-4-8" { + t.Fatalf("bucket order = [%q, %q], want [claude-fable-5, claude-opus-4-8]", + buckets[0].Model, buckets[1].Model) + } + + // opus: msg_001 (deduped to output=20) + msg_002. + opus := buckets[1].TokenUsage + if opus.InputTokens != 17 { // 10 + 7 (msg_001 counted once) + t.Errorf("opus input = %d, want 17", opus.InputTokens) + } + if opus.CacheCreationTokens != 130 { // 100 + 30 + t.Errorf("opus cache_creation = %d, want 130", opus.CacheCreationTokens) + } + if opus.CacheReadTokens != 50 { // 50 + 0 + t.Errorf("opus cache_read = %d, want 50", opus.CacheReadTokens) + } + if opus.OutputTokens != 60 { // 20 (max of msg_001) + 40 + t.Errorf("opus output = %d, want 60", opus.OutputTokens) + } + if opus.APICallCount != 2 { // msg_001 + msg_002 + t.Errorf("opus api_call_count = %d, want 2", opus.APICallCount) + } + + // fable: msg_003 only. + fable := buckets[0].TokenUsage + if fable.InputTokens != 3 || fable.CacheReadTokens != 9 || fable.OutputTokens != 12 || fable.APICallCount != 1 { + t.Errorf("fable bucket = %+v, want in=3 cacheRead=9 out=12 calls=1", fable) + } + + // Buckets must sum to the flat CalculateTokenUsage total (main transcript only). + flat, err := c.CalculateTokenUsage(data, 0) + if err != nil { + t.Fatalf("CalculateTokenUsage error: %v", err) + } + var sumIn, sumCC, sumCR, sumOut, sumCalls int + for _, b := range buckets { + sumIn += b.TokenUsage.InputTokens + sumCC += b.TokenUsage.CacheCreationTokens + sumCR += b.TokenUsage.CacheReadTokens + sumOut += b.TokenUsage.OutputTokens + sumCalls += b.TokenUsage.APICallCount + } + if sumIn != flat.InputTokens || sumCC != flat.CacheCreationTokens || + sumCR != flat.CacheReadTokens || sumOut != flat.OutputTokens || sumCalls != flat.APICallCount { + t.Errorf("bucket sums (in=%d cc=%d cr=%d out=%d calls=%d) != flat (in=%d cc=%d cr=%d out=%d calls=%d)", + sumIn, sumCC, sumCR, sumOut, sumCalls, + flat.InputTokens, flat.CacheCreationTokens, flat.CacheReadTokens, flat.OutputTokens, flat.APICallCount) + } +} + +func TestCalculateModelUsage_Offset(t *testing.T) { + t.Parallel() + + // From an offset past the first assistant turn, only later turns contribute. + data := buildJSONL( + `{"type":"user","uuid":"u1","message":{"content":"go"}}`, + `{"type":"assistant","uuid":"a1","message":{"id":"msg_001","model":"claude-opus-4-8","usage":{"input_tokens":10,"output_tokens":20}}}`, + `{"type":"user","uuid":"u2","message":{"content":"more"}}`, + `{"type":"assistant","uuid":"a2","message":{"id":"msg_002","model":"claude-fable-5","usage":{"input_tokens":15,"output_tokens":25}}}`, + ) + + c := &ClaudeCodeAgent{} + buckets, err := c.CalculateModelUsage(data, 2) + if err != nil { + t.Fatalf("CalculateModelUsage error: %v", err) + } + if len(buckets) != 1 || buckets[0].Model != "claude-fable-5" { + t.Fatalf("buckets = %+v, want single claude-fable-5", buckets) + } + if buckets[0].TokenUsage.InputTokens != 15 || buckets[0].TokenUsage.OutputTokens != 25 { + t.Errorf("bucket = %+v, want in=15 out=25", buckets[0].TokenUsage) + } +} + +func TestCalculateModelUsage_FastModeBucketsSeparately(t *testing.T) { + t.Parallel() + + // Fixture: a standard turn and a fast turn on claude-opus-4-8 (the fast turn + // is a streaming duplicate pair — same message.id, increasing output, speed + // on both rows), plus a fast turn on a second model. Fast turns must bucket + // under the model's "-fast" variant; the standard turn stays on the base id. + data, err := os.ReadFile("testdata/stream_session_fast.jsonl") + if err != nil { + t.Fatalf("read fixture: %v", err) + } + + c := &ClaudeCodeAgent{} + buckets, err := c.CalculateModelUsage(data, 0) + if err != nil { + t.Fatalf("CalculateModelUsage error: %v", err) + } + + // Deterministic order sorted by model string: + // "claude-opus-4-8" < "claude-opus-4-8-fast" < "claude-sonnet-5-fast". + if len(buckets) != 3 { + t.Fatalf("buckets = %d, want 3", len(buckets)) + } + if buckets[0].Model != "claude-opus-4-8" || + buckets[1].Model != "claude-opus-4-8-fast" || + buckets[2].Model != "claude-sonnet-5-fast" { + t.Fatalf("bucket models = [%q, %q, %q], want [claude-opus-4-8, claude-opus-4-8-fast, claude-sonnet-5-fast]", + buckets[0].Model, buckets[1].Model, buckets[2].Model) + } + + // Standard bucket: the single non-fast opus turn. + std := buckets[0].TokenUsage + if std.InputTokens != 10 || std.CacheCreationTokens != 100 || std.CacheReadTokens != 50 || + std.OutputTokens != 100 || std.APICallCount != 1 { + t.Errorf("standard bucket = %+v, want in=10 cc=100 cr=50 out=100 calls=1", std) + } + + // Fast bucket: streaming dedup kept the max-output row (200), and the speed + // rode that same kept row so it landed here, not on the base id. + fast := buckets[1].TokenUsage + if fast.InputTokens != 20 || fast.CacheCreationTokens != 200 || fast.CacheReadTokens != 80 || + fast.OutputTokens != 200 || fast.APICallCount != 1 { + t.Errorf("fast bucket = %+v, want in=20 cc=200 cr=80 out=200 calls=1", fast) + } + + // Second-model fast bucket. + sonnetFast := buckets[2].TokenUsage + if sonnetFast.InputTokens != 5 || sonnetFast.CacheReadTokens != 9 || + sonnetFast.OutputTokens != 40 || sonnetFast.APICallCount != 1 { + t.Errorf("sonnet fast bucket = %+v, want in=5 cr=9 out=40 calls=1", sonnetFast) + } + + // Buckets must still sum to the flat total: speed splits attribution but + // never changes token counts. + flat, err := c.CalculateTokenUsage(data, 0) + if err != nil { + t.Fatalf("CalculateTokenUsage error: %v", err) + } + var sumIn, sumCC, sumCR, sumOut, sumCalls int + for _, b := range buckets { + sumIn += b.TokenUsage.InputTokens + sumCC += b.TokenUsage.CacheCreationTokens + sumCR += b.TokenUsage.CacheReadTokens + sumOut += b.TokenUsage.OutputTokens + sumCalls += b.TokenUsage.APICallCount + } + if sumIn != flat.InputTokens || sumCC != flat.CacheCreationTokens || + sumCR != flat.CacheReadTokens || sumOut != flat.OutputTokens || sumCalls != flat.APICallCount { + t.Errorf("bucket sums (in=%d cc=%d cr=%d out=%d calls=%d) != flat (in=%d cc=%d cr=%d out=%d calls=%d)", + sumIn, sumCC, sumCR, sumOut, sumCalls, + flat.InputTokens, flat.CacheCreationTokens, flat.CacheReadTokens, flat.OutputTokens, flat.APICallCount) + } +} + +func TestModelKeyWithSpeed(t *testing.T) { + t.Parallel() + + cases := []struct { + model, speed, want string + }{ + {"claude-opus-4-8", "fast", "claude-opus-4-8-fast"}, + {"claude-opus-4-8", "FAST", "claude-opus-4-8-fast"}, + {"claude-opus-4-8", "standard", "claude-opus-4-8"}, + {"claude-opus-4-8", "", "claude-opus-4-8"}, + {"", "fast", ""}, + {"claude-opus-4-8-fast", "fast", "claude-opus-4-8-fast"}, + } + for _, tc := range cases { + if got := modelKeyWithSpeed(tc.model, tc.speed); got != tc.want { + t.Errorf("modelKeyWithSpeed(%q, %q) = %q, want %q", tc.model, tc.speed, got, tc.want) + } + } +} + func TestCalculateTokenUsage_IgnoresUserMessages(t *testing.T) { t.Parallel() diff --git a/cmd/entire/cli/agent/claudecode/types.go b/cmd/entire/cli/agent/claudecode/types.go index 2a8898e1b2..caa5c8c899 100644 --- a/cmd/entire/cli/agent/claudecode/types.go +++ b/cmd/entire/cli/agent/claudecode/types.go @@ -89,11 +89,18 @@ type messageUsage struct { CacheCreationInputTokens int `json:"cache_creation_input_tokens"` CacheReadInputTokens int `json:"cache_read_input_tokens"` OutputTokens int `json:"output_tokens"` + // Speed distinguishes fast-mode turns ("fast") from standard turns + // ("standard" or ""). Fast-mode turns bill at a premium, so per-model + // attribution routes them to the model's "-fast" pricing variant. + Speed string `json:"speed"` } // messageWithUsage represents an assistant message with usage data. -// Used for extracting token counts from Claude Code transcripts. +// Used for extracting token counts from Claude Code transcripts. Model is a +// sibling of usage in the raw message JSON (message.model), carrying the model +// that produced this assistant turn for per-model cost attribution. type messageWithUsage struct { ID string `json:"id"` + Model string `json:"model"` Usage messageUsage `json:"usage"` } diff --git a/cmd/entire/cli/agent/estimate_cost_test.go b/cmd/entire/cli/agent/estimate_cost_test.go new file mode 100644 index 0000000000..6eb5e81f7d --- /dev/null +++ b/cmd/entire/cli/agent/estimate_cost_test.go @@ -0,0 +1,107 @@ +package agent + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" +) + +// testTable (test-a: $1/MTok in+out, test-b: $2/MTok) is defined in +// token_usage_cost_test.go (same package). + +// EstimateCost prices the persisted per-model buckets at current rates and folds +// them into a single estimated cost, ignoring any pre-existing cost on the input. +func TestEstimateCost_PricesModelBuckets(t *testing.T) { + t.Parallel() + models := []types.ModelUsage{ + {Model: "test-a", TokenUsage: types.TokenUsage{InputTokens: 400000, OutputTokens: 20000}}, + } + cost, source := EstimateCost(nil, models, "", testTable(t)) + if cost == nil || *cost != 0.42 { + t.Fatalf("cost = %v, want 0.42", cost) + } + if source != types.CostSourceEstimated { + t.Fatalf("source = %q, want estimated", source) + } +} + +// A pre-existing cost on the buckets must be ignored: the estimate is recomputed +// from tokens at current rates (test-a => $0.42), not the stale $9.99. +func TestEstimateCost_IgnoresPreExistingCost(t *testing.T) { + t.Parallel() + stale := 9.99 + models := []types.ModelUsage{ + {Model: "test-a", TokenUsage: types.TokenUsage{InputTokens: 400000, OutputTokens: 20000, CostUSD: &stale, CostSource: types.CostSourceReported}}, + } + cost, source := EstimateCost(nil, models, "", testTable(t)) + if cost == nil || *cost != 0.42 { + t.Fatalf("cost = %v, want 0.42 (recomputed, not stale)", cost) + } + if source != types.CostSourceEstimated { + t.Fatalf("source = %q, want estimated", source) + } +} + +// With no per-model buckets, EstimateCost prices the flat usage (subagents +// flattened in) under the fallback model. +func TestEstimateCost_FlatFallbackIncludesSubagents(t *testing.T) { + t.Parallel() + usage := &types.TokenUsage{ + InputTokens: 300000, + OutputTokens: 20000, + SubagentTokens: &types.TokenUsage{ + InputTokens: 100000, + }, + } + // (300000 + 20000 + 100000) tokens at $1/MTok = $0.42. + cost, source := EstimateCost(usage, nil, "test-a", testTable(t)) + if cost == nil || *cost != 0.42 { + t.Fatalf("cost = %v, want 0.42", cost) + } + if source != types.CostSourceEstimated { + t.Fatalf("source = %q, want estimated", source) + } +} + +// A nil table means estimation is disabled: no cost, never $0. +func TestEstimateCost_NilTableNoCost(t *testing.T) { + t.Parallel() + models := []types.ModelUsage{{Model: "test-a", TokenUsage: types.TokenUsage{InputTokens: 400000}}} + cost, source := EstimateCost(nil, models, "", nil) + if cost != nil { + t.Fatalf("cost = %v, want nil for nil table", cost) + } + if source != "" { + t.Fatalf("source = %q, want empty", source) + } +} + +// An unknown model is unpriceable: no cost (never $0). +func TestEstimateCost_UnknownModelNoCost(t *testing.T) { + t.Parallel() + models := []types.ModelUsage{{Model: "who-dis", TokenUsage: types.TokenUsage{InputTokens: 400000}}} + cost, source := EstimateCost(nil, models, "", testTable(t)) + if cost != nil { + t.Fatalf("cost = %v, want nil for unknown model", cost) + } + if source != "" { + t.Fatalf("source = %q, want empty", source) + } +} + +// A priceable bucket alongside an unpriceable-but-token-bearing bucket folds to +// mixed (partial coverage): the estimate covers only the priceable tokens. +func TestEstimateCost_PartialCoverageMixed(t *testing.T) { + t.Parallel() + models := []types.ModelUsage{ + {Model: "test-a", TokenUsage: types.TokenUsage{InputTokens: 400000, OutputTokens: 20000}}, + {Model: "who-dis", TokenUsage: types.TokenUsage{InputTokens: 5000}}, + } + cost, source := EstimateCost(nil, models, "", testTable(t)) + if cost == nil || *cost != 0.42 { + t.Fatalf("cost = %v, want 0.42 (only the priceable bucket)", cost) + } + if source != types.CostSourceMixed { + t.Fatalf("source = %q, want mixed", source) + } +} diff --git a/cmd/entire/cli/agent/token_usage.go b/cmd/entire/cli/agent/token_usage.go index cd9ec1920f..8cb21da804 100644 --- a/cmd/entire/cli/agent/token_usage.go +++ b/cmd/entire/cli/agent/token_usage.go @@ -2,40 +2,459 @@ package agent import ( "context" + "fmt" "log/slog" + "strings" + "github.com/entireio/cli/cmd/entire/cli/agent/types" "github.com/entireio/cli/cmd/entire/cli/logging" + "github.com/entireio/cli/cmd/entire/cli/pricing" ) // CalculateTokenUsage calculates token usage from transcript data. // Returns nil if the agent doesn't support token calculation or on error. // Errors are debug-logged because callers treat nil token usage as "no data available". func CalculateTokenUsage(ctx context.Context, ag Agent, transcriptData []byte, transcriptLinesAtStart int, subagentsDir string) *TokenUsage { - if ag == nil { + usage, err := flatTokenUsage(ag, transcriptData, transcriptLinesAtStart, subagentsDir) + if err != nil { + logging.Debug(ctx, "failed token extraction", + slog.String("error", err.Error())) return nil } + return usage +} - // Calculate token usage - prefer SubagentAwareExtractor to include subagent tokens +// flatTokenUsage computes the flat (model-agnostic) token usage using the same +// dispatch preference as CalculateTokenUsage: SubagentAwareExtractor first (to +// include subagent tokens), then TokenCalculator, else a nil usage. It returns +// the agent's error instead of logging so cost-aware callers can decide; the +// wrapped prefix preserves the previous per-path distinction in logs. +func flatTokenUsage(ag Agent, transcriptData []byte, fromOffset int, subagentsDir string) (*TokenUsage, error) { + if ag == nil { + return nil, nil //nolint:nilnil // (nil, nil) means "no token data available", which is not an error + } if subagentExtractor, ok := AsSubagentAwareExtractor(ag); ok { - usage, err := subagentExtractor.CalculateTotalTokenUsage(transcriptData, transcriptLinesAtStart, subagentsDir) + usage, err := subagentExtractor.CalculateTotalTokenUsage(transcriptData, fromOffset, subagentsDir) if err != nil { - logging.Debug(ctx, "failed subagent aware token extraction", - slog.String("error", err.Error())) - return nil + return nil, fmt.Errorf("subagent-aware token extraction: %w", err) } - return usage + return usage, nil } - if calculator, ok := AsTokenCalculator(ag); ok { - // Fall back to basic token calculation (main transcript only) - usage, err := calculator.CalculateTokenUsage(transcriptData, transcriptLinesAtStart) + usage, err := calculator.CalculateTokenUsage(transcriptData, fromOffset) + if err != nil { + return nil, fmt.Errorf("token extraction: %w", err) + } + return usage, nil + } + return nil, nil //nolint:nilnil // (nil, nil) means "no token data available", which is not an error +} + +// CalculateUsageWithCost computes the flat TokenUsage exactly as CalculateTokenUsage +// does, then attributes cost. Per-model buckets come from a ModelUsageCalculator +// when the agent implements one; otherwise the flat usage becomes a single bucket +// under fallbackModel (an empty fallbackModel yields an unpriceable bucket). Each +// unpriced bucket is estimated via table when estimation is enabled and the +// model resolves; a nil table or a lookup miss leaves the bucket's cost nil. The +// flat cost is the sum of bucket costs and its source folds the bucket sources +// (CostSourceMixed when priced and unpriced-with-tokens buckets coexist). +// Returns the flat usage (cost fields populated) and the buckets; (nil, nil, nil) +// when the agent produces no usage. +func CalculateUsageWithCost(ag Agent, transcriptData []byte, fromOffset int, subagentsDir string, table *pricing.Table, fallbackModel string, disableEstimation bool) (*types.TokenUsage, []types.ModelUsage, error) { + fallbackModel = resolveTierFallback(fallbackModel, table) + + flat, err := flatTokenUsage(ag, transcriptData, fromOffset, subagentsDir) + if err != nil { + return nil, nil, err + } + if flat == nil { + return nil, nil, nil + } + + buckets, err := modelUsageBuckets(ag, transcriptData, fromOffset, flat, fallbackModel) + if err != nil { + return nil, nil, err + } + + applyTierVariant(buckets, fallbackModel) + + // The per-model buckets can cover fewer tokens than the flat total: an + // agent's CalculateModelUsage sees only the main transcript, while the flat + // path (SubagentAwareExtractor) may fold in extra usage the buckets never + // saw. Left alone, that shortfall would go unpriced, silently undercounting + // cost while CostSource still read "estimated". Attribute any shortfall to a + // remainder bucket under fallbackModel so the pricing pass either estimates + // it (priceable fallback) or leaves it unpriced (unpriceable fallback), in + // which case foldBucketCost's mixed rule marks coverage honestly. + if rem, ok := remainderBucket(flat, buckets, fallbackModel); ok { + buckets = append(buckets, rem) + } + + // Revert unpriceable variant ids (e.g. a fast turn on a model with no -fast + // rate) to their base so they estimate at the base rate instead of going + // unpriced under a premium-claiming id the table can't resolve. + downgradeUnpriceableVariants(buckets, table) + + priceBuckets(buckets, table, disableEstimation) + cost, source := foldBucketCost(buckets) + flat.CostUSD = cost + flat.CostSource = source + return flat, buckets, nil +} + +// PriceUsage attributes cost to an already-computed flat TokenUsage (e.g. one an +// agent reported through its stop hook, like Cursor) by treating it as a single +// bucket under model and running the same pricing pass as CalculateUsageWithCost. +// A usage that already carries a reported cost is preserved. Returns a copy of +// usage with its cost fields populated plus the single bucket; (nil, nil) when +// usage is nil. +func PriceUsage(usage *types.TokenUsage, model string, table *pricing.Table, disableEstimation bool) (*types.TokenUsage, []types.ModelUsage) { + if usage == nil { + return nil, nil + } + // Build the single bucket without aliasing usage.SubagentTokens: a shared + // subtree pointer would let a downstream in-place aggregator mutate the + // caller's usage. Estimate prices only the scalar fields, so dropping the + // subtree does not change the computed cost. A reported cost is preserved so + // priceBuckets leaves it untouched. + bucket := *usage + bucket.SubagentTokens = nil + buckets := []types.ModelUsage{{Model: model, TokenUsage: bucket}} + priceBuckets(buckets, table, disableEstimation) + cost, source := foldBucketCost(buckets) + out := *usage + out.CostUSD = cost + out.CostSource = source + return &out, buckets +} + +// EstimateCost computes a LOCAL, on-the-fly USD cost estimate for an +// already-recorded token breakdown, for DISPLAY ONLY. The CLI no longer persists +// cost — entire-api prices server-side from the token breakdown — so the token +// commands (`entire session tokens`, `entire checkpoint tokens`, `entire tokens +// profile`) call this to show a clearly-labeled local estimate alongside the +// token counts. +// +// It prefers the persisted per-model buckets, pricing each priceable bucket at +// table's current rates and folding them exactly as CalculateUsageWithCost's +// pricing pass does. When no buckets are supplied it falls back to pricing the +// flat usage (subagent tokens flattened in) under fallbackModel. It never trusts +// any pre-existing cost — every bucket is re-estimated — so the result is a pure +// current-rate estimate whose CostSource is CostSourceEstimated (or +// CostSourceMixed when some tokens are priceable and some are not). Returns +// (nil, "") when nothing is priceable (nil table, unknown models, or no billable +// tokens) so callers render no cost, never $0. +func EstimateCost(usage *types.TokenUsage, models []types.ModelUsage, fallbackModel string, table *pricing.Table) (*float64, string) { + if table == nil { + return nil, "" + } + var buckets []types.ModelUsage + switch { + case len(models) > 0: + buckets = make([]types.ModelUsage, len(models)) + for i := range models { + buckets[i] = types.ModelUsage{Model: models[i].Model, TokenUsage: bucketTokens(&models[i].TokenUsage)} + } + case usage != nil: + // No per-model breakdown: price the flat total, flattening the subagent + // subtree into the single bucket so subagent tokens are not dropped. + buckets = []types.ModelUsage{{Model: fallbackModel, TokenUsage: flattenTokenUsage(usage)}} + default: + return nil, "" + } + priceBuckets(buckets, table, false) + return foldBucketCost(buckets) +} + +// tierVariantSuffixes are the pricing-tier decorations a caller may have +// appended to fallbackModel via settings.PricingModelForAgent (e.g. the Codex +// priority service tier). They exist in the pricing table as distinct model +// ids but never appear in transcripts. This is the subset the fallback-model +// machinery (resolveTierFallback/applyTierVariant) is responsible for. +var tierVariantSuffixes = []string{"-priority"} + +// variantSuffixes lists every pricing decoration a bucket's model id may carry +// that the table might not price: "-fast" (appended by claudecode's +// modelKeyWithSpeed for fast-mode turns) and "-priority" (the Codex service-tier +// knob). downgradeUnpriceableVariants reverts an unpriceable one to its base id. +// It is the superset of tierVariantSuffixes because a speed variant reaches a +// bucket through a ModelUsageCalculator, not through the fallback model. +var variantSuffixes = []string{"-fast", "-priority"} + +// resolveTierFallback downgrades a tier-suffixed fallback model to its base id +// when the table cannot price the variant. The tier knob suffixes whatever +// model the agent reports, but only some variants have published rates (e.g. +// gpt-5.5-priority): pricing a gpt-5.3-codex session under a nonexistent +// gpt-5.3-codex-priority id would leave it entirely unpriced — worse than the +// standard-rate estimate, and the suffixed bucket id would claim a premium that +// was never applied. Falling back to the base id keeps the estimate an honest +// undercount, mirroring how fast turns on models without a published fast rate +// price at the base id. With a nil table nothing can be priced or verified, so +// the base id is kept there too. +func resolveTierFallback(fallbackModel string, table *pricing.Table) string { + for _, suffix := range tierVariantSuffixes { + base := strings.TrimSuffix(fallbackModel, suffix) + if base == fallbackModel || base == "" { + continue + } + if table != nil { + if _, ok := table.Lookup(fallbackModel); ok { + return fallbackModel + } + } + return base + } + return fallbackModel +} + +// applyTierVariant retargets per-model buckets onto the caller's tier-variant +// pricing id. When fallbackModel carries a known tier suffix (the caller opted +// the session into that service tier), a ModelUsageCalculator's buckets still +// carry the raw transcript model id — priced as-is they'd silently fall back to +// standard rates, re-introducing the tier-loss defect the suffixed fallback +// fixed on the no-calculator path. Only buckets matching the fallback's base id +// are remapped: other models have no table entry for the variant, and an +// unpriceable id would be worse than a standard-rate estimate. +func applyTierVariant(buckets []types.ModelUsage, fallbackModel string) { + for _, suffix := range tierVariantSuffixes { + base := strings.TrimSuffix(fallbackModel, suffix) + if base == fallbackModel || base == "" { + continue + } + for i := range buckets { + if buckets[i].Model == base { + buckets[i].Model = fallbackModel + } + } + } +} + +// downgradeUnpriceableVariants reverts, in place, each bucket whose model id +// carries a known variant suffix (-fast, -priority) to its base id when the table +// cannot price the variant but can price the base. A fast turn on a model with no +// published fast rate (e.g. claude-fable-5-fast) would otherwise miss the table +// entirely and go unpriced — worse than the honest base-rate estimate, under an +// id claiming a premium that was never applied. Detection stays truthful at the +// source (modelKeyWithSpeed, the tier knob); only here, at pricing time, is +// priceability known. This is the calculator-bucket analogue of +// resolveTierFallback, which only covers the fallback model. +// +// Rekeying is in place, not merging: a downgraded bucket may now share a model id +// with another bucket, but same-model buckets are already emitted today (a +// ModelUsageCalculator's per-model bucket plus the remainder bucket both sit under +// the session model), and both foldBucketCost here and strategy.accumulateModelUsage +// downstream fold buckets by model key. Leaving the duplicate is therefore the +// token- and cost-conserving choice, and it keeps the pre-existing remainder +// behavior (separate same-model buckets) unchanged. +func downgradeUnpriceableVariants(buckets []types.ModelUsage, table *pricing.Table) { + for i := range buckets { + if base, ok := downgradedVariantBase(buckets[i].Model, table); ok { + buckets[i].Model = base + } + } +} + +// downgradedVariantBase returns the base id a variant-suffixed model should be +// priced under, and whether the downgrade applies. It reports (base, true) when +// the model carries a known variant suffix and either the table is nil (nothing +// can be verified, so fall back — mirroring resolveTierFallback), or the table +// cannot price the variant but can price the base. It reports ("", false) for a +// priceable variant (keep it), a variant whose base is also unpriceable (the base +// is no better and the variant id is the more truthful label), or a model with no +// known variant suffix. +func downgradedVariantBase(model string, table *pricing.Table) (string, bool) { + for _, suffix := range variantSuffixes { + base := strings.TrimSuffix(model, suffix) + if base == model || base == "" { + continue + } + if table == nil { + return base, true + } + if _, ok := table.Lookup(model); ok { + return "", false + } + if _, ok := table.Lookup(base); ok { + return base, true + } + return "", false + } + return "", false +} + +// modelUsageBuckets returns per-model token buckets for the transcript slice. It +// prefers a ModelUsageCalculator; otherwise the whole flat usage becomes a +// single bucket under fallbackModel with its cost fields cleared (buckets carry +// token counts only). +func modelUsageBuckets(ag Agent, transcriptData []byte, fromOffset int, flat *types.TokenUsage, fallbackModel string) ([]types.ModelUsage, error) { + if calc, ok := AsModelUsageCalculator(ag); ok { + buckets, err := calc.CalculateModelUsage(transcriptData, fromOffset) if err != nil { - logging.Debug(ctx, "failed token extraction", - slog.String("error", err.Error())) - return nil + return nil, fmt.Errorf("model usage attribution: %w", err) + } + return buckets, nil + } + return []types.ModelUsage{{Model: fallbackModel, TokenUsage: bucketTokens(flat)}}, nil +} + +// remainderBucket returns a bucket carrying the per-field token shortfall +// between the flat total (INCLUDING every nested SubagentTokens subtree) and the +// sum of the per-model buckets, attributed to fallbackModel. The flat side is +// flattened via flattenTokenUsage because the per-model buckets (and, per +// bucketTokens, the fallback single bucket) carry only main-scoped scalar +// counts, while a SubagentAwareExtractor folds subagent usage into the flat +// total's SubagentTokens subtree. Without accounting for that subtree the +// subagent tokens would go entirely unpriced — e.g. a Claude Code session with +// 100k main + 500k subagent tokens would price only the 100k. Here the 500k +// shortfall becomes a remainder bucket under fallbackModel. +// +// Each field is clamped at 0 (a bucket may legitimately exceed the flat total on +// an individual field). The bool is false when no billable token shortfall +// exists: the fallback (single-bucket) path with no subagents and any +// ModelUsageCalculator whose buckets sum to the flat total both yield a zero +// remainder. Cost fields are left nil so the pricing pass estimates the +// remainder when fallbackModel is priceable, or leaves it unpriced otherwise. +func remainderBucket(flat *types.TokenUsage, buckets []types.ModelUsage, fallbackModel string) (types.ModelUsage, bool) { + var sum types.TokenUsage + for i := range buckets { + b := buckets[i].TokenUsage + sum.InputTokens += b.InputTokens + sum.CacheCreationTokens += b.CacheCreationTokens + sum.CacheReadTokens += b.CacheReadTokens + sum.OutputTokens += b.OutputTokens + sum.APICallCount += b.APICallCount + } + flatTotal := flattenTokenUsage(flat) + short := types.TokenUsage{ + InputTokens: clampNonNegative(flatTotal.InputTokens - sum.InputTokens), + CacheCreationTokens: clampNonNegative(flatTotal.CacheCreationTokens - sum.CacheCreationTokens), + CacheReadTokens: clampNonNegative(flatTotal.CacheReadTokens - sum.CacheReadTokens), + OutputTokens: clampNonNegative(flatTotal.OutputTokens - sum.OutputTokens), + APICallCount: clampNonNegative(flatTotal.APICallCount - sum.APICallCount), + } + if short.InputTokens+short.CacheCreationTokens+short.CacheReadTokens+short.OutputTokens == 0 { + return types.ModelUsage{}, false + } + return types.ModelUsage{Model: fallbackModel, TokenUsage: short}, true +} + +// flattenTokenUsage returns u's five scalar token fields summed with every +// nested SubagentTokens subtree, at arbitrary depth. It is nil-safe (nil yields +// a zero usage) and reads only token counts — the returned usage carries no cost +// fields. This is how remainderBucket recovers the true billable total from a +// flat usage whose subagent tokens live in a subtree rather than in the top-level +// scalar fields. +func flattenTokenUsage(u *types.TokenUsage) types.TokenUsage { + var out types.TokenUsage + if u == nil { + return out + } + out.InputTokens = u.InputTokens + out.CacheCreationTokens = u.CacheCreationTokens + out.CacheReadTokens = u.CacheReadTokens + out.OutputTokens = u.OutputTokens + out.APICallCount = u.APICallCount + sub := flattenTokenUsage(u.SubagentTokens) + out.InputTokens += sub.InputTokens + out.CacheCreationTokens += sub.CacheCreationTokens + out.CacheReadTokens += sub.CacheReadTokens + out.OutputTokens += sub.OutputTokens + out.APICallCount += sub.APICallCount + return out +} + +// clampNonNegative returns v, or 0 when v is negative. +func clampNonNegative(v int) int { + if v < 0 { + return 0 + } + return v +} + +// bucketTokens returns a copy of u with its cost fields cleared and its +// SubagentTokens subtree dropped, so it serves as a per-model bucket carrying +// main-scoped scalar token counts only. Nil-ing SubagentTokens is essential: +// remainderBucket attributes the subagent subtree to its own remainder bucket, +// so a fallback bucket that also carried the shared subtree pointer would both +// double-count the subagent tokens downstream and alias u's subtree (a mutation +// hazard once aggregators fold buckets in place). +func bucketTokens(u *types.TokenUsage) types.TokenUsage { + b := *u + b.CostUSD = nil + b.CostSource = "" + b.SubagentTokens = nil + return b +} + +// priceBuckets fills in each bucket's cost in place: a bucket that already has a +// reported cost is left untouched; otherwise, when estimation is enabled and the +// table resolves the bucket's model, the cost is estimated (CostSourceEstimated). +// A bucket whose model is unknown, a bucket carrying no billable tokens, or any +// bucket when the table is nil or estimation is disabled, keeps a nil cost. +func priceBuckets(buckets []types.ModelUsage, table *pricing.Table, disableEstimation bool) { + for i := range buckets { + b := &buckets[i].TokenUsage + if b.CostUSD != nil { + continue // reported cost wins (even on a zero-token bucket) + } + if disableEstimation || table == nil { + continue + } + if !bucketHasTokens(*b) { + // No billable tokens: leave the cost nil rather than fabricating a + // $0.00 "estimated" figure. Mirrors the ModelUsageCalculator path and + // keeps zero-usage checkpoints at "no cost data" instead of flipping + // aggregated provenance from reported to mixed. + continue + } + rate, ok := table.Lookup(buckets[i].Model) + if !ok { + continue + } + cost := pricing.Estimate(rate, *b) + b.CostUSD = &cost + b.CostSource = types.CostSourceEstimated + } +} + +// foldBucketCost aggregates per-bucket costs into a single (cost, source) pair. +// Costs sum via AddCostUSD and sources fold via MergeCostSource; when at least +// one bucket was priced and at least one bucket has tokens but no cost, the +// source is CostSourceMixed to signal partial coverage. +// +// Defensive normalization: a bucket carrying a non-nil CostUSD with an empty +// CostSource (currently unreachable — priceBuckets always pairs a cost it sets +// with CostSourceEstimated, and every other producer of a priced bucket is +// expected to set a source too) is treated as CostSourceEstimated here rather +// than folding into a costed-but-unlabeled total that would render with no +// reported/estimated/mixed suffix. +func foldBucketCost(buckets []types.ModelUsage) (*float64, string) { + var cost *float64 + var source string + anyPriced := false + anyUnpricedTokens := false + for i := range buckets { + b := buckets[i].TokenUsage + bSource := b.CostSource + if b.CostUSD != nil && bSource == "" { + bSource = types.CostSourceEstimated + } + source = types.MergeCostSource(source, bSource, cost, b.CostUSD) + cost = types.AddCostUSD(cost, b.CostUSD) + switch { + case b.CostUSD != nil: + anyPriced = true + case bucketHasTokens(b): + anyUnpricedTokens = true } - return usage } + if anyPriced && anyUnpricedTokens { + return cost, types.CostSourceMixed + } + return cost, source +} - return nil +// bucketHasTokens reports whether a bucket carries any billable token count. +func bucketHasTokens(u types.TokenUsage) bool { + return u.InputTokens > 0 || u.OutputTokens > 0 || u.CacheReadTokens > 0 || u.CacheCreationTokens > 0 } diff --git a/cmd/entire/cli/agent/token_usage_cost_test.go b/cmd/entire/cli/agent/token_usage_cost_test.go new file mode 100644 index 0000000000..6350b9dccd --- /dev/null +++ b/cmd/entire/cli/agent/token_usage_cost_test.go @@ -0,0 +1,647 @@ +package agent + +import ( + "errors" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/pricing" +) + +// --- Fakes --- + +// fakeTokenCalcAgent implements TokenCalculator only (no per-model attribution). +type fakeTokenCalcAgent struct { + mockBaseAgent + + usage *TokenUsage + err error +} + +func (f *fakeTokenCalcAgent) CalculateTokenUsage([]byte, int) (*TokenUsage, error) { + return f.usage, f.err +} + +// fakeModelUsageAgent implements both TokenCalculator (for the flat total) and +// ModelUsageCalculator (for per-model buckets). +type fakeModelUsageAgent struct { + mockBaseAgent + + flat *TokenUsage + buckets []types.ModelUsage + muErr error +} + +func (f *fakeModelUsageAgent) CalculateTokenUsage([]byte, int) (*TokenUsage, error) { //nolint:unparam // test mock: error result satisfies TokenCalculator + return f.flat, nil +} + +func (f *fakeModelUsageAgent) CalculateModelUsage([]byte, int) ([]types.ModelUsage, error) { + return f.buckets, f.muErr +} + +// fakeSubagentAwareAgent implements SubagentAwareExtractor (its flat total folds +// subagent usage into a SubagentTokens subtree) but NOT ModelUsageCalculator, so +// CalculateUsageWithCost takes the fallback single-bucket path (bucketTokens) +// plus the subagent remainder bucket. +type fakeSubagentAwareAgent struct { + mockBaseAgent + + usage *TokenUsage +} + +func (f *fakeSubagentAwareAgent) ExtractAllModifiedFiles([]byte, int, string) ([]string, error) { + return nil, nil +} + +func (f *fakeSubagentAwareAgent) CalculateTotalTokenUsage([]byte, int, string) (*TokenUsage, error) { //nolint:unparam // test mock: error result satisfies SubagentAwareExtractor + return f.usage, nil +} + +func ptrFloat(v float64) *float64 { return &v } + +// testTable returns a pricing table with two clean-rate override models on top +// of the embedded defaults: test-a ($1/MTok in&out) and test-b ($2/MTok in&out). +func testTable(t *testing.T) *pricing.Table { + t.Helper() + table, err := pricing.LoadTable([]pricing.ModelRate{ + {ID: "test-a", Provider: "test", InputPerMTok: 1, OutputPerMTok: 1}, + {ID: "test-b", Provider: "test", InputPerMTok: 2, OutputPerMTok: 2}, + }) + if err != nil { + t.Fatalf("LoadTable: %v", err) + } + return table +} + +func TestCalculateUsageWithCost_FallbackBucketPriced(t *testing.T) { + t.Parallel() + ag := &fakeTokenCalcAgent{usage: &TokenUsage{InputTokens: 1_000_000}} + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-a", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 1 { + t.Fatalf("buckets = %d, want 1", len(buckets)) + } + if buckets[0].Model != "test-a" { + t.Errorf("bucket model = %q, want test-a", buckets[0].Model) + } + if flat.CostUSD == nil || *flat.CostUSD != 1.0 { + t.Fatalf("flat cost = %v, want 1.0", flat.CostUSD) + } + if flat.CostSource != types.CostSourceEstimated { + t.Errorf("flat source = %q, want estimated", flat.CostSource) + } + if buckets[0].TokenUsage.CostUSD == nil || *buckets[0].TokenUsage.CostUSD != 1.0 { + t.Errorf("bucket cost = %v, want 1.0", buckets[0].TokenUsage.CostUSD) + } +} + +func TestCalculateUsageWithCost_TwoModelsSumEstimated(t *testing.T) { + t.Parallel() + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + {Model: "test-b", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if flat.CostUSD == nil || *flat.CostUSD != 3.0 { + t.Fatalf("flat cost = %v, want 3.0 (1+2)", flat.CostUSD) + } + if flat.CostSource != types.CostSourceEstimated { + t.Errorf("flat source = %q, want estimated", flat.CostSource) + } + if buckets[0].TokenUsage.CostUSD == nil || *buckets[0].TokenUsage.CostUSD != 1.0 { + t.Errorf("bucket a cost = %v, want 1.0", buckets[0].TokenUsage.CostUSD) + } + if buckets[1].TokenUsage.CostUSD == nil || *buckets[1].TokenUsage.CostUSD != 2.0 { + t.Errorf("bucket b cost = %v, want 2.0", buckets[1].TokenUsage.CostUSD) + } +} + +func TestCalculateUsageWithCost_UnknownModelMixed(t *testing.T) { + t.Parallel() + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + {Model: "who-knows", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if flat.CostUSD == nil || *flat.CostUSD != 1.0 { + t.Fatalf("flat cost = %v, want 1.0 (only test-a priced)", flat.CostUSD) + } + if flat.CostSource != types.CostSourceMixed { + t.Errorf("flat source = %q, want mixed", flat.CostSource) + } + if buckets[1].TokenUsage.CostUSD != nil { + t.Errorf("unknown bucket cost = %v, want nil", buckets[1].TokenUsage.CostUSD) + } +} + +func TestCalculateUsageWithCost_RemainderBucketPriced(t *testing.T) { + t.Parallel() + // Real extractors report the MAIN-transcript totals in the flat scalar fields + // and the subagent totals in a SubagentTokens subtree, while CalculateModelUsage + // only ever sees the main transcript. So the subagent usage surfaces solely as + // a shortfall between the flattened flat total (main + subagents) and the + // per-model bucket sum, and must be priced under fallbackModel. Here: main 1M + // under test-a, subagent 1M must land in a test-b remainder bucket. + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{ + InputTokens: 1_000_000, // main transcript + APICallCount: 3, + SubagentTokens: &TokenUsage{InputTokens: 1_000_000, APICallCount: 4}, // subagents + }, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, APICallCount: 3}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-b", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 2 { + t.Fatalf("buckets = %d, want 2 (test-a + remainder)", len(buckets)) + } + rem := buckets[1] + if rem.Model != "test-b" { + t.Errorf("remainder bucket model = %q, want test-b", rem.Model) + } + if rem.TokenUsage.InputTokens != 1_000_000 { + t.Errorf("remainder input = %d, want 1_000_000 (subagent shortfall)", rem.TokenUsage.InputTokens) + } + if rem.TokenUsage.APICallCount != 4 { + t.Errorf("remainder api_call_count = %d, want 4 (subagent api calls)", rem.TokenUsage.APICallCount) + } + if rem.TokenUsage.CostUSD == nil || *rem.TokenUsage.CostUSD != 2.0 { + t.Errorf("remainder cost = %v, want 2.0 (1M @ $2/MTok)", rem.TokenUsage.CostUSD) + } + // F7/F8: no returned bucket may carry a SubagentTokens subtree. + for i := range buckets { + if buckets[i].TokenUsage.SubagentTokens != nil { + t.Errorf("bucket %d carries SubagentTokens, want nil", i) + } + } + // test-a 1M@$1 = 1.0, remainder 1M@$2 = 2.0 → 3.0, all estimated. + if flat.CostUSD == nil || *flat.CostUSD != 3.0 { + t.Fatalf("flat cost = %v, want 3.0", flat.CostUSD) + } + if flat.CostSource != types.CostSourceEstimated { + t.Errorf("flat source = %q, want estimated", flat.CostSource) + } +} + +func TestCalculateUsageWithCost_RemainderBucketInflatedScalar(t *testing.T) { + t.Parallel() + // A ModelUsageCalculator whose buckets under-cover the flat SCALAR total (no + // subagent subtree involved) still yields a remainder for the shortfall. + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000, APICallCount: 7}, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, APICallCount: 3}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-b", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 2 { + t.Fatalf("buckets = %d, want 2 (test-a + remainder)", len(buckets)) + } + rem := buckets[1] + if rem.TokenUsage.InputTokens != 1_000_000 { + t.Errorf("remainder input = %d, want 1_000_000", rem.TokenUsage.InputTokens) + } + if rem.TokenUsage.APICallCount != 4 { + t.Errorf("remainder api_call_count = %d, want 4 (7 flat - 3 bucketed)", rem.TokenUsage.APICallCount) + } + if flat.CostUSD == nil || *flat.CostUSD != 3.0 { + t.Fatalf("flat cost = %v, want 3.0", flat.CostUSD) + } +} + +func TestCalculateUsageWithCost_RemainderBucketNestedSubagents(t *testing.T) { + t.Parallel() + // Subagents that spawn subagents: the flat total's SubagentTokens nests two + // levels deep. flattenTokenUsage must sum EVERY level so the whole subtree is + // priced in the remainder bucket, not just the first. + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{ + InputTokens: 1_000_000, // main + APICallCount: 1, + SubagentTokens: &TokenUsage{ + InputTokens: 1_000_000, // depth-1 subagent + APICallCount: 1, + SubagentTokens: &TokenUsage{InputTokens: 1_000_000, APICallCount: 1}, // depth-2 subagent + }, + }, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, APICallCount: 1}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-b", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 2 { + t.Fatalf("buckets = %d, want 2", len(buckets)) + } + rem := buckets[1] + if rem.TokenUsage.InputTokens != 2_000_000 { + t.Errorf("remainder input = %d, want 2_000_000 (both nested subagent levels)", rem.TokenUsage.InputTokens) + } + if rem.TokenUsage.APICallCount != 2 { + t.Errorf("remainder api_call_count = %d, want 2 (both subagent levels)", rem.TokenUsage.APICallCount) + } + // test-a 1M@$1 = 1.0, remainder 2M@$2 = 4.0 → 5.0. + if flat.CostUSD == nil || *flat.CostUSD != 5.0 { + t.Fatalf("flat cost = %v, want 5.0", flat.CostUSD) + } +} + +func TestCalculateUsageWithCost_SubagentRemainderFallbackNoAliasing(t *testing.T) { + t.Parallel() + // SubagentAware but NOT ModelUsage: the flat total carries a subagent subtree + // and the single bucket comes from the fallback path (bucketTokens). The + // fallback bucket must carry main-scoped scalars only (nil SubagentTokens) and + // the subagent tokens must land in the remainder bucket — never aliased, so + // two runs cannot inflate and the agent's own usage is left untouched. + ag := &fakeSubagentAwareAgent{usage: &TokenUsage{ + InputTokens: 1_000_000, // main + APICallCount: 2, + SubagentTokens: &TokenUsage{InputTokens: 1_000_000, APICallCount: 3}, // subagent + }} + + run := func() ([]types.ModelUsage, *types.TokenUsage) { + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "subdir", testTable(t), "test-a", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + return buckets, flat + } + + buckets1, flat1 := run() + if len(buckets1) != 2 { + t.Fatalf("buckets = %d, want 2 (fallback main bucket + subagent remainder)", len(buckets1)) + } + for i := range buckets1 { + if buckets1[i].TokenUsage.SubagentTokens != nil { + t.Errorf("bucket %d carries SubagentTokens, want nil", i) + } + } + if buckets1[0].TokenUsage.InputTokens != 1_000_000 { + t.Errorf("main bucket input = %d, want 1_000_000", buckets1[0].TokenUsage.InputTokens) + } + rem := buckets1[1] + if rem.TokenUsage.InputTokens != 1_000_000 || rem.TokenUsage.APICallCount != 3 { + t.Errorf("remainder = %+v, want 1M input / 3 api (subagent)", rem.TokenUsage) + } + // main 1M@$1 + subagent 1M@$1 = 2.0. + if flat1.CostUSD == nil || *flat1.CostUSD != 2.0 { + t.Fatalf("flat cost = %v, want 2.0", flat1.CostUSD) + } + + // A second run must yield identical token counts — no cross-call inflation. + buckets2, _ := run() + if buckets2[0].TokenUsage.InputTokens != buckets1[0].TokenUsage.InputTokens || + buckets2[1].TokenUsage.InputTokens != buckets1[1].TokenUsage.InputTokens { + t.Fatalf("cross-call inflation: run1=[%d,%d] run2=[%d,%d]", + buckets1[0].TokenUsage.InputTokens, buckets1[1].TokenUsage.InputTokens, + buckets2[0].TokenUsage.InputTokens, buckets2[1].TokenUsage.InputTokens) + } + // The agent's own usage subtree must not have been mutated. + if ag.usage.SubagentTokens.InputTokens != 1_000_000 { + t.Fatalf("agent usage subtree mutated: %d, want 1_000_000", ag.usage.SubagentTokens.InputTokens) + } +} + +func TestCalculateUsageWithCost_RemainderUnpriceableFallbackMixed(t *testing.T) { + t.Parallel() + // Same shortfall, but the fallbackModel is not priceable. The remainder + // bucket stays unpriced (tokens, no cost), so coverage folds to mixed. + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "who-knows", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 2 { + t.Fatalf("buckets = %d, want 2 (test-a + remainder)", len(buckets)) + } + rem := buckets[1] + if rem.Model != "who-knows" || rem.TokenUsage.InputTokens != 1_000_000 { + t.Errorf("remainder bucket = %+v, want who-knows with 1M input", rem) + } + if rem.TokenUsage.CostUSD != nil { + t.Errorf("remainder cost = %v, want nil (unpriceable)", rem.TokenUsage.CostUSD) + } + // Only test-a priced (1.0); remainder has tokens but no cost → mixed. + if flat.CostUSD == nil || *flat.CostUSD != 1.0 { + t.Fatalf("flat cost = %v, want 1.0 (test-a only)", flat.CostUSD) + } + if flat.CostSource != types.CostSourceMixed { + t.Errorf("flat source = %q, want mixed", flat.CostSource) + } +} + +func TestCalculateUsageWithCost_NoRemainderWhenBucketsSumToFlat(t *testing.T) { + t.Parallel() + // Buckets already sum to the flat total: no remainder bucket is appended. + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + {Model: "test-b", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + }, + } + + _, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-a", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 2 { + t.Fatalf("buckets = %d, want 2 (no remainder appended)", len(buckets)) + } +} + +func TestCalculateUsageWithCost_ReportedKeptEstimationOff(t *testing.T) { + t.Parallel() + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, CostUSD: ptrFloat(5.0), CostSource: types.CostSourceReported}}, + {Model: "test-b", TokenUsage: TokenUsage{InputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "", true /* disableEstimation */) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Reported bucket survives untouched. + if buckets[0].TokenUsage.CostUSD == nil || *buckets[0].TokenUsage.CostUSD != 5.0 { + t.Errorf("reported bucket cost = %v, want 5.0", buckets[0].TokenUsage.CostUSD) + } + if buckets[0].TokenUsage.CostSource != types.CostSourceReported { + t.Errorf("reported bucket source = %q, want reported", buckets[0].TokenUsage.CostSource) + } + // Estimation is off, so the priceable bucket stays unpriced. + if buckets[1].TokenUsage.CostUSD != nil { + t.Errorf("estimated-off bucket cost = %v, want nil", buckets[1].TokenUsage.CostUSD) + } + if flat.CostUSD == nil || *flat.CostUSD != 5.0 { + t.Errorf("flat cost = %v, want 5.0 (reported only)", flat.CostUSD) + } + if flat.CostSource != types.CostSourceMixed { + t.Errorf("flat source = %q, want mixed (reported + unpriced tokens)", flat.CostSource) + } +} + +func TestCalculateUsageWithCost_NilTableNoCost(t *testing.T) { + t.Parallel() + ag := &fakeTokenCalcAgent{usage: &TokenUsage{InputTokens: 1_000_000}} + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", nil /* table */, "test-a", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if flat.CostUSD != nil { + t.Errorf("flat cost = %v, want nil", flat.CostUSD) + } + if flat.CostSource != "" { + t.Errorf("flat source = %q, want empty", flat.CostSource) + } + if buckets[0].TokenUsage.CostUSD != nil { + t.Errorf("bucket cost = %v, want nil", buckets[0].TokenUsage.CostUSD) + } +} + +func TestCalculateUsageWithCost_NoUsage(t *testing.T) { + t.Parallel() + // Agent that supports neither TokenCalculator nor ModelUsageCalculator. + flat, buckets, err := CalculateUsageWithCost(&mockBaseAgent{}, nil, 0, "", testTable(t), "test-a", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if flat != nil || buckets != nil { + t.Errorf("got (%v, %v), want (nil, nil)", flat, buckets) + } +} + +func TestCalculateUsageWithCost_FlatError(t *testing.T) { + t.Parallel() + ag := &fakeTokenCalcAgent{err: errors.New("boom")} + if _, _, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-a", false); err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestCalculateUsageWithCost_ZeroTokenBucketUnpriced(t *testing.T) { + t.Parallel() + // A non-ModelUsageCalculator agent whose usage carries no billable tokens + // (only APICallCount) takes the fallback single-bucket path under a resolvable + // fallbackModel. Estimating that bucket would fabricate a $0.00 "estimated" + // cost; instead it must stay unpriced (nil cost, empty source) so a zero-usage + // checkpoint reads as "no cost data", consistent with the ModelUsage path. + ag := &fakeTokenCalcAgent{usage: &TokenUsage{APICallCount: 5}} + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", testTable(t), "test-a", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(buckets) != 1 { + t.Fatalf("buckets = %d, want 1 (fallback bucket, no remainder)", len(buckets)) + } + if buckets[0].Model != "test-a" { + t.Errorf("bucket model = %q, want test-a", buckets[0].Model) + } + if buckets[0].TokenUsage.CostUSD != nil { + t.Errorf("zero-token bucket cost = %v, want nil (no $0 estimate)", buckets[0].TokenUsage.CostUSD) + } + if buckets[0].TokenUsage.CostSource != "" { + t.Errorf("zero-token bucket source = %q, want empty", buckets[0].TokenUsage.CostSource) + } + if flat.CostUSD != nil { + t.Errorf("flat cost = %v, want nil", flat.CostUSD) + } + if flat.CostSource != "" { + t.Errorf("flat source = %q, want empty (not estimated)", flat.CostSource) + } +} + +func TestPriceUsage_ZeroTokenUsageMergesAsReported(t *testing.T) { + t.Parallel() + // End-to-end coupling for the zero-token guard: a zero-usage checkpoint priced + // under a resolvable model must come back unpriced (nil cost, empty source), + // so that when it is later aggregated with a reported checkpoint the merged + // provenance stays "reported" instead of flipping to "mixed" — which a + // fabricated $0.00 "estimated" figure would have caused. + zero, _ := PriceUsage(&TokenUsage{APICallCount: 4}, "test-a", testTable(t), false) + if zero.CostUSD != nil { + t.Fatalf("zero-usage priced cost = %v, want nil (no $0 estimate)", zero.CostUSD) + } + if zero.CostSource != "" { + t.Fatalf("zero-usage priced source = %q, want empty", zero.CostSource) + } + + reported := &TokenUsage{InputTokens: 1_000_000, CostUSD: ptrFloat(6.25), CostSource: types.CostSourceReported} + if got := types.MergeCostSourceUsages(reported, zero); got != types.CostSourceReported { + t.Errorf("merge = %q, want reported (zero-usage checkpoint must not flip provenance)", got) + } + if got := types.MergeCostSourceUsages(zero, reported); got != types.CostSourceReported { + t.Errorf("merge (swapped) = %q, want reported", got) + } +} + +func TestPriceUsage_EstimatesUnderModel(t *testing.T) { + t.Parallel() + priced, buckets := PriceUsage(&TokenUsage{InputTokens: 1_000_000}, "test-b", testTable(t), false) + if priced.CostUSD == nil || *priced.CostUSD != 2.0 { + t.Fatalf("cost = %v, want 2.0", priced.CostUSD) + } + if priced.CostSource != types.CostSourceEstimated { + t.Errorf("source = %q, want estimated", priced.CostSource) + } + if len(buckets) != 1 || buckets[0].Model != "test-b" { + t.Errorf("buckets = %+v, want single test-b", buckets) + } +} + +func TestPriceUsage_KeepsReported(t *testing.T) { + t.Parallel() + in := &TokenUsage{InputTokens: 1_000_000, CostUSD: ptrFloat(9.0), CostSource: types.CostSourceReported} + priced, _ := PriceUsage(in, "test-a", testTable(t), false) + if priced.CostUSD == nil || *priced.CostUSD != 9.0 { + t.Fatalf("cost = %v, want 9.0 (reported kept)", priced.CostUSD) + } + if priced.CostSource != types.CostSourceReported { + t.Errorf("source = %q, want reported", priced.CostSource) + } +} + +func TestPriceUsage_NilUsage(t *testing.T) { + t.Parallel() + priced, buckets := PriceUsage(nil, "test-a", testTable(t), false) + if priced != nil || buckets != nil { + t.Errorf("got (%v, %v), want (nil, nil)", priced, buckets) + } +} + +// TestCalculateUsageWithCost_ClaudeFastVariantPricedAtPremium proves the +// claudecode fast-mode buckets (Model "-fast") price at the fast premium +// through the REAL embedded table: claude-opus-4-8-fast bills at 10/50, double +// the standard claude-opus-4-8 5/25. The bucket keys themselves are produced by +// claudecode.CalculateModelUsage (covered in the claudecode package, which must +// not import pricing); this is the priced half of that feature. +func TestCalculateUsageWithCost_ClaudeFastVariantPricedAtPremium(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + if err != nil { + t.Fatalf("LoadTable: %v", err) + } + + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000, OutputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "claude-opus-4-8", TokenUsage: TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + {Model: "claude-opus-4-8-fast", TokenUsage: TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", table, "", false) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // standard: 1M@$5 + 1M@$25 = 30.0 + if buckets[0].Model != "claude-opus-4-8" || buckets[0].TokenUsage.CostUSD == nil || *buckets[0].TokenUsage.CostUSD != 30.0 { + t.Errorf("standard bucket = %+v (cost %v), want claude-opus-4-8 @ 30.0", buckets[0].Model, buckets[0].TokenUsage.CostUSD) + } + // fast: 1M@$10 + 1M@$50 = 60.0 (double the standard) + if buckets[1].Model != "claude-opus-4-8-fast" || buckets[1].TokenUsage.CostUSD == nil || *buckets[1].TokenUsage.CostUSD != 60.0 { + t.Errorf("fast bucket = %+v (cost %v), want claude-opus-4-8-fast @ 60.0", buckets[1].Model, buckets[1].TokenUsage.CostUSD) + } + if flat.CostUSD == nil || *flat.CostUSD != 90.0 { + t.Fatalf("flat cost = %v, want 90.0", flat.CostUSD) + } + if flat.CostSource != types.CostSourceEstimated { + t.Errorf("flat source = %q, want estimated", flat.CostSource) + } +} + +// TestRemainderBucket_APICallCountOnlyShortfallNoRemainder proves a shortfall +// that is only an APICallCount delta (all four token fields fully covered by +// the per-model buckets) does not produce a spurious zero-token remainder +// bucket. +func TestRemainderBucket_APICallCountOnlyShortfallNoRemainder(t *testing.T) { + t.Parallel() + flat := &TokenUsage{InputTokens: 1_000_000, APICallCount: 5} + buckets := []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, APICallCount: 2}}, + } + + _, ok := remainderBucket(flat, buckets, "test-a") + if ok { + t.Fatalf("remainderBucket returned a bucket for an APICallCount-only shortfall, want none") + } +} + +// TestRemainderBucket_RealTokenShortfallStillReturned proves a genuine token +// shortfall (not just an APICallCount delta) still produces a remainder +// bucket, guarding against an overly broad fix to the APICallCount exclusion. +func TestRemainderBucket_RealTokenShortfallStillReturned(t *testing.T) { + t.Parallel() + flat := &TokenUsage{InputTokens: 2_000_000, APICallCount: 5} + buckets := []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, APICallCount: 2}}, + } + + rem, ok := remainderBucket(flat, buckets, "test-a") + if !ok { + t.Fatalf("remainderBucket returned no bucket for a real token shortfall") + } + if rem.TokenUsage.InputTokens != 1_000_000 { + t.Errorf("remainder InputTokens = %d, want 1_000_000", rem.TokenUsage.InputTokens) + } +} + +// TestFoldBucketCost_EmptySourceWithCostNormalizesToEstimated proves a bucket +// carrying a non-nil CostUSD but an empty CostSource (currently unreachable via +// priceBuckets, but defended against as a boundary case) folds into a total +// labeled CostSourceEstimated rather than an unlabeled "" source. +func TestFoldBucketCost_EmptySourceWithCostNormalizesToEstimated(t *testing.T) { + t.Parallel() + buckets := []types.ModelUsage{ + {Model: "test-a", TokenUsage: TokenUsage{InputTokens: 1_000_000, CostUSD: ptrFloat(5.0), CostSource: ""}}, + } + + cost, source := foldBucketCost(buckets) + if cost == nil || *cost != 5.0 { + t.Fatalf("cost = %v, want 5.0", cost) + } + if source != types.CostSourceEstimated { + t.Errorf("source = %q, want estimated", source) + } +} diff --git a/cmd/entire/cli/agent/token_usage_tier_test.go b/cmd/entire/cli/agent/token_usage_tier_test.go new file mode 100644 index 0000000000..016ae8a0de --- /dev/null +++ b/cmd/entire/cli/agent/token_usage_tier_test.go @@ -0,0 +1,146 @@ +package agent + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/pricing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestApplyTierVariant covers the bucket retargeting that keeps a +// tier-suffixed fallback model effective when an agent's +// ModelUsageCalculator produces buckets under raw transcript ids. +func TestApplyTierVariant(t *testing.T) { + t.Parallel() + + mk := func(models ...string) []types.ModelUsage { + buckets := make([]types.ModelUsage, 0, len(models)) + for _, m := range models { + buckets = append(buckets, types.ModelUsage{Model: m}) + } + return buckets + } + models := func(buckets []types.ModelUsage) []string { + out := make([]string, 0, len(buckets)) + for _, b := range buckets { + out = append(out, b.Model) + } + return out + } + + tests := []struct { + name string + buckets []types.ModelUsage + fallback string + want []string + }{ + { + name: "matching base is retargeted onto the variant", + buckets: mk("gpt-5.5"), + fallback: "gpt-5.5-priority", + want: []string{"gpt-5.5-priority"}, + }, + { + name: "only the matching bucket moves in a mixed session", + buckets: mk("gpt-5.5", "gpt-5.4", "", "gpt-5.5-priority"), + fallback: "gpt-5.5-priority", + want: []string{"gpt-5.5-priority", "gpt-5.4", "", "gpt-5.5-priority"}, + }, + { + name: "unsuffixed fallback is a no-op", + buckets: mk("gpt-5.5", ""), + fallback: "gpt-5.5", + want: []string{"gpt-5.5", ""}, + }, + { + name: "empty fallback is a no-op", + buckets: mk("gpt-5.5"), + fallback: "", + want: []string{"gpt-5.5"}, + }, + { + name: "bare suffix has no base and is a no-op", + buckets: mk("gpt-5.5", ""), + fallback: "-priority", + want: []string{"gpt-5.5", ""}, + }, + { + name: "non-matching base leaves everything raw", + buckets: mk("claude-opus-4-8"), + fallback: "gpt-5.5-priority", + want: []string{"claude-opus-4-8"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + applyTierVariant(tt.buckets, tt.fallback) + assert.Equal(t, tt.want, models(tt.buckets)) + }) + } +} + +// TestResolveTierFallback covers the downgrade of a tier-suffixed fallback +// model to its base id when the pricing table has no entry for the variant, so +// the session prices as an honest standard-rate undercount instead of going +// entirely unpriced under a premium-claiming id. +func TestResolveTierFallback(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + require.NoError(t, err) + + tests := []struct { + name string + fallback string + table *pricing.Table + want string + }{ + { + name: "priceable variant is kept", + fallback: "gpt-5.5-priority", + table: table, + want: "gpt-5.5-priority", + }, + { + name: "unpriceable variant downgrades to its base", + fallback: "gpt-5.3-codex-priority", + table: table, + want: "gpt-5.3-codex", + }, + { + name: "unsuffixed model passes through", + fallback: "gpt-5.3-codex", + table: table, + want: "gpt-5.3-codex", + }, + { + name: "nil table keeps the base id", + fallback: "gpt-5.5-priority", + table: nil, + want: "gpt-5.5", + }, + { + name: "bare suffix passes through", + fallback: "-priority", + table: table, + want: "-priority", + }, + { + name: "empty model passes through", + fallback: "", + table: table, + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, tt.want, resolveTierFallback(tt.fallback, tt.table)) + }) + } +} diff --git a/cmd/entire/cli/agent/token_usage_variant_downgrade_test.go b/cmd/entire/cli/agent/token_usage_variant_downgrade_test.go new file mode 100644 index 0000000000..025210afe6 --- /dev/null +++ b/cmd/entire/cli/agent/token_usage_variant_downgrade_test.go @@ -0,0 +1,197 @@ +package agent + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/pricing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestDowngradedVariantBase covers the priceability rule that decides whether a +// variant-suffixed model id is reverted to its base for pricing: downgrade only +// when the base is a strictly better (priceable) target than the unpriceable +// variant, or when the table is nil. +func TestDowngradedVariantBase(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + require.NoError(t, err) + + tests := []struct { + name string + model string + table *pricing.Table + wantBase string + wantOK bool + }{ + {"unpriceable fast variant downgrades to base", "claude-fable-5-fast", table, "claude-fable-5", true}, + {"priceable fast variant is kept", "claude-opus-4-8-fast", table, "", false}, + {"unpriceable priority variant downgrades to base", "gpt-5.3-codex-priority", table, "gpt-5.3-codex", true}, + {"priceable priority variant is kept", "gpt-5.5-priority", table, "", false}, + {"no variant suffix is a no-op", "claude-fable-5", table, "", false}, + {"variant with unpriceable base is kept truthful", "totally-unknown-model-fast", table, "", false}, + {"bare suffix has no base", "-fast", table, "", false}, + {"empty model is a no-op", "", table, "", false}, + {"nil table downgrades the variant", "claude-fable-5-fast", nil, "claude-fable-5", true}, + {"nil table downgrades a priority variant too", "gpt-5.5-priority", nil, "gpt-5.5", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + base, ok := downgradedVariantBase(tt.model, tt.table) + assert.Equal(t, tt.wantOK, ok) + assert.Equal(t, tt.wantBase, base) + }) + } +} + +// TestDowngradeUnpriceableVariants covers the in-place bucket rekeying: only +// unpriceable-variant buckets move to their base; priceable variants and plain +// ids stay put. +func TestDowngradeUnpriceableVariants(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + require.NoError(t, err) + + models := func(buckets []types.ModelUsage) []string { + out := make([]string, 0, len(buckets)) + for _, b := range buckets { + out = append(out, b.Model) + } + return out + } + mk := func(ids ...string) []types.ModelUsage { + buckets := make([]types.ModelUsage, 0, len(ids)) + for _, id := range ids { + buckets = append(buckets, types.ModelUsage{Model: id}) + } + return buckets + } + + t.Run("only the unpriceable fast bucket moves", func(t *testing.T) { + t.Parallel() + buckets := mk("claude-opus-4-8-fast", "claude-fable-5-fast", "claude-fable-5", "") + downgradeUnpriceableVariants(buckets, table) + assert.Equal(t, []string{"claude-opus-4-8-fast", "claude-fable-5", "claude-fable-5", ""}, models(buckets)) + }) + + t.Run("nil table downgrades every variant", func(t *testing.T) { + t.Parallel() + buckets := mk("claude-fable-5-fast", "gpt-5.5-priority", "claude-opus-4-8") + downgradeUnpriceableVariants(buckets, nil) + assert.Equal(t, []string{"claude-fable-5", "gpt-5.5", "claude-opus-4-8"}, models(buckets)) + }) +} + +// TestCalculateUsageWithCost_UnpriceableFastDowngradesToBase is the regression +// for the finding: a fast-mode turn on a model with no published -fast rate must +// estimate at the base rate under the base id, not go unpriced (nil) under the +// unpriceable -fast id. claude-fable-5 has no -fast entry; its base bills $10/$50. +func TestCalculateUsageWithCost_UnpriceableFastDowngradesToBase(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + require.NoError(t, err) + + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}, + buckets: []types.ModelUsage{ + {Model: "claude-fable-5-fast", TokenUsage: TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", table, "", false) + require.NoError(t, err) + require.Len(t, buckets, 1) + // Rekeyed to the base id and priced at the base rate, not left unpriced. + assert.Equal(t, "claude-fable-5", buckets[0].Model) + require.NotNil(t, buckets[0].TokenUsage.CostUSD) + assert.InDelta(t, 60.0, *buckets[0].TokenUsage.CostUSD, 1e-9) // 1M@$10 + 1M@$50 + assert.Equal(t, types.CostSourceEstimated, buckets[0].TokenUsage.CostSource) + require.NotNil(t, flat.CostUSD) + assert.InDelta(t, 60.0, *flat.CostUSD, 1e-9) + assert.Equal(t, types.CostSourceEstimated, flat.CostSource) +} + +// TestCalculateUsageWithCost_MixedFastPriceableAndUnpriceable proves a session +// mixing a priceable fast model with an unpriceable one handles each correctly +// and conserves the total: opus-4-8-fast keeps its premium, fable-5-fast falls +// back to the fable base. +func TestCalculateUsageWithCost_MixedFastPriceableAndUnpriceable(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + require.NoError(t, err) + + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000, OutputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "claude-opus-4-8-fast", TokenUsage: TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + {Model: "claude-fable-5-fast", TokenUsage: TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", table, "", false) + require.NoError(t, err) + require.Len(t, buckets, 2) + + // Priceable fast variant keeps its premium id and rate: 1M@$10 + 1M@$50 = 60. + assert.Equal(t, "claude-opus-4-8-fast", buckets[0].Model) + require.NotNil(t, buckets[0].TokenUsage.CostUSD) + assert.InDelta(t, 60.0, *buckets[0].TokenUsage.CostUSD, 1e-9) + + // Unpriceable fast variant downgrades to the base id and rate: also $60 here. + assert.Equal(t, "claude-fable-5", buckets[1].Model) + require.NotNil(t, buckets[1].TokenUsage.CostUSD) + assert.InDelta(t, 60.0, *buckets[1].TokenUsage.CostUSD, 1e-9) + + // Conservation: flat cost is the sum, all estimated. + require.NotNil(t, flat.CostUSD) + assert.InDelta(t, 120.0, *flat.CostUSD, 1e-9) + assert.Equal(t, types.CostSourceEstimated, flat.CostSource) +} + +// TestCalculateUsageWithCost_DowngradedFastCollidesWithRemainderConserved proves +// the leave-duplicate choice: when a downgraded fast bucket ends up under the same +// base id as the remainder bucket, both are kept (as same-model buckets already +// are today) and the total is conserved — foldBucketCost sums them and downstream +// folds by model key. +func TestCalculateUsageWithCost_DowngradedFastCollidesWithRemainderConserved(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + require.NoError(t, err) + + // flat exceeds the single calculator bucket, so a remainder bucket is emitted + // under fallbackModel "claude-fable-5"; the calculator's fast bucket downgrades + // onto that same base id. + ag := &fakeModelUsageAgent{ + flat: &TokenUsage{InputTokens: 2_000_000, OutputTokens: 2_000_000}, + buckets: []types.ModelUsage{ + {Model: "claude-fable-5-fast", TokenUsage: TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + }, + } + + flat, buckets, err := CalculateUsageWithCost(ag, nil, 0, "", table, "claude-fable-5", false) + require.NoError(t, err) + require.Len(t, buckets, 2) + + var totalIn, totalOut int + for _, b := range buckets { + assert.Equal(t, "claude-fable-5", b.Model, "both buckets under the base id after downgrade") + require.NotNil(t, b.TokenUsage.CostUSD) + assert.InDelta(t, 60.0, *b.TokenUsage.CostUSD, 1e-9) + totalIn += b.TokenUsage.InputTokens + totalOut += b.TokenUsage.OutputTokens + } + assert.Equal(t, 2_000_000, totalIn, "input tokens conserved across the duplicate buckets") + assert.Equal(t, 2_000_000, totalOut, "output tokens conserved across the duplicate buckets") + + require.NotNil(t, flat.CostUSD) + assert.InDelta(t, 120.0, *flat.CostUSD, 1e-9) + assert.Equal(t, types.CostSourceEstimated, flat.CostSource) +} diff --git a/cmd/entire/cli/agent/types.go b/cmd/entire/cli/agent/types.go index 741c13abd4..a81cdaa6b6 100644 --- a/cmd/entire/cli/agent/types.go +++ b/cmd/entire/cli/agent/types.go @@ -51,3 +51,8 @@ type SessionChange struct { // contract can reference it without importing the full agent package. The // alias keeps existing agent.TokenUsage references working. type TokenUsage = types.TokenUsage + +// ModelUsage pairs a model identifier with its token usage. Like TokenUsage it +// lives in the leaf agent/types package; the alias lets agent-package callers +// reference agent.ModelUsage. +type ModelUsage = types.ModelUsage diff --git a/cmd/entire/cli/agent/types/token_usage.go b/cmd/entire/cli/agent/types/token_usage.go index 5dacbb0d30..fbc150288d 100644 --- a/cmd/entire/cli/agent/types/token_usage.go +++ b/cmd/entire/cli/agent/types/token_usage.go @@ -1,5 +1,17 @@ package types +// Cost source provenance values for TokenUsage.CostSource. +const ( + // CostSourceReported means the cost came directly from an agent-reported + // figure (authoritative). + CostSourceReported = "reported" + // CostSourceEstimated means the cost was derived from a pricing estimate. + CostSourceEstimated = "estimated" + // CostSourceMixed means the aggregated cost combined differing provenances + // (e.g. some reported, some estimated). + CostSourceMixed = "mixed" +) + // TokenUsage represents aggregated token usage for a checkpoint. // This is agent-agnostic and can be populated by any agent that tracks token usage. type TokenUsage struct { @@ -15,6 +27,109 @@ type TokenUsage struct { APICallCount int `json:"api_call_count"` // SubagentTokens contains token usage from spawned subagents (if any) SubagentTokens *TokenUsage `json:"subagent_tokens,omitempty"` + // CostUSD is the total cost in USD for this usage, when known. + // nil means unknown (no agent-reported cost and no estimate) — never treat as $0. + CostUSD *float64 `json:"cost_usd,omitempty"` + // CostSource records provenance: CostSourceReported, CostSourceEstimated, or CostSourceMixed. + CostSource string `json:"cost_source,omitempty"` +} + +// ModelUsage pairs a model identifier with its token usage. It is used by +// per-model cost accounting (later chunks); harmless to carry now. +type ModelUsage struct { + Model string `json:"model"` + TokenUsage TokenUsage `json:"token_usage"` +} + +// AddCostUSD sums two optional cost values, returning a NEW pointer that never +// aliases either input. +// +// - nil + nil -> nil +// - nil + x -> copy of x +// - x + y -> x + y +func AddCostUSD(a, b *float64) *float64 { + if a == nil && b == nil { + return nil + } + var sum float64 + if a != nil { + sum += *a + } + if b != nil { + sum += *b + } + return &sum +} + +// MergeCostSource merges two cost-source provenance labels. A side whose cost +// is nil contributes no source (its label is ignored, since it carries no +// cost). Both effective sides empty -> ""; equal non-empty -> that value; +// differing non-empty -> CostSourceMixed. +func MergeCostSource(a, b string, aCost, bCost *float64) string { + if aCost == nil { + a = "" + } + if bCost == nil { + b = "" + } + switch { + case a == "" && b == "": + return "" + case a == "": + return b + case b == "": + return a + case a == b: + return a + default: + return CostSourceMixed + } +} + +// MergeCostSourceUsages merges the cost-source provenance of two usages. It +// behaves like MergeCostSource over the two sides' (CostSource, CostUSD) pairs +// but additionally treats a token-bearing side with no cost as a distinct +// provenance: when exactly one side carries a non-nil cost and the other side +// has any nonzero token field with a nil cost, the merge is CostSourceMixed. +// This mirrors foldBucketCost's partial-coverage rule, so a priced usage merged +// with an unpriced-but-token-bearing usage honestly reports mixed coverage +// instead of masquerading as fully priced. Either usage may be nil (contributes +// nothing). APICallCount and SubagentTokens are not treated as billable tokens. +func MergeCostSourceUsages(a, b *TokenUsage) string { + var aSource, bSource string + var aCost, bCost *float64 + if a != nil { + aSource, aCost = a.CostSource, a.CostUSD + } + if b != nil { + bSource, bCost = b.CostSource, b.CostUSD + } + + // Exactly one side priced while the other carries unpriced tokens -> mixed. + if aCost != nil && bCost == nil && usageHasBillableTokens(b) { + return CostSourceMixed + } + if bCost != nil && aCost == nil && usageHasBillableTokens(a) { + return CostSourceMixed + } + return MergeCostSource(aSource, bSource, aCost, bCost) +} + +// usageHasBillableTokens reports whether u carries any nonzero billable token +// count (input, cache-creation, cache-read, or output) at any depth. It recurses +// into the SubagentTokens subtree so a step whose tokens live entirely in an +// unpriced subagent subtree still counts as token-bearing — otherwise +// MergeCostSourceUsages would miss the mixed-coverage case and report a merged +// cost as fully "estimated" instead of "mixed". Mirrors hasTokenUsageData / +// flattenTokenUsage, which also recurse. It is nil-safe. +func usageHasBillableTokens(u *TokenUsage) bool { + if u == nil { + return false + } + if u.InputTokens != 0 || u.CacheCreationTokens != 0 || u.CacheReadTokens != 0 || u.OutputTokens != 0 { + return true + } + return usageHasBillableTokens(u.SubagentTokens) } // AddTokenUsage returns the sum of a and b, recursing into subagent usage. diff --git a/cmd/entire/cli/agent/types/token_usage_test.go b/cmd/entire/cli/agent/types/token_usage_test.go index 54f6cefc1f..fcd3b2d7f0 100644 --- a/cmd/entire/cli/agent/types/token_usage_test.go +++ b/cmd/entire/cli/agent/types/token_usage_test.go @@ -1,6 +1,210 @@ package types -import "testing" +import ( + "encoding/json" + "strings" + "testing" +) + +func costPtr(v float64) *float64 { return &v } + +func TestAddCostUSD_NilNil(t *testing.T) { + t.Parallel() + if got := AddCostUSD(nil, nil); got != nil { + t.Fatalf("AddCostUSD(nil, nil) = %v, want nil", got) + } +} + +func TestAddCostUSD_NilVal(t *testing.T) { + t.Parallel() + b := costPtr(1.25) + got := AddCostUSD(nil, b) + if got == nil || *got != 1.25 { + t.Fatalf("AddCostUSD(nil, 1.25) = %v, want 1.25", got) + } + // Must be a copy, not an alias of the input. + if got == b { + t.Fatal("AddCostUSD(nil, b) aliased the input pointer") + } + *got = 99 + if *b != 1.25 { + t.Fatalf("mutating result mutated input b: b=%v", *b) + } + + // Symmetric: value on the a side. + a := costPtr(2.5) + got = AddCostUSD(a, nil) + if got == nil || *got != 2.5 { + t.Fatalf("AddCostUSD(2.5, nil) = %v, want 2.5", got) + } + if got == a { + t.Fatal("AddCostUSD(a, nil) aliased the input pointer") + } +} + +func TestAddCostUSD_ValVal(t *testing.T) { + t.Parallel() + a := costPtr(1.5) + b := costPtr(2.25) + got := AddCostUSD(a, b) + if got == nil || *got != 3.75 { + t.Fatalf("AddCostUSD(1.5, 2.25) = %v, want 3.75", got) + } + if got == a || got == b { + t.Fatal("AddCostUSD returned an aliased input pointer") + } + // Mutating result must not touch inputs. + *got = 0 + if *a != 1.5 || *b != 2.25 { + t.Fatalf("inputs mutated: a=%v b=%v", *a, *b) + } +} + +func TestMergeCostSource_Matrix(t *testing.T) { + t.Parallel() + nz := costPtr(1) // any non-nil cost + cases := []struct { + name string + a, b string + aCost, bCost *float64 + want string + }{ + {"both empty sources", "", "", nz, nz, ""}, + {"both nil cost non-empty source", CostSourceReported, CostSourceEstimated, nil, nil, ""}, + {"reported+reported", CostSourceReported, CostSourceReported, nz, nz, CostSourceReported}, + {"reported+estimated -> mixed", CostSourceReported, CostSourceEstimated, nz, nz, CostSourceMixed}, + {"estimated + empty(nil cost) -> estimated", CostSourceEstimated, "", nz, nil, CostSourceEstimated}, + {"empty(nil cost) + reported -> reported", "", CostSourceReported, nil, nz, CostSourceReported}, + {"reported + estimated(nil cost) ignores b -> reported", CostSourceReported, CostSourceEstimated, nz, nil, CostSourceReported}, + {"estimated(nil cost) + reported ignores a -> reported", CostSourceEstimated, CostSourceReported, nil, nz, CostSourceReported}, + {"mixed + reported -> mixed", CostSourceMixed, CostSourceReported, nz, nz, CostSourceMixed}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := MergeCostSource(tc.a, tc.b, tc.aCost, tc.bCost); got != tc.want { + t.Fatalf("MergeCostSource(%q,%q,%v,%v) = %q, want %q", + tc.a, tc.b, tc.aCost, tc.bCost, got, tc.want) + } + }) + } +} + +func TestMergeCostSourceUsages_Matrix(t *testing.T) { + t.Parallel() + priced := func(source string) *TokenUsage { + return &TokenUsage{InputTokens: 100, CostUSD: costPtr(0.5), CostSource: source} + } + tokensNoCost := &TokenUsage{InputTokens: 200} // token-bearing, unpriced + cacheTokensNoCost := &TokenUsage{CacheReadTokens: 200} // token-bearing (cache), unpriced + emptyNoCost := &TokenUsage{} // no tokens, no cost + apiOnlyNoCost := &TokenUsage{APICallCount: 3} // api count is not billable tokens + estimatedNoCost := &TokenUsage{CostSource: CostSourceEstimated} // label but nil cost, no tokens + + cases := []struct { + name string + a, b *TokenUsage + want string + }{ + {"priced + unpriced-with-tokens -> mixed", priced(CostSourceEstimated), tokensNoCost, CostSourceMixed}, + {"unpriced-with-tokens + priced -> mixed (order)", tokensNoCost, priced(CostSourceEstimated), CostSourceMixed}, + {"priced + unpriced-with-cache-tokens -> mixed", priced(CostSourceReported), cacheTokensNoCost, CostSourceMixed}, + {"priced + empty usage -> estimated", priced(CostSourceEstimated), emptyNoCost, CostSourceEstimated}, + {"priced + api-only usage -> reported (api not billable)", priced(CostSourceReported), apiOnlyNoCost, CostSourceReported}, + {"priced + nil -> that source", priced(CostSourceEstimated), nil, CostSourceEstimated}, + {"reported + estimated (both priced) -> mixed", priced(CostSourceReported), priced(CostSourceEstimated), CostSourceMixed}, + {"reported + reported (both priced) -> reported", priced(CostSourceReported), priced(CostSourceReported), CostSourceReported}, + {"empty + empty -> empty", emptyNoCost, emptyNoCost, ""}, + {"tokens-no-cost + estimated-label-no-cost -> empty", tokensNoCost, estimatedNoCost, ""}, + {"both nil -> empty", nil, nil, ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := MergeCostSourceUsages(tc.a, tc.b); got != tc.want { + t.Fatalf("MergeCostSourceUsages = %q, want %q", got, tc.want) + } + }) + } +} + +// A priced usage merged with a step whose tokens live ENTIRELY in an unpriced +// SubagentTokens subtree must report mixed coverage: the subtree is billable but +// carries no cost, so the merged provenance is "mixed", not "estimated". +// Mutation check: a non-recursive usageHasBillableTokens sees only the top-level +// scalars (all zero here) and misses the subtree, so the merge would report +// "estimated". +func TestMergeCostSourceUsages_SubagentOnlyTokensAreBillable(t *testing.T) { + t.Parallel() + priced := &TokenUsage{InputTokens: 100, CostUSD: costPtr(1.0), CostSource: CostSourceEstimated} + subagentOnly := &TokenUsage{SubagentTokens: &TokenUsage{InputTokens: 50}} + + if got := MergeCostSourceUsages(priced, subagentOnly); got != CostSourceMixed { + t.Fatalf("MergeCostSourceUsages(priced, subagent-only) = %q, want mixed", got) + } + // Order-independent. + if got := MergeCostSourceUsages(subagentOnly, priced); got != CostSourceMixed { + t.Fatalf("MergeCostSourceUsages(subagent-only, priced) = %q, want mixed", got) + } +} + +func TestTokenUsage_JSON_NilCostOmitted(t *testing.T) { + t.Parallel() + u := TokenUsage{InputTokens: 10, OutputTokens: 5} + data, err := json.Marshal(u) + if err != nil { + t.Fatalf("marshal: %v", err) + } + s := string(data) + if strings.Contains(s, "cost_usd") { + t.Fatalf("expected no cost_usd key for nil cost, got %s", s) + } + if strings.Contains(s, "cost_source") { + t.Fatalf("expected no cost_source key for empty source, got %s", s) + } +} + +func TestTokenUsage_JSON_RoundTripWithCost(t *testing.T) { + t.Parallel() + u := TokenUsage{ + InputTokens: 10, + OutputTokens: 5, + CostUSD: costPtr(0.42), + CostSource: CostSourceReported, + } + data, err := json.Marshal(u) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if !strings.Contains(string(data), "cost_usd") { + t.Fatalf("expected cost_usd key, got %s", data) + } + var got TokenUsage + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if got.CostUSD == nil || *got.CostUSD != 0.42 { + t.Fatalf("CostUSD round-trip failed: %v", got.CostUSD) + } + if got.CostSource != CostSourceReported { + t.Fatalf("CostSource round-trip failed: %q", got.CostSource) + } +} + +func TestTokenUsage_JSON_LegacyWithoutCostFields(t *testing.T) { + t.Parallel() + legacy := `{"input_tokens":10,"cache_creation_tokens":0,"cache_read_tokens":0,"output_tokens":5,"api_call_count":1}` + var got TokenUsage + if err := json.Unmarshal([]byte(legacy), &got); err != nil { + t.Fatalf("unmarshal legacy: %v", err) + } + if got.CostUSD != nil { + t.Fatalf("legacy CostUSD = %v, want nil", got.CostUSD) + } + if got.CostSource != "" { + t.Fatalf("legacy CostSource = %q, want empty", got.CostSource) + } +} func TestAddTokenUsage(t *testing.T) { t.Parallel() diff --git a/cmd/entire/cli/checkpoint/aggregate_token_cost_test.go b/cmd/entire/cli/checkpoint/aggregate_token_cost_test.go new file mode 100644 index 0000000000..b5b3592e81 --- /dev/null +++ b/cmd/entire/cli/checkpoint/aggregate_token_cost_test.go @@ -0,0 +1,83 @@ +package checkpoint + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" +) + +func costPtr(v float64) *float64 { return &v } + +// aggregateTokenUsage must sum CostUSD and merge CostSource across the two +// operands. (It intentionally does not recurse SubagentTokens — pre-existing +// behavior unchanged by the cost plumbing.) +func TestAggregateTokenUsage_SumsCost(t *testing.T) { + t.Parallel() + a := &agent.TokenUsage{InputTokens: 1, CostUSD: costPtr(0.75), CostSource: types.CostSourceReported} + b := &agent.TokenUsage{InputTokens: 2, CostUSD: costPtr(0.25), CostSource: types.CostSourceReported} + got := aggregateTokenUsage(a, b) + if got.CostUSD == nil || *got.CostUSD != 1.0 { + t.Fatalf("CostUSD = %v, want 1.0", got.CostUSD) + } + if got.CostSource != types.CostSourceReported { + t.Fatalf("CostSource = %q, want reported", got.CostSource) + } +} + +func TestAggregateTokenUsage_MixedSource(t *testing.T) { + t.Parallel() + a := &agent.TokenUsage{CostUSD: costPtr(1), CostSource: types.CostSourceReported} + b := &agent.TokenUsage{CostUSD: costPtr(1), CostSource: types.CostSourceEstimated} + got := aggregateTokenUsage(a, b) + if got.CostSource != types.CostSourceMixed { + t.Fatalf("CostSource = %q, want mixed", got.CostSource) + } +} + +func TestAggregateTokenUsage_NilCostSideIgnoredForSource(t *testing.T) { + t.Parallel() + // b has an estimated label but nil cost, so it contributes no source. + a := &agent.TokenUsage{CostUSD: costPtr(1), CostSource: types.CostSourceReported} + b := &agent.TokenUsage{CostSource: types.CostSourceEstimated} + got := aggregateTokenUsage(a, b) + if got.CostUSD == nil || *got.CostUSD != 1 { + t.Fatalf("CostUSD = %v, want 1", got.CostUSD) + } + if got.CostSource != types.CostSourceReported { + t.Fatalf("CostSource = %q, want reported", got.CostSource) + } +} + +// A costed session summed with an uncosted but token-bearing session must fold +// to "mixed": the total cost is the one known figure, but coverage is partial. +func TestAggregateTokenUsage_PricedPlusUnpricedTokensMixed(t *testing.T) { + t.Parallel() + sessionA := &agent.TokenUsage{InputTokens: 1000, CostUSD: costPtr(0.50), CostSource: types.CostSourceEstimated} + sessionB := &agent.TokenUsage{InputTokens: 200} // token-bearing, no cost + got := aggregateTokenUsage(sessionA, sessionB) + if got.CostUSD == nil || *got.CostUSD != 0.50 { + t.Fatalf("CostUSD = %v, want 0.50", got.CostUSD) + } + if got.CostSource != types.CostSourceMixed { + t.Fatalf("CostSource = %q, want mixed", got.CostSource) + } + // Order must not matter. + rev := aggregateTokenUsage(sessionB, sessionA) + if rev.CostSource != types.CostSourceMixed { + t.Fatalf("reversed CostSource = %q, want mixed", rev.CostSource) + } +} + +func TestAggregateTokenUsage_BothNilCost(t *testing.T) { + t.Parallel() + a := &agent.TokenUsage{InputTokens: 1} + b := &agent.TokenUsage{InputTokens: 2} + got := aggregateTokenUsage(a, b) + if got.CostUSD != nil { + t.Fatalf("CostUSD = %v, want nil", got.CostUSD) + } + if got.CostSource != "" { + t.Fatalf("CostSource = %q, want empty", got.CostSource) + } +} diff --git a/cmd/entire/cli/checkpoint/cost_not_persisted_test.go b/cmd/entire/cli/checkpoint/cost_not_persisted_test.go new file mode 100644 index 0000000000..f9a53e3a09 --- /dev/null +++ b/cmd/entire/cli/checkpoint/cost_not_persisted_test.go @@ -0,0 +1,125 @@ +package checkpoint + +import ( + "context" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" + "github.com/entireio/cli/redact" +) + +func fptr(v float64) *float64 { return &v } + +// TestWriteCommitted_DoesNotPersistCost is the CLI-side cost-ownership invariant: +// a checkpoint write must persist the full token breakdown the platform prices +// from (flat counts, subagent subtree, and per-model buckets with the four token +// fields + model id) but MUST NOT persist any cost — cost_usd/cost_source are +// nil/empty at every level (flat, subagent, per-model) in both the session +// metadata and the aggregated root summary, even when the write request carries +// cost. entire-api owns server-side pricing from these tokens. +func TestWriteCommitted_DoesNotPersistCost(t *testing.T) { + repo, _ := setupBranchTestRepo(t) + store := NewGitStore(repo, DefaultV1Refs()) + cpID := id.MustCheckpointID("c05700000001") + + err := store.Write(context.Background(), Session{ + CheckpointID: cpID, + SessionID: "cost-session", + Strategy: "manual-commit", + Transcript: redact.AlreadyRedacted([]byte(`{"message":"x"}`)), + Model: "claude-opus-4-8", + TokenUsage: &agent.TokenUsage{ + InputTokens: 1000, + OutputTokens: 200, + CacheReadTokens: 300, + CacheCreationTokens: 40, + APICallCount: 5, + CostUSD: fptr(1.23), + CostSource: types.CostSourceEstimated, + SubagentTokens: &agent.TokenUsage{ + InputTokens: 10, + OutputTokens: 2, + CostUSD: fptr(0.05), + CostSource: types.CostSourceReported, + }, + }, + ModelUsage: []types.ModelUsage{ + {Model: "claude-opus-4-8", TokenUsage: agent.TokenUsage{ + InputTokens: 1000, + OutputTokens: 200, + CacheReadTokens: 300, + CacheCreationTokens: 40, + CostUSD: fptr(1.23), + CostSource: types.CostSourceEstimated, + }}, + }, + AuthorName: "Test Author", + AuthorEmail: "test@example.com", + }) + if err != nil { + t.Fatalf("Write() error = %v", err) + } + + meta, err := store.ReadSessionMetadata(context.Background(), cpID, 0) + if err != nil { + t.Fatalf("ReadSessionMetadata() error = %v", err) + } + assertUsageTokensNoCost(t, "session flat", meta.TokenUsage, 1000, 200, 300, 40) + if meta.TokenUsage.APICallCount != 5 { + t.Errorf("session flat APICallCount = %d, want 5", meta.TokenUsage.APICallCount) + } + + // Subagent subtree: token counts preserved, cost stripped. + if meta.TokenUsage.SubagentTokens == nil { + t.Fatal("subagent tokens must be persisted") + } + assertUsageTokensNoCost(t, "session subagent", meta.TokenUsage.SubagentTokens, 10, 2, 0, 0) + + // Per-model breakdown: model id + four token fields preserved, cost stripped. + if len(meta.ModelUsage) != 1 { + t.Fatalf("len(meta.ModelUsage) = %d, want 1", len(meta.ModelUsage)) + } + if meta.ModelUsage[0].Model != "claude-opus-4-8" { + t.Errorf("model id = %q, want claude-opus-4-8", meta.ModelUsage[0].Model) + } + mu := meta.ModelUsage[0].TokenUsage + assertUsageTokensNoCost(t, "session per-model", &mu, 1000, 200, 300, 40) + + // Aggregated root summary: token counts preserved, cost stripped. + summary, err := store.Read(context.Background(), cpID) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + if summary == nil || summary.TokenUsage == nil { + t.Fatal("summary token usage must be persisted") + } + assertUsageTokensNoCost(t, "summary flat", summary.TokenUsage, 1000, 200, 300, 40) + for i := range summary.ModelUsage { + if summary.ModelUsage[i].Model == "" { + t.Errorf("summary per-model bucket %d has empty model id", i) + } + b := summary.ModelUsage[i].TokenUsage + if b.CostUSD != nil || b.CostSource != "" { + t.Errorf("summary per-model cost must not be persisted: %+v", b) + } + } +} + +func assertUsageTokensNoCost(t *testing.T, label string, u *agent.TokenUsage, input, output, cacheRead, cacheWrite int) { + t.Helper() + if u == nil { + t.Fatalf("%s: usage must be persisted", label) + } + if u.InputTokens != input || u.OutputTokens != output || u.CacheReadTokens != cacheRead || u.CacheCreationTokens != cacheWrite { + t.Errorf("%s: token counts = in=%d out=%d cr=%d cw=%d, want in=%d out=%d cr=%d cw=%d", + label, u.InputTokens, u.OutputTokens, u.CacheReadTokens, u.CacheCreationTokens, input, output, cacheRead, cacheWrite) + } + if u.CostUSD != nil { + t.Errorf("%s: cost_usd must NOT be persisted, got %v", label, *u.CostUSD) + } + if u.CostSource != "" { + t.Errorf("%s: cost_source must NOT be persisted, got %q", label, u.CostSource) + } +} diff --git a/cmd/entire/cli/checkpoint/fsstore/fsstore.go b/cmd/entire/cli/checkpoint/fsstore/fsstore.go index b1169d1444..eb4ecb2b1f 100644 --- a/cmd/entire/cli/checkpoint/fsstore/fsstore.go +++ b/cmd/entire/cli/checkpoint/fsstore/fsstore.go @@ -310,6 +310,8 @@ func metadataFromWriteOptions(opts cp.WriteOptions) cp.Metadata { // Contract: a zero CreatedAt means "use the current time". createdAt = time.Now() } + // WithoutCost mirrors the git store: no checkpoint backend persists cost + // (entire-api prices server-side); cost is a display-only local estimate. return cp.Metadata{ CheckpointID: opts.CheckpointID, SessionID: opts.SessionID, @@ -338,7 +340,7 @@ func metadataFromWriteOptions(opts cp.WriteOptions) cp.Metadata { ReviewPrompt: opts.ReviewPrompt, InvestigateRunID: opts.InvestigateRunID, InvestigateTopic: opts.InvestigateTopic, - } + }.WithoutCost() } func upsertSession(sessions []storedSession, session storedSession) []storedSession { diff --git a/cmd/entire/cli/checkpoint/merge_model_usage_test.go b/cmd/entire/cli/checkpoint/merge_model_usage_test.go new file mode 100644 index 0000000000..84e875c7f5 --- /dev/null +++ b/cmd/entire/cli/checkpoint/merge_model_usage_test.go @@ -0,0 +1,99 @@ +package checkpoint + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" +) + +// mergeModelUsage is the per-model aggregation primitive behind the root +// CheckpointSummary. Two sessions with an overlapping model must produce a single +// merged bucket (tokens summed, costs folded) plus the union of distinct models, +// sorted by model for deterministic output. +func TestMergeModelUsage_OverlappingModelsMerge(t *testing.T) { + t.Parallel() + + // Session 0's per-model list. + sessionA := []types.ModelUsage{ + {Model: "opus", TokenUsage: types.TokenUsage{InputTokens: 10, OutputTokens: 2, CostUSD: costPtr(0.5), CostSource: types.CostSourceEstimated}}, + {Model: "haiku", TokenUsage: types.TokenUsage{InputTokens: 4, CostUSD: costPtr(0.1), CostSource: types.CostSourceEstimated}}, + } + // Session 1's per-model list: opus overlaps, sonnet is new. + sessionB := []types.ModelUsage{ + {Model: "opus", TokenUsage: types.TokenUsage{InputTokens: 6, OutputTokens: 1, CostUSD: costPtr(0.25), CostSource: types.CostSourceEstimated}}, + {Model: "sonnet", TokenUsage: types.TokenUsage{InputTokens: 8, CostUSD: costPtr(0.2), CostSource: types.CostSourceReported}}, + } + + got := mergeModelUsage(sessionA, sessionB) + + // Deterministic order: haiku, opus, sonnet. + wantOrder := []string{"haiku", "opus", "sonnet"} + if len(got) != len(wantOrder) { + t.Fatalf("len = %d, want %d (%v)", len(got), len(wantOrder), got) + } + for i, model := range wantOrder { + if got[i].Model != model { + t.Fatalf("got[%d].Model = %q, want %q", i, got[i].Model, model) + } + } + + // opus is the merged bucket. + opus := got[1].TokenUsage + if opus.InputTokens != 16 || opus.OutputTokens != 3 { + t.Fatalf("opus tokens = in %d out %d, want in 16 out 3", opus.InputTokens, opus.OutputTokens) + } + if opus.CostUSD == nil || *opus.CostUSD != 0.75 { + t.Fatalf("opus cost = %v, want 0.75", opus.CostUSD) + } + if opus.CostSource != types.CostSourceEstimated { + t.Fatalf("opus source = %q, want estimated", opus.CostSource) + } +} + +// A model priced by one session (reported) but not the other folds to mixed only +// when both sides carry cost; a nil-cost side contributes no source label. +func TestMergeModelUsage_MixedAndNilCostSource(t *testing.T) { + t.Parallel() + + a := []types.ModelUsage{{Model: "opus", TokenUsage: types.TokenUsage{InputTokens: 1, CostUSD: costPtr(1.0), CostSource: types.CostSourceReported}}} + // Same model, estimated cost -> mixed. + b := []types.ModelUsage{{Model: "opus", TokenUsage: types.TokenUsage{InputTokens: 1, CostUSD: costPtr(0.5), CostSource: types.CostSourceEstimated}}} + got := mergeModelUsage(a, b) + if len(got) != 1 { + t.Fatalf("len = %d, want 1", len(got)) + } + if got[0].TokenUsage.CostSource != types.CostSourceMixed { + t.Fatalf("source = %q, want mixed", got[0].TokenUsage.CostSource) + } + if *got[0].TokenUsage.CostUSD != 1.5 { + t.Fatalf("cost = %v, want 1.5", got[0].TokenUsage.CostUSD) + } +} + +// Empty inputs yield nil so the omitempty JSON tag drops model_usage entirely. +func TestMergeModelUsage_EmptyYieldsNil(t *testing.T) { + t.Parallel() + if got := mergeModelUsage(nil, nil); got != nil { + t.Fatalf("nil+nil = %v, want nil", got) + } + // One empty side just returns the other, sorted. + only := []types.ModelUsage{{Model: "opus", TokenUsage: types.TokenUsage{InputTokens: 3}}} + got := mergeModelUsage(nil, only) + if len(got) != 1 || got[0].Model != "opus" || got[0].TokenUsage.InputTokens != 3 { + t.Fatalf("nil+one = %v, want single opus", got) + } +} + +// The merge must not alias the input buckets' cost pointers. +func TestMergeModelUsage_NoAliasing(t *testing.T) { + t.Parallel() + in := []types.ModelUsage{{Model: "opus", TokenUsage: types.TokenUsage{InputTokens: 5, CostUSD: costPtr(0.5)}}} + got := mergeModelUsage(in, nil) + if got[0].TokenUsage.CostUSD == in[0].TokenUsage.CostUSD { + t.Fatal("merged cost pointer aliases input bucket") + } + got[0].TokenUsage.InputTokens = 999 + if in[0].TokenUsage.InputTokens != 5 { + t.Fatalf("mutating merged result mutated input: %d", in[0].TokenUsage.InputTokens) + } +} diff --git a/cmd/entire/cli/checkpoint/persistent.go b/cmd/entire/cli/checkpoint/persistent.go index 6d82820c75..a9c760e526 100644 --- a/cmd/entire/cli/checkpoint/persistent.go +++ b/cmd/entire/cli/checkpoint/persistent.go @@ -260,7 +260,7 @@ func (s *treeWriter) applyAttributionBackfill(ctx context.Context, existing *obj } summary.CombinedAttribution = combinedAttribution - metadataJSON, err := jsonutil.MarshalIndentWithNewline(summary, "", " ") + metadataJSON, err := jsonutil.MarshalIndentWithNewline(summary.WithoutCost(), "", " ") if err != nil { return plumbing.ZeroHash, fmt.Errorf("failed to marshal checkpoint summary: %w", err) } @@ -313,7 +313,7 @@ func (s *treeWriter) applySummaryBackfill(ctx context.Context, existing *object. existingMetadata.Summary = RedactSummary(summary) - metadataJSON, err := jsonutil.MarshalIndentWithNewline(existingMetadata, "", " ") + metadataJSON, err := jsonutil.MarshalIndentWithNewline(existingMetadata.WithoutCost(), "", " ") if err != nil { return plumbing.ZeroHash, "", fmt.Errorf("failed to marshal metadata: %w", err) } @@ -431,7 +431,7 @@ func (s *treeWriter) applyTranscriptBackfill(ctx context.Context, opts UpdateOpt if sess.CompactTranscript != compactPath || sess.AssetsManifest != manifestPath { sess.CompactTranscript = compactPath sess.AssetsManifest = manifestPath - summaryJSON, err := jsonutil.MarshalIndentWithNewline(checkpointSummary, "", " ") + summaryJSON, err := jsonutil.MarshalIndentWithNewline(checkpointSummary.WithoutCost(), "", " ") if err != nil { return plumbing.ZeroHash, fmt.Errorf("failed to marshal checkpoint summary: %w", err) } @@ -752,6 +752,7 @@ func (s *treeWriter) writeSessionToSubdirectory(ctx context.Context, opts WriteO TranscriptLinesAtStart: opts.CheckpointTranscriptStart, // Deprecated: kept for backward compat CompactTranscriptStart: compactTranscriptStart, TokenUsage: opts.TokenUsage, + ModelUsage: opts.ModelUsage, SkillEventsVersion: skillEventsVersion(opts.SkillEvents), SkillEvents: opts.SkillEvents, SessionMetrics: opts.SessionMetrics, @@ -766,7 +767,7 @@ func (s *treeWriter) writeSessionToSubdirectory(ctx context.Context, opts WriteO InvestigateTopic: opts.InvestigateTopic, } - metadataJSON, err := jsonutil.MarshalIndentWithNewline(sessionMetadata, "", " ") + metadataJSON, err := jsonutil.MarshalIndentWithNewline(sessionMetadata.WithoutCost(), "", " ") if err != nil { return filePaths, fmt.Errorf("failed to marshal session metadata: %w", err) } @@ -788,7 +789,7 @@ func (s *treeWriter) writeSessionToSubdirectory(ctx context.Context, opts WriteO // writeCheckpointSummary writes the root-level CheckpointSummary with aggregated statistics. // sessions is the complete sessions array (already built by the caller). func (s *treeWriter) writeCheckpointSummary(opts WriteOptions, basePath string, entries map[string]object.TreeEntry, sessions []SessionFilePaths) error { - checkpointsCount, filesTouched, tokenUsage, err := s.reaggregateFromEntries(basePath, len(sessions), entries) + checkpointsCount, filesTouched, tokenUsage, modelUsage, err := s.reaggregateFromEntries(basePath, len(sessions), entries) if err != nil { return fmt.Errorf("failed to aggregate session stats: %w", err) } @@ -828,13 +829,14 @@ func (s *treeWriter) writeCheckpointSummary(opts WriteOptions, basePath string, FilesTouched: filesTouched, Sessions: sessions, TokenUsage: tokenUsage, + ModelUsage: modelUsage, CombinedAttribution: combinedAttribution, HasReview: hasReview, HasInvestigation: hasInvestigation, Imported: imported, } - metadataJSON, err := jsonutil.MarshalIndentWithNewline(summary, "", " ") + metadataJSON, err := jsonutil.MarshalIndentWithNewline(summary.WithoutCost(), "", " ") if err != nil { return fmt.Errorf("failed to marshal checkpoint summary: %w", err) } @@ -919,28 +921,30 @@ func (s *treeWriter) findSessionIndex(ctx context.Context, basePath string, exis } // reaggregateFromEntries reads all session metadata from the entries map and -// reaggregates CheckpointsCount, FilesTouched, and TokenUsage. -func (s *treeWriter) reaggregateFromEntries(basePath string, sessionCount int, entries map[string]object.TreeEntry) (int, []string, *agent.TokenUsage, error) { +// reaggregates CheckpointsCount, FilesTouched, TokenUsage, and per-model ModelUsage. +func (s *treeWriter) reaggregateFromEntries(basePath string, sessionCount int, entries map[string]object.TreeEntry) (int, []string, *agent.TokenUsage, []types.ModelUsage, error) { var totalCount int var allFiles []string var totalTokens *agent.TokenUsage + var totalModels []types.ModelUsage for i := range sessionCount { path := checkpointSubtreePath(basePath, strconv.Itoa(i), paths.MetadataFileName) entry, exists := entries[path] if !exists { - return 0, nil, nil, fmt.Errorf("session %d metadata not found at %s", i, path) + return 0, nil, nil, nil, fmt.Errorf("session %d metadata not found at %s", i, path) } meta, err := s.readMetadataFromBlob(entry.Hash) if err != nil { - return 0, nil, nil, fmt.Errorf("failed to read session %d metadata: %w", i, err) + return 0, nil, nil, nil, fmt.Errorf("failed to read session %d metadata: %w", i, err) } totalCount += meta.CheckpointsCount allFiles = mergeFilesTouched(allFiles, meta.FilesTouched) totalTokens = aggregateTokenUsage(totalTokens, meta.TokenUsage) + totalModels = mergeModelUsage(totalModels, meta.ModelUsage) } - return totalCount, allFiles, totalTokens, nil + return totalCount, allFiles, totalTokens, totalModels, nil } func checkpointCreatedAt(opts WriteOptions) time.Time { @@ -990,12 +994,14 @@ func aggregateTokenUsage(a, b *agent.TokenUsage) *agent.TokenUsage { return nil } result := &agent.TokenUsage{} + var aCost, bCost *float64 if a != nil { result.InputTokens = a.InputTokens result.CacheCreationTokens = a.CacheCreationTokens result.CacheReadTokens = a.CacheReadTokens result.OutputTokens = a.OutputTokens result.APICallCount = a.APICallCount + aCost = a.CostUSD } if b != nil { result.InputTokens += b.InputTokens @@ -1003,10 +1009,47 @@ func aggregateTokenUsage(a, b *agent.TokenUsage) *agent.TokenUsage { result.CacheReadTokens += b.CacheReadTokens result.OutputTokens += b.OutputTokens result.APICallCount += b.APICallCount + bCost = b.CostUSD } + result.CostUSD = types.AddCostUSD(aCost, bCost) + // MergeCostSourceUsages folds a priced side with an unpriced-but-token-bearing + // side to mixed, so a summary that combines a costed session with an + // uncosted token-bearing one reports partial coverage honestly. + result.CostSource = types.MergeCostSourceUsages(a, b) return result } +// mergeModelUsage merges two per-model usage lists keyed by model, summing token +// counts and folding costs with the same AddCostUSD/MergeCostSource rules as +// aggregateTokenUsage. The result is sorted by model for deterministic output. +// Returns nil when both inputs are empty so the omitempty JSON tag drops the field. +func mergeModelUsage(a, b []types.ModelUsage) []types.ModelUsage { + if len(a) == 0 && len(b) == 0 { + return nil + } + byModel := make(map[string]*agent.TokenUsage, len(a)+len(b)) + fold := func(list []types.ModelUsage) { + for i := range list { + usage := list[i].TokenUsage + byModel[list[i].Model] = aggregateTokenUsage(byModel[list[i].Model], &usage) + } + } + fold(a) + fold(b) + + models := make([]string, 0, len(byModel)) + for model := range byModel { + models = append(models, model) + } + sort.Strings(models) + + out := make([]types.ModelUsage, 0, len(models)) + for _, model := range models { + out = append(out, types.ModelUsage{Model: model, TokenUsage: *byModel[model]}) + } + return out +} + // writeTranscript writes the transcript, compact transcript, and content hash // to the checkpoint entries. The compact transcript.jsonl (the full compacted // session) is written into the tree and pushed alongside full.jsonl. Returns @@ -1796,7 +1839,7 @@ func (s *treeWriter) updateSessionMetadata(sessionDir string, entries map[string } mutate(metadata) - metadataJSON, err := jsonutil.MarshalIndentWithNewline(metadata, "", " ") + metadataJSON, err := jsonutil.MarshalIndentWithNewline(metadata.WithoutCost(), "", " ") if err != nil { return fmt.Errorf("marshal session metadata: %w", err) } diff --git a/cmd/entire/cli/checkpoint_tokens.go b/cmd/entire/cli/checkpoint_tokens.go index ac07ef4cdb..f05e7c8c8a 100644 --- a/cmd/entire/cli/checkpoint_tokens.go +++ b/cmd/entire/cli/checkpoint_tokens.go @@ -9,8 +9,10 @@ import ( "strings" "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" "github.com/entireio/cli/cmd/entire/cli/checkpoint" "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" + "github.com/entireio/cli/cmd/entire/cli/pricing" "github.com/spf13/cobra" ) @@ -42,7 +44,9 @@ type checkpointTokensComparison struct { CacheWrite *checkpointTokensMetricDelta `json:"cache_write,omitempty"` Output *checkpointTokensMetricDelta `json:"output,omitempty"` APICalls *checkpointTokensMetricDelta `json:"api_calls,omitempty"` + Cost *checkpointTokensCostDelta `json:"cost,omitempty"` CacheReadCaveat string `json:"cache_read_caveat,omitempty"` + CostCaveat string `json:"cost_caveat,omitempty"` Qualification string `json:"qualification"` Limitations []string `json:"limitations,omitempty"` } @@ -55,6 +59,20 @@ type checkpointTokensMetricDelta struct { Direction string `json:"direction"` } +// checkpointTokensCostDelta mirrors checkpointTokensMetricDelta for USD cost. It +// is only populated when both compared checkpoints carry a non-nil cost. +type checkpointTokensCostDelta struct { + Baseline float64 `json:"baseline"` + Current float64 `json:"current"` + Change float64 `json:"change"` + ChangePercent *float64 `json:"change_percent,omitempty"` + Direction string `json:"direction"` +} + +// costNotComparableCaveat explains why a compared cost delta was omitted: cost +// is present on one side but not the other, so no honest delta can be shown. +const costNotComparableCaveat = "cost not comparable: missing on one side" + const ( checkpointComparisonStatusUnavailable = "unavailable" checkpointComparisonStatusObservedReduction = "observed_reduction" @@ -155,7 +173,7 @@ func loadCheckpointTokensReport(ctx context.Context, cmd *cobra.Command, checkpo return checkpointTokensReport{}, lookup, err } - return buildCheckpointTokensReport(cpID, summary, metas, metadataWarnings), lookup, nil + return buildCheckpointTokensReport(cpID, summary, metas, metadataWarnings, loadDisplayPricingTable(ctx)), lookup, nil } func readCheckpointTokenSessionMetadata(ctx context.Context, store reviewContextSessionMetadataReader, cpID id.CheckpointID, sessionCount int) ([]*checkpoint.Metadata, int, error) { @@ -181,7 +199,7 @@ func readCheckpointTokenSessionMetadata(ctx context.Context, store reviewContext return metas, warnings, nil } -func buildCheckpointTokensReport(cpID id.CheckpointID, summary *checkpoint.CheckpointSummary, metas []*checkpoint.Metadata, metadataWarnings int) checkpointTokensReport { +func buildCheckpointTokensReport(cpID id.CheckpointID, summary *checkpoint.CheckpointSummary, metas []*checkpoint.Metadata, metadataWarnings int, table *pricing.Table) checkpointTokensReport { report := checkpointTokensReport{ CheckpointID: cpID.String(), Source: "committed_checkpoint", @@ -215,6 +233,7 @@ func buildCheckpointTokensReport(cpID id.CheckpointID, summary *checkpoint.Check usage := checkpointTokenUsage(summary, metas, metadataWarnings > 0) if tokens := buildSessionTokensUsage(usage); tokens != nil { report.Tokens = tokens + applyLocalCostEstimate(tokens, usage, checkpointModelUsage(summary, metas, metadataWarnings > 0), report.Model, table) if tokens.SubagentTotal > 0 { report.Contributors = append(report.Contributors, sessionTokensContributor{ Kind: "subagents", @@ -356,17 +375,40 @@ func checkpointTokenUsage(summary *checkpoint.CheckpointSummary, metas []*checkp return sessionUsage } +// checkpointModelUsage returns the per-model token buckets for a checkpoint, +// mirroring checkpointTokenUsage's source preference so the local cost estimate +// prices the same tokens the report displays: the per-session buckets when +// metadata read cleanly, else the aggregated root-summary buckets. Buckets carry +// only token counts (cost is a display-only local estimate applied downstream). +func checkpointModelUsage(summary *checkpoint.CheckpointSummary, metas []*checkpoint.Metadata, metadataReadWarning bool) []types.ModelUsage { + var sessionBuckets []types.ModelUsage + for _, meta := range metas { + if meta != nil { + sessionBuckets = append(sessionBuckets, meta.ModelUsage...) + } + } + if !metadataReadWarning && len(sessionBuckets) > 0 { + return sessionBuckets + } + if summary != nil && len(summary.ModelUsage) > 0 { + return summary.ModelUsage + } + return sessionBuckets +} + func addCheckpointTokenUsage(a, b *agent.TokenUsage) *agent.TokenUsage { if a == nil && b == nil { return nil } result := &agent.TokenUsage{} + var aCost, bCost *float64 if a != nil { result.InputTokens = a.InputTokens result.CacheCreationTokens = a.CacheCreationTokens result.CacheReadTokens = a.CacheReadTokens result.OutputTokens = a.OutputTokens result.APICallCount = a.APICallCount + aCost = a.CostUSD } if b != nil { result.InputTokens = saturatingIntAdd(result.InputTokens, b.InputTokens) @@ -374,8 +416,13 @@ func addCheckpointTokenUsage(a, b *agent.TokenUsage) *agent.TokenUsage { result.CacheReadTokens = saturatingIntAdd(result.CacheReadTokens, b.CacheReadTokens) result.OutputTokens = saturatingIntAdd(result.OutputTokens, b.OutputTokens) result.APICallCount = saturatingIntAdd(result.APICallCount, b.APICallCount) + bCost = b.CostUSD } result.SubagentTokens = addCheckpointTokenUsage(tokenUsageSubagents(a), tokenUsageSubagents(b)) + result.CostUSD = types.AddCostUSD(aCost, bCost) + // MergeCostSourceUsages folds a priced side with an unpriced-but-token-bearing + // side to mixed, mirroring the per-bucket coverage rule. + result.CostSource = types.MergeCostSourceUsages(a, b) return result } @@ -424,6 +471,7 @@ func buildCheckpointTokensComparison(target, baseline checkpointTokensReport) *c comparison.Output = buildCheckpointMetricDelta(baseline.Tokens.Output, target.Tokens.Output) comparison.APICalls = buildCheckpointMetricDelta(baseline.Tokens.APICalls, target.Tokens.APICalls) comparison.CacheReadCaveat = checkpointComparisonCacheReadCaveat(comparison.CacheRead) + comparison.Cost, comparison.CostCaveat = buildCheckpointCostComparison(baseline.Tokens.CostUSD, target.Tokens.CostUSD) comparison.Status = checkpointComparisonStatus(comparison.Total) comparison.Qualification = checkpointComparisonQualification(comparison.Status) if classes := checkpointCostProxyPressureIncreased(comparison); len(classes) > 0 { @@ -474,6 +522,52 @@ func buildCheckpointMetricDelta(baseline, current int) *checkpointTokensMetricDe return delta } +// buildCheckpointCostComparison returns the cost delta and an accompanying +// caveat. A delta is produced only when BOTH sides carry a non-nil cost; +// otherwise it returns (nil, caveat) when exactly one side has cost (so the +// asymmetry is explained) and (nil, "") when neither side has cost (cost simply +// is not tracked for these checkpoints, so nothing is said). +func buildCheckpointCostComparison(baseline, current *float64) (*checkpointTokensCostDelta, string) { + switch { + case baseline != nil && current != nil: + return buildCheckpointCostDelta(*baseline, *current), "" + case baseline != nil || current != nil: + return nil, costNotComparableCaveat + default: + return nil, "" + } +} + +func buildCheckpointCostDelta(baseline, current float64) *checkpointTokensCostDelta { + change := current - baseline + delta := &checkpointTokensCostDelta{ + Baseline: baseline, + Current: current, + Change: change, + Direction: checkpointDeltaDirectionFloat(change), + } + if baseline != 0 { + percent := (change / baseline) * 100 + delta.ChangePercent = &percent + } + return delta +} + +// checkpointDeltaDirectionFloat classifies a USD cost change using the same +// down/up/unchanged semantics as checkpointDeltaDirection, but at cent +// granularity so sub-cent float noise reads as "unchanged". +func checkpointDeltaDirectionFloat(change float64) string { + cents := change * 100 + switch { + case cents <= -0.5: + return checkpointDeltaDirectionDown + case cents >= 0.5: + return checkpointDeltaDirectionUp + default: + return checkpointDeltaDirectionUnchanged + } +} + func saturatingIntSub(a, b int) int { if b < 0 { if b == minInt() { @@ -574,6 +668,7 @@ func writeCheckpointTokensText(w io.Writer, report checkpointTokensReport) { } writeTokenUsageSection(w, report.Tokens) + writeTokenCostLine(w, report.Tokens) writeCheckpointTokenComparison(w, report.Comparison) if len(report.Recommendations) > 0 { writeTokenRecommendations(w, report.Recommendations) @@ -632,6 +727,9 @@ func writeCheckpointTokenComparison(w io.Writer, comparison *checkpointTokensCom if comparison.CacheReadCaveat != "" { fmt.Fprintf(w, "Caveat: %s\n", comparison.CacheReadCaveat) } + if comparison.CostCaveat != "" { + fmt.Fprintf(w, "Caveat: %s\n", comparison.CostCaveat) + } if comparison.Status != checkpointComparisonStatusUnavailable { fmt.Fprintf(w, "Total tokens: %s\n", formatCheckpointMetricDelta(comparison.Total, formatTokenCount)) fmt.Fprintf(w, "Input: %s\n", formatCheckpointMetricDelta(comparison.Input, formatTokenCount)) @@ -639,6 +737,9 @@ func writeCheckpointTokenComparison(w io.Writer, comparison *checkpointTokensCom fmt.Fprintf(w, "Cache write: %s\n", formatCheckpointMetricDelta(comparison.CacheWrite, formatTokenCount)) fmt.Fprintf(w, "Output: %s\n", formatCheckpointMetricDelta(comparison.Output, formatTokenCount)) fmt.Fprintf(w, "API calls: %s\n", formatCheckpointMetricDelta(comparison.APICalls, formatPlainCount)) + if comparison.Cost != nil { + fmt.Fprintf(w, "Cost: %s\n", formatCheckpointCostDelta(comparison.Cost)) + } } fmt.Fprintln(w) fmt.Fprintln(w, "Qualification") @@ -660,6 +761,21 @@ func formatCheckpointMetricDelta(delta *checkpointTokensMetricDelta, formatValue return fmt.Sprintf("%s %s (%s -> %s)", delta.Direction, formatPercent(absFloat(*delta.ChangePercent)), from, to) } +func formatCheckpointCostDelta(delta *checkpointTokensCostDelta) string { + if delta == nil { + return checkpointComparisonStatusUnavailable + } + from := formatCostAmount(delta.Baseline) + to := formatCostAmount(delta.Current) + if delta.Direction == checkpointDeltaDirectionUnchanged { + return fmt.Sprintf("unchanged (%s -> %s)", from, to) + } + if delta.ChangePercent == nil { + return fmt.Sprintf("%s (%s -> %s)", delta.Direction, from, to) + } + return fmt.Sprintf("%s %s (%s -> %s)", delta.Direction, formatPercent(absFloat(*delta.ChangePercent)), from, to) +} + func formatPlainCount(value int) string { return strconv.Itoa(value) } diff --git a/cmd/entire/cli/checkpoint_tokens_cost_test.go b/cmd/entire/cli/checkpoint_tokens_cost_test.go new file mode 100644 index 0000000000..5e67bff07d --- /dev/null +++ b/cmd/entire/cli/checkpoint_tokens_cost_test.go @@ -0,0 +1,245 @@ +package cli + +import ( + "bytes" + "strings" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/checkpoint" + "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" + "github.com/entireio/cli/cmd/entire/cli/pricing" +) + +func costPtr(v float64) *float64 { return &v } + +// testDisplayPricingTable is a deterministic table for the display-side local +// cost estimate: "test-model" prices every token class at $1/MTok (provider +// "test" has no explicit cache rates, so cache tokens fall back to the input +// rate). A usage of input+output = N tokens therefore estimates to N/1e6 USD. +func testDisplayPricingTable(t *testing.T) *pricing.Table { + t.Helper() + table, err := pricing.LoadTable([]pricing.ModelRate{ + {ID: "test-model", Provider: "test", InputPerMTok: 1, OutputPerMTok: 1}, + }) + if err != nil { + t.Fatalf("LoadTable: %v", err) + } + return table +} + +// addCheckpointTokenUsage must sum CostUSD, merge CostSource, and carry cost +// through the SubagentTokens recursion. +func TestAddCheckpointTokenUsage_CarriesCost(t *testing.T) { + t.Parallel() + a := &agent.TokenUsage{ + InputTokens: 10, + CostUSD: costPtr(0.50), + CostSource: types.CostSourceReported, + SubagentTokens: &agent.TokenUsage{ + InputTokens: 1, + CostUSD: costPtr(0.125), + CostSource: types.CostSourceReported, + }, + } + b := &agent.TokenUsage{ + InputTokens: 20, + CostUSD: costPtr(0.25), + CostSource: types.CostSourceEstimated, + SubagentTokens: &agent.TokenUsage{ + InputTokens: 2, + CostUSD: costPtr(0.375), + CostSource: types.CostSourceReported, + }, + } + got := addCheckpointTokenUsage(a, b) + if got.CostUSD == nil || *got.CostUSD != 0.75 { + t.Fatalf("CostUSD = %v, want 0.75", got.CostUSD) + } + if got.CostSource != types.CostSourceMixed { + t.Fatalf("CostSource = %q, want mixed", got.CostSource) + } + if got.SubagentTokens == nil { + t.Fatal("expected SubagentTokens carried") + } + if got.SubagentTokens.CostUSD == nil || *got.SubagentTokens.CostUSD != 0.5 { + t.Fatalf("subagent CostUSD = %v, want 0.5", *got.SubagentTokens.CostUSD) + } + if got.SubagentTokens.CostSource != types.CostSourceReported { + t.Fatalf("subagent CostSource = %q, want reported", got.SubagentTokens.CostSource) + } +} + +func TestAddCheckpointTokenUsage_NilCostStaysNil(t *testing.T) { + t.Parallel() + a := &agent.TokenUsage{InputTokens: 1} + b := &agent.TokenUsage{InputTokens: 2} + got := addCheckpointTokenUsage(a, b) + if got.CostUSD != nil { + t.Fatalf("CostUSD = %v, want nil", got.CostUSD) + } + if got.CostSource != "" { + t.Fatalf("CostSource = %q, want empty", got.CostSource) + } +} + +// buildCheckpointTokensReport must recompute a LOCAL cost estimate from the +// persisted token breakdown (the CLI no longer persists cost) and label it as a +// local estimate. Here the per-model breakdown carries 420000 priceable tokens +// under test-model ($1/MTok) => $0.42. +func TestBuildCheckpointTokensReport_RecomputesLocalCostEstimate(t *testing.T) { + t.Parallel() + cpID := id.MustCheckpointID("abc123abc123") + report := buildCheckpointTokensReport( + cpID, + &checkpoint.CheckpointSummary{ + CheckpointID: cpID, + Sessions: []checkpoint.SessionFilePaths{{Metadata: "0/metadata.json"}}, + }, + []*checkpoint.Metadata{ + { + SessionID: "cost-session", + Agent: "Claude Code", + Model: "test-model", + // Persisted usage carries NO cost (stripped on write); only tokens. + TokenUsage: &agent.TokenUsage{InputTokens: 400000, OutputTokens: 20000}, + ModelUsage: []types.ModelUsage{ + {Model: "test-model", TokenUsage: agent.TokenUsage{InputTokens: 400000, OutputTokens: 20000}}, + }, + }, + }, + 0, + testDisplayPricingTable(t), + ) + if report.Tokens == nil || report.Tokens.CostUSD == nil || *report.Tokens.CostUSD != 0.42 { + t.Fatalf("expected recomputed cost 0.42, got %+v", report.Tokens) + } + if report.Tokens.CostSource != types.CostSourceEstimated { + t.Fatalf("cost source = %q, want estimated", report.Tokens.CostSource) + } + + var buf bytes.Buffer + writeCheckpointTokensText(&buf, report) + out := buf.String() + if !strings.Contains(out, "Cost: $0.42 (estimated locally)") { + t.Fatalf("expected local-estimate cost line, got:\n%s", out) + } + if !strings.Contains(out, localCostEstimateNote) { + t.Fatalf("expected local-estimate note, got:\n%s", out) + } +} + +// With no pricing table (estimation disabled) and no priceable model, no cost is +// shown — never $0. +func TestBuildCheckpointTokensReport_NoCostWhenUnpriceable(t *testing.T) { + t.Parallel() + cpID := id.MustCheckpointID("abc123abc124") + report := buildCheckpointTokensReport( + cpID, + &checkpoint.CheckpointSummary{ + CheckpointID: cpID, + Sessions: []checkpoint.SessionFilePaths{{Metadata: "0/metadata.json"}}, + }, + []*checkpoint.Metadata{ + { + SessionID: "cost-session", + Agent: "Claude Code", + Model: "test-model", + TokenUsage: &agent.TokenUsage{InputTokens: 400000, OutputTokens: 20000}, + }, + }, + 0, + nil, // estimation disabled + ) + if report.Tokens == nil { + t.Fatal("expected token usage") + } + if report.Tokens.CostUSD != nil { + t.Fatalf("expected no cost when unpriceable, got %v", report.Tokens.CostUSD) + } + var buf bytes.Buffer + writeCheckpointTokensText(&buf, report) + if strings.Contains(buf.String(), "Cost:") { + t.Fatalf("expected no cost line, got:\n%s", buf.String()) + } +} + +// --compare emits a cost delta only when both checkpoints carry a cost. +func TestBuildCheckpointTokensComparison_CostDeltaBothSides(t *testing.T) { + t.Parallel() + baseline := checkpointTokensReport{ + CheckpointID: "aaa111bbb222", + Tokens: &sessionTokensUsage{Total: 1000, Input: 1000, CostUSD: costPtr(0.50), CostSource: types.CostSourceEstimated}, + } + target := checkpointTokensReport{ + CheckpointID: "bbb222ccc333", + Tokens: &sessionTokensUsage{Total: 800, Input: 800, CostUSD: costPtr(0.30), CostSource: types.CostSourceEstimated}, + } + + cmp := buildCheckpointTokensComparison(target, baseline) + if cmp.Cost == nil { + t.Fatal("expected cost delta") + } + if cmp.CostCaveat != "" { + t.Fatalf("expected no cost caveat, got %q", cmp.CostCaveat) + } + if cmp.Cost.Direction != checkpointDeltaDirectionDown { + t.Fatalf("direction = %q, want down", cmp.Cost.Direction) + } + if cmp.Cost.Baseline != 0.50 || cmp.Cost.Current != 0.30 { + t.Fatalf("baseline/current = %v/%v, want 0.50/0.30", cmp.Cost.Baseline, cmp.Cost.Current) + } + if cmp.Cost.ChangePercent == nil || *cmp.Cost.ChangePercent > -39.9 || *cmp.Cost.ChangePercent < -40.1 { + t.Fatalf("change percent = %v, want ~-40", cmp.Cost.ChangePercent) + } + + var buf bytes.Buffer + writeCheckpointTokenComparison(&buf, cmp) + if !strings.Contains(buf.String(), "Cost: down 40% ($0.50 -> $0.30)") { + t.Fatalf("expected cost delta line, got:\n%s", buf.String()) + } +} + +// --compare omits the delta and notes the asymmetry when exactly one side has cost. +func TestBuildCheckpointTokensComparison_CostCaveatWhenOneSideMissing(t *testing.T) { + t.Parallel() + baseline := checkpointTokensReport{ + CheckpointID: "aaa111bbb222", + Tokens: &sessionTokensUsage{Total: 1000, Input: 1000, CostUSD: costPtr(0.50), CostSource: types.CostSourceEstimated}, + } + target := checkpointTokensReport{ + CheckpointID: "bbb222ccc333", + Tokens: &sessionTokensUsage{Total: 800, Input: 800}, + } + + cmp := buildCheckpointTokensComparison(target, baseline) + if cmp.Cost != nil { + t.Fatalf("expected no cost delta, got %+v", cmp.Cost) + } + if cmp.CostCaveat != costNotComparableCaveat { + t.Fatalf("cost caveat = %q, want %q", cmp.CostCaveat, costNotComparableCaveat) + } + + var buf bytes.Buffer + writeCheckpointTokenComparison(&buf, cmp) + out := buf.String() + if !strings.Contains(out, "Caveat: cost not comparable: missing on one side") { + t.Fatalf("expected cost caveat line, got:\n%s", out) + } + if strings.Contains(out, "Cost: ") { + t.Fatalf("expected no cost delta line, got:\n%s", out) + } +} + +// When neither side has cost, cost is simply not tracked: no delta, no caveat noise. +func TestBuildCheckpointTokensComparison_NoCostNoiseWhenBothMissing(t *testing.T) { + t.Parallel() + baseline := checkpointTokensReport{CheckpointID: "aaa111bbb222", Tokens: &sessionTokensUsage{Total: 1000, Input: 1000}} + target := checkpointTokensReport{CheckpointID: "bbb222ccc333", Tokens: &sessionTokensUsage{Total: 800, Input: 800}} + + cmp := buildCheckpointTokensComparison(target, baseline) + if cmp.Cost != nil || cmp.CostCaveat != "" { + t.Fatalf("expected no cost delta and no caveat, got cost=%+v caveat=%q", cmp.Cost, cmp.CostCaveat) + } +} diff --git a/cmd/entire/cli/config.go b/cmd/entire/cli/config.go index 9db5b68519..c247d8aac2 100644 --- a/cmd/entire/cli/config.go +++ b/cmd/entire/cli/config.go @@ -7,6 +7,7 @@ import ( "github.com/entireio/cli/cmd/entire/cli/agent" "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/pricing" "github.com/entireio/cli/cmd/entire/cli/settings" "github.com/entireio/cli/cmd/entire/cli/strategy" @@ -53,6 +54,14 @@ func SaveEntireSettingsLocal(ctx context.Context, s *settings.EntireSettings) er return nil } +// LoadPricingTable builds the model pricing table (embedded defaults plus any +// configured overrides) and reports whether cost estimation is disabled. It is +// the cli-package entry point over settings.LoadPricingTable and never fails: on +// error the table is nil and callers omit cost rather than failing the hook. +func LoadPricingTable(ctx context.Context) (*pricing.Table, bool) { + return settings.LoadPricingTable(ctx) +} + // IsEnabled returns whether Entire is currently enabled. // Returns true by default if settings cannot be loaded. func IsEnabled(ctx context.Context) (bool, error) { diff --git a/cmd/entire/cli/labs.go b/cmd/entire/cli/labs.go index 606d382d36..5126f66bdc 100644 --- a/cmd/entire/cli/labs.go +++ b/cmd/entire/cli/labs.go @@ -40,6 +40,11 @@ var experimentalCommands = []experimentalCommandInfo{ Invocation: "entire tokens profile", Summary: "Aggregate token usage across committed checkpoints", }, + { + CommandPath: []string{"tokens", "pricing-refresh"}, + Invocation: "entire tokens pricing-refresh", + Summary: "Refresh the remote pricing cache (opt in via pricing.remote)", + }, { CommandPath: []string{"session", "tokens"}, Invocation: "entire session tokens", diff --git a/cmd/entire/cli/lifecycle.go b/cmd/entire/cli/lifecycle.go index a1c2ddab82..96076c5e61 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -912,14 +912,41 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev transcriptLinesAtStart = preState.TranscriptOffset } - // Resolve token usage. Hook-provided counts (e.g., Cursor's stop hook, - // which is the only authoritative source for Cursor sessions because the - // JSONL transcript has no usage fields) take precedence; otherwise fall - // back to transcript-based computation, preferring SubagentAwareExtractor - // to include subagent tokens. - tokenUsage := event.TokenUsage - if tokenUsage == nil { - tokenUsage = agent.CalculateTokenUsage(ctx, ag, transcriptData, transcriptLinesAtStart, subagentsDir) + // Resolve token usage and attribute cost. Hook-provided counts (e.g., + // Cursor's stop hook, which is the only authoritative source for Cursor + // sessions because the JSONL transcript has no usage fields) take precedence + // and are priced directly as a single bucket under the session model; + // otherwise usage is computed from the transcript (preferring + // SubagentAwareExtractor to include subagent tokens) and priced per model. + // The per-model buckets are threaded through StepContext.ModelUsage so the + // strategy can accumulate them into checkpoint-scoped per-model metadata. + table, disableEstimation := LoadPricingTable(ctx) + // Opt-in daily remote pricing refresh: when enabled and the cache is stale, + // spawn a detached worker to refresh it. Never fetches inline on this + // turn-end hook path; ShouldRefresh gates on the 24h backoff so the common + // case is a cheap cache stat. + maybeSpawnPricingRefresh(ctx) + // Codex "priority" service tier bills at a premium; when opted in, price the + // turn under the model's "-priority" variant. Only Codex reads settings here, + // so every other agent's turn-end skips the extra load. + pricingModel := event.Model + if agentType == agent.AgentTypeCodex && event.Model != "" { + if s, sErr := LoadEntireSettings(ctx); sErr == nil { + pricingModel = s.PricingModelForAgent(agentType, event.Model) + } + } + var tokenUsage *agent.TokenUsage + var buckets []agent.ModelUsage + if event.TokenUsage != nil { + tokenUsage, buckets = agent.PriceUsage(event.TokenUsage, pricingModel, table, disableEstimation) + } else { + var costErr error + tokenUsage, buckets, costErr = agent.CalculateUsageWithCost(ag, transcriptData, transcriptLinesAtStart, subagentsDir, table, pricingModel, disableEstimation) + if costErr != nil { + logging.Debug(logCtx, "failed usage-with-cost extraction", + slog.String("error", costErr.Error())) + tokenUsage, buckets = nil, nil + } } // Build fully-populated step context and delegate to strategy @@ -938,6 +965,7 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev StepTranscriptIdentifier: transcriptIdentifierAtStart, StepTranscriptStart: transcriptLinesAtStart, TokenUsage: tokenUsage, + ModelUsage: buckets, } if err := strat.SaveStep(ctx, stepCtx); err != nil { diff --git a/cmd/entire/cli/lifecycle_pricing_test.go b/cmd/entire/cli/lifecycle_pricing_test.go new file mode 100644 index 0000000000..aece96063a --- /dev/null +++ b/cmd/entire/cli/lifecycle_pricing_test.go @@ -0,0 +1,50 @@ +package cli + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/pricing" + "github.com/entireio/cli/cmd/entire/cli/settings" +) + +// tierSettings builds an EntireSettings carrying only a Codex service tier, the +// single input the pricing seam reads. +func tierSettings(tier string) *settings.EntireSettings { + return &settings.EntireSettings{Pricing: &settings.PricingSettings{CodexServiceTier: tier}} +} + +// TestCodexPriorityTierPricesAtPremium ties the service-tier knob to the actual +// priced outcome through the real embedded table and the lifecycle-facing seam +// (settings.EntireSettings.PricingModelForAgent): with "priority" a Codex turn on +// gpt-5.5 prices under gpt-5.5-priority (12.5/75), without it under gpt-5.5 +// (5/30). This is the turn-end analogue of the condensation regression. +func TestCodexPriorityTierPricesAtPremium(t *testing.T) { + t.Parallel() + + table, err := pricing.LoadTable(nil) + if err != nil { + t.Fatalf("LoadTable: %v", err) + } + usage := &agent.TokenUsage{InputTokens: 1_000_000, OutputTokens: 1_000_000} + + // Priority knob -> gpt-5.5-priority: 1M@$12.5 + 1M@$75 = 87.5. + prioModel := tierSettings("priority").PricingModelForAgent(agent.AgentTypeCodex, "gpt-5.5") + if prioModel != "gpt-5.5-priority" { + t.Fatalf("priority model = %q, want gpt-5.5-priority", prioModel) + } + prioPriced, _ := agent.PriceUsage(usage, prioModel, table, false) + if prioPriced.CostUSD == nil || *prioPriced.CostUSD != 87.5 { + t.Fatalf("priority cost = %v, want 87.5", prioPriced.CostUSD) + } + + // No knob -> gpt-5.5: 1M@$5 + 1M@$30 = 35. + stdModel := tierSettings("").PricingModelForAgent(agent.AgentTypeCodex, "gpt-5.5") + if stdModel != "gpt-5.5" { + t.Fatalf("standard model = %q, want gpt-5.5", stdModel) + } + stdPriced, _ := agent.PriceUsage(usage, stdModel, table, false) + if stdPriced.CostUSD == nil || *stdPriced.CostUSD != 35.0 { + t.Fatalf("standard cost = %v, want 35.0", stdPriced.CostUSD) + } +} diff --git a/cmd/entire/cli/pricing/cursor_composer25_test.go b/cmd/entire/cli/pricing/cursor_composer25_test.go new file mode 100644 index 0000000000..5c748c93df --- /dev/null +++ b/cmd/entire/cli/pricing/cursor_composer25_test.go @@ -0,0 +1,81 @@ +package pricing + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestCursorComposer25 covers the composer-2.5 standard entry added because real +// Cursor CLI sessions report the bare id "composer-2.5" (the current default), +// which previously resolved to no rate and left cost nil. +func TestCursorComposer25(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + // Exact id resolves to the standard rate (same $0.50/$2.50 as composer-2). + rate, ok := tbl.Lookup("composer-2.5") + require.True(t, ok, "composer-2.5 must resolve") + assert.Equal(t, "composer-2.5", rate.ID) + assert.Equal(t, "cursor", rate.Provider) + assert.InDelta(t, 0.5, rate.InputPerMTok, 1e-9) + assert.InDelta(t, 2.5, rate.OutputPerMTok, 1e-9) + require.NotNil(t, rate.CacheReadPerMTok) + assert.InDelta(t, 0.05, *rate.CacheReadPerMTok, 1e-9) + require.NotNil(t, rate.CacheWritePerMTok) + assert.InDelta(t, 0.5, *rate.CacheWritePerMTok, 1e-9) + + // Estimate matches hand math: + // 100000*0.5 + 40000*2.5 + 10000*0.05 + 5000*0.5 = 153000 -> $0.153. + usage := types.TokenUsage{ + InputTokens: 100_000, + OutputTokens: 40_000, + CacheReadTokens: 10_000, + CacheCreationTokens: 5_000, + } + assert.InDelta(t, 0.153, Estimate(rate, usage), 1e-9) + + // Alias forms resolve to composer-2.5. + for _, q := range []string{"cursor/composer-2.5", "composer-2.5-2026-05-18"} { + got, ok := tbl.Lookup(q) + require.Truef(t, ok, "%q must resolve", q) + assert.Equalf(t, "composer-2.5", got.ID, "%q should resolve to composer-2.5", q) + } +} + +// TestCursorComposer25NoCrossMatch guards the exact-match-wins boundary between +// composer-2.5 and the pre-existing composer-2.5-fast and composer-2 entries: a +// glob from one family must never swallow the other's ids. +func TestCursorComposer25NoCrossMatch(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + cases := []struct { + query string + wantID string + }{ + {"composer-2.5", "composer-2.5"}, // exact standard + {"composer-2.5-fast", "composer-2.5-fast"}, // exact fast wins over composer-2.5 glob + {"composer-2.5-2026-05-18", "composer-2.5"}, // dated standard -> standard + {"composer-2.5-fast-2026-05-18", "composer-2.5-fast"}, // dated fast -> fast + {"composer-2", "composer-2"}, // unrelated base id unaffected + } + for _, tc := range cases { + got, ok := tbl.Lookup(tc.query) + require.Truef(t, ok, "%q must resolve", tc.query) + assert.Equalf(t, tc.wantID, got.ID, "Lookup(%q)", tc.query) + } + + // composer-2.5 and composer-2.5-fast are distinct entries with distinct rates: + // standard $2.50 output vs fast $15.00 output. + std, _ := tbl.Lookup("composer-2.5") + fast, _ := tbl.Lookup("composer-2.5-fast") + assert.InDelta(t, 2.5, std.OutputPerMTok, 1e-9) + assert.InDelta(t, 15.0, fast.OutputPerMTok, 1e-9) +} diff --git a/cmd/entire/cli/pricing/doc.go b/cmd/entire/cli/pricing/doc.go new file mode 100644 index 0000000000..ee100437db --- /dev/null +++ b/cmd/entire/cli/pricing/doc.go @@ -0,0 +1,53 @@ +// Package pricing provides a small, embedded table of per-model token prices +// and a helper to estimate the USD cost of a checkpoint's token usage. +// +// The rate table is maintained as JSON files under models/ and compiled into +// the binary via go:embed. To change or add a price, edit the appropriate JSON +// file (grouped by provider) and open a pull request. Operators can layer +// per-repository overrides through settings, which LoadTable applies on top of +// the embedded defaults (replacing an entry whose id matches, or appending a +// new one). +// +// An opt-in remote layer (remote.go, gated by the pricing.remote setting) can +// additionally refresh a cached rate table from a remote source roughly once a +// day (RefreshRemoteCache), layered on top of the embedded defaults the same +// way settings overrides are. There is no inline fetch on the hook/condensation +// path: the refresh is spawned as a detached background worker +// (maybeSpawnPricingRefresh) and never blocks a foreground command — LoadTable +// only ever reads the on-disk cache the worker last wrote. +// +// Lookup performs exact-id and alias matching only, with no implicit +// model-family fallback: an unknown model resolves to no rate, and callers must +// treat that as "no estimate available" rather than guessing a price. Aliases +// cover the id spellings seen across providers and tools: bare ids, Anthropic +// dated ids (claude-haiku-4-5-20251001), slash-prefixed ids +// (anthropic/claude-sonnet-5), Bedrock ids and regional inference profiles +// (anthropic.claude-opus-4-8-v1:0, us.anthropic.claude-opus-4-8-v1:0), and +// Vertex ids (claude-opus-4-5@20251101). +// +// A few billing nuances are deliberately approximated: +// +// - Long context: a 1M-context request on a current-generation model bills at +// the model's standard rate — there is no long-context premium — so the +// "[1m]" ids Claude Code emits (e.g. claude-fable-5[1m]) share the base rate. +// Because path.Match would read a literal "[1m]" alias as a character class, +// Lookup strips a trailing "[...]" suffix from the query rather than relying +// on such an alias. +// +// - Cache writes: the derived cache-write (1.25x input) and cache-read (0.1x +// input) multipliers are Anthropic-only economics and are applied only when a +// rate's provider is "anthropic" and it omits an explicit cache rate. For any +// other provider a missing cache rate falls back to the full input rate (1.0x), +// billing cached tokens as normal input so an estimate never silently +// undercharges; the non-Anthropic embedded tables set explicit cache rates, +// so that fallback is only a safety net. The 1.25x write multiplier itself +// assumes the 5-minute-TTL premium; a 1-hour-TTL cache write actually bills +// 2x input, so Anthropic estimates undercount when 1h writes are present. This +// is a known under-estimate until per-TTL buckets are parsed; Claude Code +// transcripts report cache_creation.ephemeral_1h_input_tokens and +// ephemeral_5m_input_tokens separately. +// +// - Fast mode: turns billed at the fast-mode premium (usage.speed == "fast") +// are priced here at standard rates — another known under-estimate — because +// that premium is not published in-table. +package pricing diff --git a/cmd/entire/cli/pricing/embed.go b/cmd/entire/cli/pricing/embed.go new file mode 100644 index 0000000000..804b727e9a --- /dev/null +++ b/cmd/entire/cli/pricing/embed.go @@ -0,0 +1,9 @@ +package pricing + +import "embed" + +// modelsFS holds the embedded per-provider pricing tables (models/*.json), +// parsed by LoadTable. +// +//go:embed models/*.json +var modelsFS embed.FS diff --git a/cmd/entire/cli/pricing/google_gemini_test.go b/cmd/entire/cli/pricing/google_gemini_test.go new file mode 100644 index 0000000000..41785ee900 --- /dev/null +++ b/cmd/entire/cli/pricing/google_gemini_test.go @@ -0,0 +1,85 @@ +package pricing + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestGoogleGemini31And35 covers the gemini-3.1-pro and gemini-3.5-flash entries +// added because Cursor CLI and Gemini CLI report these ids today, while the table +// previously carried only the 3.0 generation. +func TestGoogleGemini31And35(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + proSample := types.TokenUsage{InputTokens: 50_000, OutputTokens: 20_000, CacheReadTokens: 8_000, CacheCreationTokens: 4_000} + + // gemini-3.1-pro: $2/$12, cache read $0.20. + pro, ok := tbl.Lookup("gemini-3.1-pro") + require.True(t, ok, "gemini-3.1-pro must resolve") + assert.Equal(t, "gemini-3.1-pro", pro.ID) + assert.InDelta(t, 2.0, pro.InputPerMTok, 1e-9) + assert.InDelta(t, 12.0, pro.OutputPerMTok, 1e-9) + require.NotNil(t, pro.CacheReadPerMTok) + assert.InDelta(t, 0.2, *pro.CacheReadPerMTok, 1e-9) + // 50000*2 + 20000*12 + 8000*0.2 + 4000*2 = 349600 -> $0.3496. + assert.InDelta(t, 0.3496, Estimate(pro, proSample), 1e-9) + + // gemini-3.5-flash: $1.50/$9, cache read $0.15. + flash, ok := tbl.Lookup("gemini-3.5-flash") + require.True(t, ok, "gemini-3.5-flash must resolve") + assert.Equal(t, "gemini-3.5-flash", flash.ID) + assert.InDelta(t, 1.5, flash.InputPerMTok, 1e-9) + assert.InDelta(t, 9.0, flash.OutputPerMTok, 1e-9) + require.NotNil(t, flash.CacheReadPerMTok) + assert.InDelta(t, 0.15, *flash.CacheReadPerMTok, 1e-9) + // 50000*1.5 + 20000*9 + 8000*0.15 + 4000*1.5 = 262200 -> $0.2622. + assert.InDelta(t, 0.2622, Estimate(flash, proSample), 1e-9) + + // Alias forms resolve. + for q, wantID := range map[string]string{ + "google/gemini-3.1-pro": "gemini-3.1-pro", + "gemini-3.1-pro-preview": "gemini-3.1-pro", + "google/gemini-3.5-flash": "gemini-3.5-flash", + "gemini-3.5-flash-preview": "gemini-3.5-flash", + } { + got, ok := tbl.Lookup(q) + require.Truef(t, ok, "%q must resolve", q) + assert.Equalf(t, wantID, got.ID, "Lookup(%q)", q) + } +} + +// TestGoogleGeminiDottedGlobsDoNotCrossMatch guards the boundary between the +// dotted 3.1/3.5 ids and the 3.0 globs: path.Match's "3.1"/"3.5" must not fall +// into the "gemini-3-pro-*" / "gemini-3-flash-*" entries, and the reverse. +func TestGoogleGeminiDottedGlobsDoNotCrossMatch(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + cases := []struct { + query string + wantID string + }{ + // 3.0 dated variants stay on the 3.0 entries (new dotted globs don't swallow them). + {"gemini-3-pro-preview", "gemini-3-pro"}, + {"gemini-3-flash-preview", "gemini-3-flash"}, + // dotted dated variants stay on the dotted entries (3.0 globs don't swallow them). + {"gemini-3.1-pro-preview", "gemini-3.1-pro"}, + {"gemini-3.5-flash-preview", "gemini-3.5-flash"}, + // exact base ids unaffected. + {"gemini-3-pro", "gemini-3-pro"}, + {"gemini-3-flash", "gemini-3-flash"}, + } + for _, tc := range cases { + got, ok := tbl.Lookup(tc.query) + require.Truef(t, ok, "%q must resolve", tc.query) + assert.Equalf(t, tc.wantID, got.ID, "Lookup(%q)", tc.query) + } +} diff --git a/cmd/entire/cli/pricing/models/anthropic.json b/cmd/entire/cli/pricing/models/anthropic.json new file mode 100644 index 0000000000..2636ff892d --- /dev/null +++ b/cmd/entire/cli/pricing/models/anthropic.json @@ -0,0 +1,387 @@ +{ + "schema_version": 1, + "models": [ + { + "id": "claude-fable-5", + "provider": "anthropic", + "aliases": [ + "claude-fable-5-2*", + "anthropic/claude-fable-5*", + "anthropic.claude-fable-5", + "anthropic.claude-fable-5-*", + "us.anthropic.claude-fable-5*", + "eu.anthropic.claude-fable-5*", + "apac.anthropic.claude-fable-5*", + "claude-fable-5@*", + "global.anthropic.claude-fable-5*", + "claude-fable-5-v*@*" + ], + "input_per_mtok": 10, + "output_per_mtok": 50, + "effective_date": "2026-06-24" + }, + { + "id": "claude-mythos-5", + "provider": "anthropic", + "aliases": [ + "claude-mythos-5-2*", + "anthropic/claude-mythos-5*", + "anthropic.claude-mythos-5", + "anthropic.claude-mythos-5-*", + "us.anthropic.claude-mythos-5*", + "eu.anthropic.claude-mythos-5*", + "apac.anthropic.claude-mythos-5*", + "claude-mythos-5@*", + "global.anthropic.claude-mythos-5*", + "claude-mythos-5-v*@*" + ], + "input_per_mtok": 10, + "output_per_mtok": 50, + "effective_date": "2026-06-24" + }, + { + "id": "claude-opus-4-8-fast", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-8-2*-fast", + "anthropic/claude-opus-4-8-fast", + "anthropic.claude-opus-4-8-fast", + "us.anthropic.claude-opus-4-8-fast*", + "eu.anthropic.claude-opus-4-8-fast*", + "apac.anthropic.claude-opus-4-8-fast*", + "claude-opus-4-8-fast@*", + "global.anthropic.claude-opus-4-8-fast*", + "anthropic/claude-opus-4.8-fast", + "claude-opus-4.8-fast" + ], + "input_per_mtok": 10, + "output_per_mtok": 50, + "effective_date": "2026-07-10" + }, + { + "id": "claude-opus-4-8", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-8-2*", + "anthropic/claude-opus-4-8*", + "anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-8-*", + "us.anthropic.claude-opus-4-8*", + "eu.anthropic.claude-opus-4-8*", + "apac.anthropic.claude-opus-4-8*", + "claude-opus-4-8@*", + "global.anthropic.claude-opus-4-8*", + "claude-opus-4-8-v*@*" + ], + "input_per_mtok": 5, + "output_per_mtok": 25, + "effective_date": "2026-06-24" + }, + { + "id": "claude-opus-4-7", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-7-2*", + "anthropic/claude-opus-4-7*", + "anthropic.claude-opus-4-7", + "anthropic.claude-opus-4-7-*", + "us.anthropic.claude-opus-4-7*", + "eu.anthropic.claude-opus-4-7*", + "apac.anthropic.claude-opus-4-7*", + "claude-opus-4-7@*", + "global.anthropic.claude-opus-4-7*", + "claude-opus-4-7-v*@*" + ], + "input_per_mtok": 5, + "output_per_mtok": 25, + "effective_date": "2026-03-01" + }, + { + "id": "claude-opus-4-6", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-6-2*", + "anthropic/claude-opus-4-6*", + "anthropic.claude-opus-4-6", + "anthropic.claude-opus-4-6-*", + "us.anthropic.claude-opus-4-6*", + "eu.anthropic.claude-opus-4-6*", + "apac.anthropic.claude-opus-4-6*", + "claude-opus-4-6@*", + "global.anthropic.claude-opus-4-6*", + "claude-opus-4-6-v*@*" + ], + "input_per_mtok": 5, + "output_per_mtok": 25, + "effective_date": "2026-01-01" + }, + { + "id": "claude-opus-4-5", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-5-2*", + "anthropic/claude-opus-4-5*", + "anthropic.claude-opus-4-5", + "anthropic.claude-opus-4-5-*", + "us.anthropic.claude-opus-4-5*", + "eu.anthropic.claude-opus-4-5*", + "apac.anthropic.claude-opus-4-5*", + "claude-opus-4-5@*", + "global.anthropic.claude-opus-4-5*", + "claude-opus-4-5-v*@*" + ], + "input_per_mtok": 5, + "output_per_mtok": 25, + "effective_date": "2025-11-01" + }, + { + "id": "claude-opus-4-1", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-1-2*", + "anthropic/claude-opus-4-1*", + "anthropic.claude-opus-4-1", + "anthropic.claude-opus-4-1-*", + "us.anthropic.claude-opus-4-1*", + "eu.anthropic.claude-opus-4-1*", + "apac.anthropic.claude-opus-4-1*", + "claude-opus-4-1@*", + "global.anthropic.claude-opus-4-1*", + "claude-opus-4-1-v*@*" + ], + "input_per_mtok": 15, + "output_per_mtok": 75, + "effective_date": "2025-08-05" + }, + { + "id": "claude-opus-4-0", + "provider": "anthropic", + "aliases": [ + "claude-opus-4-0-2*", + "anthropic/claude-opus-4-0*", + "anthropic.claude-opus-4-0", + "anthropic.claude-opus-4-0-*", + "us.anthropic.claude-opus-4-0*", + "eu.anthropic.claude-opus-4-0*", + "apac.anthropic.claude-opus-4-0*", + "claude-opus-4-0@*", + "global.anthropic.claude-opus-4-0*", + "claude-opus-4-0-v*@*" + ], + "input_per_mtok": 15, + "output_per_mtok": 75, + "effective_date": "2025-05-14" + }, + { + "id": "claude-opus-3", + "provider": "anthropic", + "aliases": [ + "claude-3-opus", + "claude-opus-3-2*", + "claude-3-opus-2*", + "anthropic/claude-opus-3*", + "anthropic/claude-3-opus", + "anthropic.claude-opus-3*", + "anthropic.claude-3-opus*", + "us.anthropic.claude-opus-3*", + "us.anthropic.claude-3-opus*", + "eu.anthropic.claude-opus-3*", + "eu.anthropic.claude-3-opus*", + "apac.anthropic.claude-opus-3*", + "apac.anthropic.claude-3-opus*", + "claude-opus-3@*", + "claude-3-opus@*", + "global.anthropic.claude-opus-3*", + "claude-opus-3-v*@*" + ], + "input_per_mtok": 15, + "output_per_mtok": 75, + "effective_date": "2024-02-29" + }, + { + "id": "claude-sonnet-5", + "provider": "anthropic", + "aliases": [ + "claude-sonnet-5-2*", + "anthropic/claude-sonnet-5*", + "anthropic.claude-sonnet-5", + "anthropic.claude-sonnet-5-*", + "us.anthropic.claude-sonnet-5*", + "eu.anthropic.claude-sonnet-5*", + "apac.anthropic.claude-sonnet-5*", + "claude-sonnet-5@*", + "global.anthropic.claude-sonnet-5*", + "claude-sonnet-5-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2026-06-24" + }, + { + "id": "claude-sonnet-4-6", + "provider": "anthropic", + "aliases": [ + "claude-sonnet-4-6-2*", + "anthropic/claude-sonnet-4-6*", + "anthropic.claude-sonnet-4-6", + "anthropic.claude-sonnet-4-6-*", + "us.anthropic.claude-sonnet-4-6*", + "eu.anthropic.claude-sonnet-4-6*", + "apac.anthropic.claude-sonnet-4-6*", + "claude-sonnet-4-6@*", + "global.anthropic.claude-sonnet-4-6*", + "claude-sonnet-4-6-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2026-01-01" + }, + { + "id": "claude-sonnet-4-5", + "provider": "anthropic", + "aliases": [ + "claude-sonnet-4-5-2*", + "anthropic/claude-sonnet-4-5*", + "anthropic.claude-sonnet-4-5", + "anthropic.claude-sonnet-4-5-*", + "us.anthropic.claude-sonnet-4-5*", + "eu.anthropic.claude-sonnet-4-5*", + "apac.anthropic.claude-sonnet-4-5*", + "claude-sonnet-4-5@*", + "global.anthropic.claude-sonnet-4-5*", + "claude-sonnet-4-5-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2025-09-29" + }, + { + "id": "claude-sonnet-4-0", + "provider": "anthropic", + "aliases": [ + "claude-sonnet-4-0-2*", + "anthropic/claude-sonnet-4-0*", + "anthropic.claude-sonnet-4-0", + "anthropic.claude-sonnet-4-0-*", + "us.anthropic.claude-sonnet-4-0*", + "eu.anthropic.claude-sonnet-4-0*", + "apac.anthropic.claude-sonnet-4-0*", + "claude-sonnet-4-0@*", + "global.anthropic.claude-sonnet-4-0*", + "claude-sonnet-4-0-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2025-05-14" + }, + { + "id": "claude-3-7-sonnet", + "provider": "anthropic", + "aliases": [ + "claude-3-7-sonnet-2*", + "anthropic/claude-3-7-sonnet*", + "anthropic.claude-3-7-sonnet*", + "us.anthropic.claude-3-7-sonnet*", + "eu.anthropic.claude-3-7-sonnet*", + "apac.anthropic.claude-3-7-sonnet*", + "claude-3-7-sonnet@*", + "global.anthropic.claude-3-7-sonnet*", + "claude-3-7-sonnet-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2025-02-19" + }, + { + "id": "claude-3-5-sonnet", + "provider": "anthropic", + "aliases": [ + "claude-3-5-sonnet-2*", + "anthropic/claude-3-5-sonnet*", + "anthropic.claude-3-5-sonnet*", + "us.anthropic.claude-3-5-sonnet*", + "eu.anthropic.claude-3-5-sonnet*", + "apac.anthropic.claude-3-5-sonnet*", + "claude-3-5-sonnet@*", + "global.anthropic.claude-3-5-sonnet*", + "claude-3-5-sonnet-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2024-06-20" + }, + { + "id": "claude-3-sonnet", + "provider": "anthropic", + "aliases": [ + "claude-3-sonnet-2*", + "anthropic/claude-3-sonnet*", + "anthropic.claude-3-sonnet*", + "us.anthropic.claude-3-sonnet*", + "eu.anthropic.claude-3-sonnet*", + "apac.anthropic.claude-3-sonnet*", + "claude-3-sonnet@*", + "global.anthropic.claude-3-sonnet*", + "claude-3-sonnet-v*@*" + ], + "input_per_mtok": 3, + "output_per_mtok": 15, + "effective_date": "2024-02-29" + }, + { + "id": "claude-haiku-4-5", + "provider": "anthropic", + "aliases": [ + "claude-haiku-4-5-2*", + "anthropic/claude-haiku-4-5*", + "anthropic.claude-haiku-4-5", + "anthropic.claude-haiku-4-5-*", + "us.anthropic.claude-haiku-4-5*", + "eu.anthropic.claude-haiku-4-5*", + "apac.anthropic.claude-haiku-4-5*", + "claude-haiku-4-5@*", + "global.anthropic.claude-haiku-4-5*", + "claude-haiku-4-5-v*@*" + ], + "input_per_mtok": 1, + "output_per_mtok": 5, + "effective_date": "2025-10-01" + }, + { + "id": "claude-3-5-haiku", + "provider": "anthropic", + "aliases": [ + "claude-3-5-haiku-2*", + "anthropic/claude-3-5-haiku*", + "anthropic.claude-3-5-haiku*", + "us.anthropic.claude-3-5-haiku*", + "eu.anthropic.claude-3-5-haiku*", + "apac.anthropic.claude-3-5-haiku*", + "claude-3-5-haiku@*", + "global.anthropic.claude-3-5-haiku*", + "claude-3-5-haiku-v*@*" + ], + "input_per_mtok": 0.8, + "output_per_mtok": 4, + "effective_date": "2024-10-22" + }, + { + "id": "claude-3-haiku", + "provider": "anthropic", + "aliases": [ + "claude-3-haiku-2*", + "anthropic/claude-3-haiku*", + "anthropic.claude-3-haiku*", + "us.anthropic.claude-3-haiku*", + "eu.anthropic.claude-3-haiku*", + "apac.anthropic.claude-3-haiku*", + "claude-3-haiku@*", + "global.anthropic.claude-3-haiku*", + "claude-3-haiku-v*@*" + ], + "input_per_mtok": 0.25, + "output_per_mtok": 1.25, + "effective_date": "2024-03-07" + } + ] +} \ No newline at end of file diff --git a/cmd/entire/cli/pricing/models/cursor.json b/cmd/entire/cli/pricing/models/cursor.json new file mode 100644 index 0000000000..2d5a13b429 --- /dev/null +++ b/cmd/entire/cli/pricing/models/cursor.json @@ -0,0 +1,57 @@ +{ + "schema_version": 1, + "models": [ + { + "id": "composer-2", + "provider": "cursor", + "aliases": [ + "cursor/composer-2", + "composer-2-2*" + ], + "input_per_mtok": 0.5, + "output_per_mtok": 2.5, + "cache_read_per_mtok": 0.05, + "cache_write_per_mtok": 0.5, + "effective_date": "2026-07-10" + }, + { + "id": "composer-2-fast", + "provider": "cursor", + "aliases": [ + "cursor/composer-2-fast", + "composer-2-fast-2*" + ], + "input_per_mtok": 1.5, + "output_per_mtok": 7.5, + "cache_read_per_mtok": 0.15, + "cache_write_per_mtok": 1.5, + "effective_date": "2026-07-10" + }, + { + "id": "composer-2.5", + "provider": "cursor", + "aliases": [ + "cursor/composer-2.5", + "composer-2.5-2*" + ], + "input_per_mtok": 0.5, + "output_per_mtok": 2.5, + "cache_read_per_mtok": 0.05, + "cache_write_per_mtok": 0.5, + "effective_date": "2026-07-10" + }, + { + "id": "composer-2.5-fast", + "provider": "cursor", + "aliases": [ + "cursor/composer-2.5-fast", + "composer-2.5-fast-2*" + ], + "input_per_mtok": 3.0, + "output_per_mtok": 15.0, + "cache_read_per_mtok": 0.3, + "cache_write_per_mtok": 3.0, + "effective_date": "2026-07-10" + } + ] +} diff --git a/cmd/entire/cli/pricing/models/google.json b/cmd/entire/cli/pricing/models/google.json new file mode 100644 index 0000000000..09dfe484f6 --- /dev/null +++ b/cmd/entire/cli/pricing/models/google.json @@ -0,0 +1,45 @@ +{ + "schema_version": 1, + "models": [ + { + "id": "gemini-3-pro", + "provider": "google", + "aliases": ["gemini-3-pro-*", "google/gemini-3-pro"], + "input_per_mtok": 2, + "output_per_mtok": 12, + "cache_read_per_mtok": 0.2, + "cache_write_per_mtok": 2, + "effective_date": "2026-01-01" + }, + { + "id": "gemini-3-flash", + "provider": "google", + "aliases": ["gemini-3-flash-*", "google/gemini-3-flash"], + "input_per_mtok": 0.3, + "output_per_mtok": 2.5, + "cache_read_per_mtok": 0.03, + "cache_write_per_mtok": 0.3, + "effective_date": "2026-01-01" + }, + { + "id": "gemini-3.1-pro", + "provider": "google", + "aliases": ["gemini-3.1-pro-*", "google/gemini-3.1-pro"], + "input_per_mtok": 2, + "output_per_mtok": 12, + "cache_read_per_mtok": 0.2, + "cache_write_per_mtok": 2, + "effective_date": "2026-07-10" + }, + { + "id": "gemini-3.5-flash", + "provider": "google", + "aliases": ["gemini-3.5-flash-*", "google/gemini-3.5-flash"], + "input_per_mtok": 1.5, + "output_per_mtok": 9, + "cache_read_per_mtok": 0.15, + "cache_write_per_mtok": 1.5, + "effective_date": "2026-07-10" + } + ] +} diff --git a/cmd/entire/cli/pricing/models/openai.json b/cmd/entire/cli/pricing/models/openai.json new file mode 100644 index 0000000000..2d49373bec --- /dev/null +++ b/cmd/entire/cli/pricing/models/openai.json @@ -0,0 +1,122 @@ +{ + "schema_version": 1, + "models": [ + { + "id": "gpt-5.5", + "provider": "openai", + "aliases": [ + "gpt-5.5-2*", + "openai/gpt-5.5" + ], + "input_per_mtok": 5, + "output_per_mtok": 30, + "cache_read_per_mtok": 0.5, + "cache_write_per_mtok": 6.25, + "effective_date": "2026-07-10" + }, + { + "id": "gpt-5.5-priority", + "provider": "openai", + "aliases": [ + "gpt-5.5-priority-2*", + "openai/gpt-5.5-priority" + ], + "input_per_mtok": 12.5, + "output_per_mtok": 75, + "cache_read_per_mtok": 1.25, + "cache_write_per_mtok": 15.625, + "effective_date": "2026-07-10" + }, + { + "id": "gpt-5.4", + "provider": "openai", + "aliases": [ + "gpt-5.4-*", + "openai/gpt-5.4" + ], + "input_per_mtok": 2.5, + "output_per_mtok": 15, + "cache_read_per_mtok": 0.25, + "cache_write_per_mtok": 3.125, + "effective_date": "2026-07-10" + }, + { + "id": "gpt-5.3-codex", + "provider": "openai", + "aliases": [ + "gpt-5.3-codex-2*", + "openai/gpt-5.3-codex", + "openai/gpt-5.3-codex-2*" + ], + "input_per_mtok": 1.75, + "output_per_mtok": 14, + "cache_read_per_mtok": 0.175, + "cache_write_per_mtok": 2.1875, + "effective_date": "2026-07-10" + }, + { + "id": "gpt-5", + "provider": "openai", + "aliases": [ + "gpt-5-*", + "openai/gpt-5" + ], + "input_per_mtok": 1.25, + "output_per_mtok": 10, + "cache_read_per_mtok": 0.125, + "cache_write_per_mtok": 1.25, + "effective_date": "2025-08-07" + }, + { + "id": "gpt-5.6-sol", + "provider": "openai", + "aliases": [ + "gpt-5.6-sol-2*", + "openai/gpt-5.6-sol", + "openai/gpt-5.6-sol-2*", + "gpt-5.6-sol@*", + "global.openai.gpt-5.6-sol", + "global.openai.gpt-5.6-sol-2*" + ], + "input_per_mtok": 5.0, + "output_per_mtok": 30.0, + "cache_read_per_mtok": 0.5, + "cache_write_per_mtok": 6.25, + "effective_date": "2026-07-09" + }, + { + "id": "gpt-5.6-terra", + "provider": "openai", + "aliases": [ + "gpt-5.6-terra-2*", + "openai/gpt-5.6-terra", + "openai/gpt-5.6-terra-2*", + "gpt-5.6-terra@*", + "global.openai.gpt-5.6-terra", + "global.openai.gpt-5.6-terra-2*" + ], + "input_per_mtok": 2.5, + "output_per_mtok": 15.0, + "cache_read_per_mtok": 0.25, + "cache_write_per_mtok": 3.125, + "effective_date": "2026-07-09" + }, + { + "id": "gpt-5.6-luna", + "provider": "openai", + "aliases": [ + "gpt-5.6-luna-2*", + "openai/gpt-5.6-luna", + "openai/gpt-5.6-luna-2*", + "gpt-5.6-luna@*", + "global.openai.gpt-5.6-luna", + "global.openai.gpt-5.6-luna-2*" + ], + "input_per_mtok": 1.0, + "output_per_mtok": 6.0, + "cache_read_per_mtok": 0.1, + "cache_write_per_mtok": 1.25, + "effective_date": "2026-07-09" + } + ] +} diff --git a/cmd/entire/cli/pricing/remote.go b/cmd/entire/cli/pricing/remote.go new file mode 100644 index 0000000000..ff93322e2a --- /dev/null +++ b/cmd/entire/cli/pricing/remote.go @@ -0,0 +1,346 @@ +package pricing + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "net/http" + "os" + "path/filepath" + "time" + + "github.com/entireio/cli/cmd/entire/cli/internal/flock" + "github.com/entireio/cli/cmd/entire/cli/logging" + "github.com/entireio/cli/cmd/entire/cli/versioninfo" + "github.com/entireio/cli/internal/entireclient/userdirs" +) + +// RemoteCache is the on-disk shape of the cached remote pricing table. It wraps +// the same fileSchema the embedded models/*.json files use — so validateRate and +// the schema_version contract are shared, not reimplemented — with the fetch +// bookkeeping the background refresh needs: FetchedAt drives the 24h refresh +// backoff, ETag drives conditional (If-None-Match) requests, SourceURL records +// where the table came from for diagnostics. +type RemoteCache struct { + FetchedAt time.Time `json:"fetched_at"` + ETag string `json:"etag,omitempty"` + SourceURL string `json:"source_url,omitempty"` + Doc *fileSchema `json:"doc,omitempty"` +} + +// remoteCacheFileName is the cache file basename. It lives beside the discovery +// caches under the per-user cache dir (userdirs.Cache()). +const remoteCacheFileName = "pricing_remote.json" + +// remoteCachePath returns the absolute path of the remote pricing cache file. +// Path resolution goes through userdirs.Cache — the single implementation every +// cache consumer shares — so the remote pricing cache sits under ~/.cache/entire +// (or $XDG_CACHE_HOME/entire) with nodes.json and the other discovery caches. +func remoteCachePath() string { + return filepath.Join(userdirs.Cache(), remoteCacheFileName) +} + +// loadRemoteCache reads and parses the cache file. A missing file, a corrupt +// file, or any read error all return nil: the remote layer is purely additive, +// so "no usable cache" degrades to "embedded defaults only" and a damaged cache +// self-heals on the next successful refresh (which overwrites it atomically) +// rather than being rewritten here on the read path. +func loadRemoteCache() *RemoteCache { + data, err := os.ReadFile(remoteCachePath()) // #nosec G304 -- path derived from userdirs, not user input + if err != nil { + return nil + } + var rc RemoteCache + if err := json.Unmarshal(data, &rc); err != nil { + return nil + } + return &rc +} + +// LoadRemoteEntries returns the valid model rates from the cached remote pricing +// table, or nil when there is nothing usable. It is read-only and performs NO +// network I/O — the daily refresh (RefreshRemoteCache) is what populates the +// cache. It never errors: a missing or corrupt cache, an absent Doc, an +// unsupported schema_version, and individually invalid entries all degrade to +// "fewer (or zero) remote entries", so LoadPricingTable can layer whatever is +// valid on top of the embedded defaults without a failure path. +func LoadRemoteEntries(ctx context.Context) []ModelRate { + rc := loadRemoteCache() + if rc == nil || rc.Doc == nil { + return nil + } + if rc.Doc.SchemaVersion != 1 { + logging.Debug(ctx, "pricing: ignoring remote cache with unsupported schema_version", + slog.Int("schema_version", rc.Doc.SchemaVersion)) + return nil + } + valid := make([]ModelRate, 0, len(rc.Doc.Models)) + for i := range rc.Doc.Models { + m := rc.Doc.Models[i] + if err := validateRate(m); err != nil { + // Drop the single bad entry, keep the rest: one malformed row from + // the remote table must not sink the whole merge. + logging.Debug(ctx, "pricing: dropping invalid remote entry", + slog.String("id", m.ID), slog.String("error", err.Error())) + continue + } + valid = append(valid, m) + } + if len(valid) == 0 { + return nil + } + return valid +} + +// RemoteFetchedAt returns the time the remote pricing cache was last written, or +// the zero time when there is no cache. It exists for staleness diagnostics +// (the LoadPricingTable debug log and the manual refresh command). +func RemoteFetchedAt() time.Time { + rc := loadRemoteCache() + if rc == nil { + return time.Time{} + } + return rc.FetchedAt +} + +// RemoteURL is the default source for the remote pricing table. It is a var so +// tests can point it at an httptest server; the ENTIRE_PRICING_URL environment +// variable overrides it at fetch time for self-hosted setups. +var RemoteURL = "https://entire.io/pricing/v1/models.json" + +const ( + // remoteRefreshInterval is the minimum cache age before a refresh is + // attempted. Every fetch outcome — success, 304, or failure — resets + // FetchedAt, so a broken endpoint is retried at most once per interval. + remoteRefreshInterval = 24 * time.Hour + // remoteMaxBytes caps the response body read (1 MiB), mirroring versioncheck. + remoteMaxBytes = 1 << 20 +) + +// remoteFetchTimeout bounds the whole refresh request (dial + read), matching +// the version check's 2s budget so a slow endpoint never stalls a background +// worker. It is a var so tests can shrink it to exercise the timeout path fast. +var remoteFetchTimeout = 2 * time.Second + +// Refresh outcomes reported by RefreshResult.Outcome. +const ( + refreshOutcomeUpdated = "updated" // 200 with a usable doc: Doc+ETag replaced. + refreshOutcomeNotModified = "not-modified" // 304: prior Doc kept, FetchedAt bumped. + refreshOutcomeUnchanged = "unchanged" // error/404/garbage: prior Doc kept, FetchedAt bumped. + refreshOutcomeSkipped = "skipped" // throttled: fresh cache, no request made. +) + +// RefreshResult summarizes what a refresh attempt did, for the manual +// `entire tokens pricing-refresh` report. +type RefreshResult struct { + SourceURL string + Outcome string + Entries int + FetchedAt time.Time + ETag string +} + +// Staleness returns how old the cache is relative to now, or 0 when FetchedAt +// is unset. +func (r RefreshResult) Staleness() time.Duration { + if r.FetchedAt.IsZero() { + return 0 + } + return time.Since(r.FetchedAt) +} + +// remoteURL returns the effective source URL, honoring the ENTIRE_PRICING_URL +// override read at fetch time. +func remoteURL() string { + if u := os.Getenv("ENTIRE_PRICING_URL"); u != "" { + return u + } + return RemoteURL +} + +// ShouldRefresh reports whether the remote pricing cache is stale (older than +// the 24h interval) or absent. Trigger sites call it to gate a background +// refresh spawn so a fresh cache skips the work — and the spawn — entirely. +func ShouldRefresh() bool { + rc := loadRemoteCache() + if rc == nil { + return true + } + return time.Since(rc.FetchedAt) >= remoteRefreshInterval +} + +// RefreshRemoteCache fetches the remote pricing table and updates the local +// cache, honoring the 24h throttle. It is the entry point the detached +// __refresh_pricing worker calls. See refreshRemoteCache for the full contract. +func RefreshRemoteCache(ctx context.Context) error { + _, err := refreshRemoteCache(ctx, false) + return err +} + +// RefreshRemoteCacheForce is RefreshRemoteCache without the freshness throttle: +// it always performs the network fetch and returns a RefreshResult describing +// the outcome. The manual `entire tokens pricing-refresh` command uses it. +func RefreshRemoteCacheForce(ctx context.Context) (RefreshResult, error) { + return refreshRemoteCache(ctx, true) +} + +// refreshRemoteCache is the shared implementation. It takes a cross-process +// flock for herd control, so several concurrently-spawned workers collapse to a +// single network request: the first takes the lock and refreshes, the rest take +// the lock, observe the freshly-written FetchedAt, and skip (unless force). +// +// Failure is absorbed, not surfaced: a network error, timeout, non-200 status, +// oversized/garbage body, or a doc that fails the schema/sanity check all keep +// the previous Doc+ETag and only bump FetchedAt — so a broken endpoint degrades +// to "keep serving the last good table" and is retried at most once per 24h. A +// 304 keeps the Doc and bumps FetchedAt. It returns an error only for a +// genuinely local failure (cache-dir create, lock, or write). +func refreshRemoteCache(ctx context.Context, force bool) (RefreshResult, error) { + res := RefreshResult{SourceURL: remoteURL(), Outcome: refreshOutcomeSkipped} + + // Cheap pre-lock throttle: skip the lock entirely when the cache is fresh + // and we are not forcing. The under-lock re-check below is the authoritative + // herd-control guard; this just avoids lock contention on the common path. + if !force && !ShouldRefresh() { + fillResultFromCache(&res, loadRemoteCache()) + return res, nil + } + + if err := os.MkdirAll(userdirs.Cache(), 0o700); err != nil { + return res, fmt.Errorf("create cache dir: %w", err) + } + path := remoteCachePath() + release, err := flock.Acquire(path + ".lock") + if err != nil { + return res, fmt.Errorf("lock remote pricing cache: %w", err) + } + defer release() + + // Re-read under the lock. If another process refreshed while we waited for + // the lock, its FetchedAt is now fresh and we skip the fetch (unless forced). + // This is what keeps a burst of spawned workers to a single request. + prev := loadRemoteCache() + if !force && prev != nil && time.Since(prev.FetchedAt) < remoteRefreshInterval { + fillResultFromCache(&res, prev) + return res, nil + } + if prev == nil { + prev = &RemoteCache{} + } + + updated, outcome := fetchRemote(ctx, prev) + if err := writeRemoteCache(path, updated); err != nil { + return res, err + } + res.Outcome = outcome + fillResultFromCache(&res, updated) + return res, nil +} + +// fillResultFromCache copies the reporting fields out of a cache snapshot. +func fillResultFromCache(res *RefreshResult, rc *RemoteCache) { + if rc == nil { + return + } + res.FetchedAt = rc.FetchedAt + res.ETag = rc.ETag + res.Entries = countValidEntries(rc.Doc) +} + +// fetchRemote performs the conditional HTTP request and folds the outcome into a +// new RemoteCache derived from prev. Every outcome bumps FetchedAt. It never +// returns an error — the worst case is "keep prev's Doc+ETag". +func fetchRemote(ctx context.Context, prev *RemoteCache) (*RemoteCache, string) { + url := remoteURL() + out := &RemoteCache{ + FetchedAt: time.Now(), + ETag: prev.ETag, + SourceURL: url, + Doc: prev.Doc, + } + + ctx, cancel := context.WithTimeout(ctx, remoteFetchTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return out, refreshOutcomeUnchanged + } + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", versioninfo.UserAgent()) + if prev.ETag != "" { + req.Header.Set("If-None-Match", prev.ETag) + } + + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + return out, refreshOutcomeUnchanged + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotModified { + return out, refreshOutcomeNotModified + } + if resp.StatusCode != http.StatusOK { + return out, refreshOutcomeUnchanged + } + + body, err := io.ReadAll(io.LimitReader(resp.Body, remoteMaxBytes)) + if err != nil { + return out, refreshOutcomeUnchanged + } + var doc fileSchema + if err := json.Unmarshal(body, &doc); err != nil { + return out, refreshOutcomeUnchanged + } + if !remoteDocUsable(&doc) { + return out, refreshOutcomeUnchanged + } + out.Doc = &doc + out.ETag = resp.Header.Get("ETag") + return out, refreshOutcomeUpdated +} + +// remoteDocUsable is the sanity gate a freshly-fetched doc must pass before it +// replaces the cached table: the supported schema_version and at least one entry +// that survives validateRate. This mirrors LoadRemoteEntries' read-time contract +// so a doc that would contribute zero usable rates is never stored. +func remoteDocUsable(doc *fileSchema) bool { + return countValidEntries(doc) > 0 +} + +// countValidEntries counts the entries in doc that pass validateRate under the +// supported schema_version; a nil or wrong-version doc counts zero. +func countValidEntries(doc *fileSchema) int { + if doc == nil || doc.SchemaVersion != 1 { + return 0 + } + n := 0 + for i := range doc.Models { + if validateRate(doc.Models[i]) == nil { + n++ + } + } + return n +} + +// writeRemoteCache writes rc atomically (temp + rename) so a concurrent reader +// never observes a half-written file. +func writeRemoteCache(path string, rc *RemoteCache) error { + data, err := json.MarshalIndent(rc, "", " ") + if err != nil { + return fmt.Errorf("marshal remote pricing cache: %w", err) + } + tmp := path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return fmt.Errorf("write remote pricing cache tmp: %w", err) + } + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + return fmt.Errorf("rename remote pricing cache: %w", err) + } + return nil +} diff --git a/cmd/entire/cli/pricing/remote_refresh_test.go b/cmd/entire/cli/pricing/remote_refresh_test.go new file mode 100644 index 0000000000..58107f8666 --- /dev/null +++ b/cmd/entire/cli/pricing/remote_refresh_test.go @@ -0,0 +1,235 @@ +package pricing + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const validRemoteDoc = `{"schema_version":1,"models":[` + + `{"id":"gpt-5.5","provider":"openai","input_per_mtok":999,"output_per_mtok":1000}]}` + +// priorInputRate is the gpt-5.5 input rate a seeded (pre-existing) cache carries, +// distinct from the embedded (5) and the fetched (999) rates so a test can tell +// which one survived. +const priorInputRate = 111.0 + +// seedStaleCache writes a well-formed cache with a Doc pricing gpt-5.5 at +// priorInputRate and a FetchedAt old enough to be stale, so a refresh proceeds. +func seedStaleCache(t *testing.T, etag string) { + t.Helper() + writeRemoteCacheFile(t, &RemoteCache{ + FetchedAt: time.Now().Add(-48 * time.Hour), + ETag: etag, + Doc: &fileSchema{ + SchemaVersion: 1, + Models: []ModelRate{ + {ID: "gpt-5.5", Provider: "openai", InputPerMTok: priorInputRate, OutputPerMTok: 1000}, + }, + }, + }) +} + +func priorInput(t *testing.T) float64 { + t.Helper() + rc := loadRemoteCache() + require.NotNil(t, rc) + require.NotNil(t, rc.Doc) + require.NotEmpty(t, rc.Doc.Models) + return rc.Doc.Models[0].InputPerMTok +} + +func TestRefreshRemoteCache_200StoresDoc(t *testing.T) { + isolateCache(t) + var hits atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + w.Header().Set("ETag", `"v1"`) + w.WriteHeader(http.StatusOK) + //nolint:errcheck // test helper + io.WriteString(w, validRemoteDoc) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + require.NoError(t, RefreshRemoteCache(context.Background())) + assert.Equal(t, int32(1), hits.Load()) + + rc := loadRemoteCache() + require.NotNil(t, rc) + require.NotNil(t, rc.Doc) + assert.Equal(t, 1, rc.Doc.SchemaVersion) + assert.Equal(t, `"v1"`, rc.ETag) + assert.False(t, rc.FetchedAt.IsZero()) + assert.Equal(t, srv.URL, rc.SourceURL) + + entries := LoadRemoteEntries(context.Background()) + require.Len(t, entries, 1) + assert.InDelta(t, 999.0, entries[0].InputPerMTok, 1e-9) +} + +func TestRefreshRemoteCache_404KeepsPrior(t *testing.T) { + isolateCache(t) + seedStaleCache(t, "seed") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + res, err := RefreshRemoteCacheForce(context.Background()) + require.NoError(t, err) + assert.Equal(t, refreshOutcomeUnchanged, res.Outcome) + assert.InDelta(t, priorInputRate, priorInput(t), 1e-9, "prior doc must be kept on 404") + assert.WithinDuration(t, time.Now(), loadRemoteCache().FetchedAt, time.Minute, "FetchedAt must be bumped") +} + +func TestRefreshRemoteCache_TimeoutKeepsPrior(t *testing.T) { + isolateCache(t) + seedStaleCache(t, "seed") + old := remoteFetchTimeout + remoteFetchTimeout = 100 * time.Millisecond + defer func() { remoteFetchTimeout = old }() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(400 * time.Millisecond) + w.WriteHeader(http.StatusOK) + //nolint:errcheck // test helper; client has already timed out + io.WriteString(w, validRemoteDoc) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + res, err := RefreshRemoteCacheForce(context.Background()) + require.NoError(t, err) + assert.Equal(t, refreshOutcomeUnchanged, res.Outcome) + assert.InDelta(t, priorInputRate, priorInput(t), 1e-9, "prior doc must be kept on timeout") + assert.WithinDuration(t, time.Now(), loadRemoteCache().FetchedAt, time.Minute) +} + +func TestRefreshRemoteCache_GarbageKeepsPrior(t *testing.T) { + isolateCache(t) + seedStaleCache(t, "seed") + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + //nolint:errcheck // test helper + io.WriteString(w, "{ this is not valid json") + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + res, err := RefreshRemoteCacheForce(context.Background()) + require.NoError(t, err) + assert.Equal(t, refreshOutcomeUnchanged, res.Outcome) + assert.InDelta(t, priorInputRate, priorInput(t), 1e-9, "prior doc must be kept on garbage body") +} + +func TestRefreshRemoteCache_304KeepsDocBumpsFetchedAt(t *testing.T) { + isolateCache(t) + seedStaleCache(t, `"seed-etag"`) + var sawIfNoneMatch string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sawIfNoneMatch = r.Header.Get("If-None-Match") + if sawIfNoneMatch == `"seed-etag"` { + w.WriteHeader(http.StatusNotModified) + return + } + w.WriteHeader(http.StatusOK) + //nolint:errcheck // test helper + io.WriteString(w, validRemoteDoc) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + res, err := RefreshRemoteCacheForce(context.Background()) + require.NoError(t, err) + assert.Equal(t, `"seed-etag"`, sawIfNoneMatch, "stored ETag must be sent as If-None-Match") + assert.Equal(t, refreshOutcomeNotModified, res.Outcome) + rc := loadRemoteCache() + assert.InDelta(t, priorInputRate, rc.Doc.Models[0].InputPerMTok, 1e-9, "prior doc must be kept on 304") + assert.Equal(t, `"seed-etag"`, rc.ETag, "ETag must be retained on 304") + assert.WithinDuration(t, time.Now(), rc.FetchedAt, time.Minute, "FetchedAt must be bumped on 304") +} + +func TestRefreshRemoteCache_SchemaVersion99NotStored(t *testing.T) { + isolateCache(t) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + //nolint:errcheck // test helper + io.WriteString(w, `{"schema_version":99,"models":[`+ + `{"id":"gpt-5.5","provider":"openai","input_per_mtok":999,"output_per_mtok":1000}]}`) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + res, err := RefreshRemoteCacheForce(context.Background()) + require.NoError(t, err) + assert.Equal(t, refreshOutcomeUnchanged, res.Outcome, "unsupported schema doc must not be stored") + assert.Nil(t, LoadRemoteEntries(context.Background()), "no mergeable entries from a v99 doc") +} + +func TestRefreshRemoteCache_ThrottleSkipsWhenFresh(t *testing.T) { + isolateCache(t) + // Fresh cache: FetchedAt now. + writeRemoteCacheFile(t, &RemoteCache{ + FetchedAt: time.Now(), + Doc: &fileSchema{ + SchemaVersion: 1, + Models: []ModelRate{{ID: "gpt-5.5", Provider: "openai", InputPerMTok: priorInputRate, OutputPerMTok: 1000}}, + }, + }) + var hits atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + require.NoError(t, RefreshRemoteCache(context.Background())) + assert.Equal(t, int32(0), hits.Load(), "a fresh cache must skip the fetch entirely") +} + +func TestRefreshRemoteCacheForce_BypassesThrottle(t *testing.T) { + isolateCache(t) + writeRemoteCacheFile(t, &RemoteCache{ + FetchedAt: time.Now(), + Doc: &fileSchema{ + SchemaVersion: 1, + Models: []ModelRate{{ID: "gpt-5.5", Provider: "openai", InputPerMTok: priorInputRate, OutputPerMTok: 1000}}, + }, + }) + var hits atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + w.Header().Set("ETag", `"v2"`) + w.WriteHeader(http.StatusOK) + //nolint:errcheck // test helper + io.WriteString(w, validRemoteDoc) + })) + defer srv.Close() + t.Setenv("ENTIRE_PRICING_URL", srv.URL) + + res, err := RefreshRemoteCacheForce(context.Background()) + require.NoError(t, err) + assert.Equal(t, int32(1), hits.Load(), "force must bypass the throttle and fetch") + assert.Equal(t, refreshOutcomeUpdated, res.Outcome) + assert.InDelta(t, 999.0, priorInput(t), 1e-9, "forced refresh must replace the doc") +} + +func TestShouldRefresh(t *testing.T) { + isolateCache(t) + assert.True(t, ShouldRefresh(), "missing cache should refresh") + + writeRemoteCacheFile(t, &RemoteCache{FetchedAt: time.Now()}) + assert.False(t, ShouldRefresh(), "fresh cache should not refresh") + + writeRemoteCacheFile(t, &RemoteCache{FetchedAt: time.Now().Add(-48 * time.Hour)}) + assert.True(t, ShouldRefresh(), "stale cache should refresh") +} diff --git a/cmd/entire/cli/pricing/remote_test.go b/cmd/entire/cli/pricing/remote_test.go new file mode 100644 index 0000000000..d9a2b3adfb --- /dev/null +++ b/cmd/entire/cli/pricing/remote_test.go @@ -0,0 +1,115 @@ +package pricing + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeRemoteCacheFile marshals rc to the isolated remote cache path. Call +// isolateCache(t) first so it lands under a throwaway XDG_CACHE_HOME. +func writeRemoteCacheFile(t *testing.T, rc *RemoteCache) { + t.Helper() + dir := filepath.Dir(remoteCachePath()) + require.NoError(t, os.MkdirAll(dir, 0o755)) + data, err := json.Marshal(rc) + require.NoError(t, err) + require.NoError(t, os.WriteFile(remoteCachePath(), data, 0o644)) +} + +// writeRemoteCacheRaw writes raw bytes to the cache path (for corrupt-file cases). +func writeRemoteCacheRaw(t *testing.T, data []byte) { + t.Helper() + dir := filepath.Dir(remoteCachePath()) + require.NoError(t, os.MkdirAll(dir, 0o755)) + require.NoError(t, os.WriteFile(remoteCachePath(), data, 0o644)) +} + +// isolateCache points userdirs.Cache at a throwaway dir for the test. +func isolateCache(t *testing.T) { + t.Helper() + t.Setenv("XDG_CACHE_HOME", t.TempDir()) +} + +func TestLoadRemoteEntries_MissingCache(t *testing.T) { + isolateCache(t) + assert.Nil(t, LoadRemoteEntries(context.Background())) +} + +func TestLoadRemoteEntries_CorruptCache(t *testing.T) { + isolateCache(t) + writeRemoteCacheRaw(t, []byte("{not json")) + assert.Nil(t, LoadRemoteEntries(context.Background())) +} + +func TestLoadRemoteEntries_NilDoc(t *testing.T) { + isolateCache(t) + writeRemoteCacheFile(t, &RemoteCache{FetchedAt: time.Now()}) + assert.Nil(t, LoadRemoteEntries(context.Background())) +} + +func TestLoadRemoteEntries_UnsupportedSchemaVersion(t *testing.T) { + isolateCache(t) + writeRemoteCacheFile(t, &RemoteCache{ + FetchedAt: time.Now(), + Doc: &fileSchema{ + SchemaVersion: 99, + Models: []ModelRate{ + {ID: "gpt-5.5", Provider: "openai", InputPerMTok: 7, OutputPerMTok: 11}, + }, + }, + }) + assert.Nil(t, LoadRemoteEntries(context.Background())) +} + +func TestLoadRemoteEntries_DropsInvalidKeepsValid(t *testing.T) { + isolateCache(t) + writeRemoteCacheFile(t, &RemoteCache{ + FetchedAt: time.Now(), + Doc: &fileSchema{ + SchemaVersion: 1, + Models: []ModelRate{ + {ID: "good-model", Provider: "test", InputPerMTok: 3, OutputPerMTok: 9}, + {ID: "bad-zero-input", Provider: "test", InputPerMTok: 0, OutputPerMTok: 9}, + {ID: "", Provider: "test", InputPerMTok: 3, OutputPerMTok: 9}, + }, + }, + }) + entries := LoadRemoteEntries(context.Background()) + require.Len(t, entries, 1) + assert.Equal(t, "good-model", entries[0].ID) +} + +func TestLoadRemoteEntries_ValidDoc(t *testing.T) { + isolateCache(t) + writeRemoteCacheFile(t, &RemoteCache{ + FetchedAt: time.Now(), + Doc: &fileSchema{ + SchemaVersion: 1, + Models: []ModelRate{ + {ID: "gpt-5.5", Provider: "openai", InputPerMTok: 999, OutputPerMTok: 1000}, + {ID: "new-model", Provider: "test", InputPerMTok: 1, OutputPerMTok: 2}, + }, + }, + }) + entries := LoadRemoteEntries(context.Background()) + require.Len(t, entries, 2) + assert.Equal(t, "gpt-5.5", entries[0].ID) + assert.InDelta(t, 999.0, entries[0].InputPerMTok, 1e-9) +} + +func TestRemoteFetchedAt(t *testing.T) { + isolateCache(t) + assert.True(t, RemoteFetchedAt().IsZero(), "missing cache should report zero time") + + when := time.Now().Add(-3 * time.Hour).UTC().Truncate(time.Second) + writeRemoteCacheFile(t, &RemoteCache{FetchedAt: when}) + got := RemoteFetchedAt().UTC().Truncate(time.Second) + assert.True(t, got.Equal(when), "RemoteFetchedAt = %v, want %v", got, when) +} diff --git a/cmd/entire/cli/pricing/table.go b/cmd/entire/cli/pricing/table.go new file mode 100644 index 0000000000..df14308ebb --- /dev/null +++ b/cmd/entire/cli/pricing/table.go @@ -0,0 +1,302 @@ +package pricing + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io/fs" + "log/slog" + "path" + "strings" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/logging" +) + +// ModelRate is the price of a single model, in USD per million tokens (MTok). +// CacheReadPerMTok and CacheWritePerMTok are optional: when nil, Estimate +// derives them from InputPerMTok using the Anthropic cache multipliers. +type ModelRate struct { + // ID is the canonical model identifier (e.g. "claude-opus-4-8"). + ID string `json:"id"` + // Provider is the vendor that serves the model (e.g. "anthropic"). + Provider string `json:"provider"` + // Aliases are additional identifiers that resolve to this model. Each may + // be a literal id or a shell-style glob (e.g. "claude-opus-4-8-*"). + Aliases []string `json:"aliases"` + // InputPerMTok is the price of fresh (non-cached) input tokens per MTok. + InputPerMTok float64 `json:"input_per_mtok"` + // OutputPerMTok is the price of output tokens per MTok. + OutputPerMTok float64 `json:"output_per_mtok"` + // CacheReadPerMTok is the price of cache-read tokens per MTok, if set. + CacheReadPerMTok *float64 `json:"cache_read_per_mtok"` + // CacheWritePerMTok is the price of cache-write tokens per MTok, if set. + CacheWritePerMTok *float64 `json:"cache_write_per_mtok"` + // EffectiveDate is the ISO date the price took effect (informational). + EffectiveDate string `json:"effective_date"` +} + +// fileSchema is the on-disk shape of each embedded models/*.json file. +type fileSchema struct { + SchemaVersion int `json:"schema_version"` + Models []ModelRate `json:"models"` +} + +// Table is a resolved set of model rates, keyed for lookup by id and alias. +type Table struct { + models []ModelRate + index map[string]int +} + +const ( + // AnthropicCacheReadMultiplier is the fraction of the input price charged + // for cache-read tokens when a model omits an explicit cache-read rate. + AnthropicCacheReadMultiplier = 0.1 + // AnthropicCacheWriteMultiplier is the multiple of the input price charged + // for cache-write tokens when a model omits an explicit cache-write rate. + AnthropicCacheWriteMultiplier = 1.25 +) + +// LoadTable parses every embedded pricing file, validates each entry, then +// layers the given overrides on top: an override whose id matches an existing +// entry replaces it, otherwise it is appended. Overrides are validated too. A +// nil or empty overrides slice yields the embedded defaults unchanged. +func LoadTable(overrides []ModelRate) (*Table, error) { + entries, err := fs.ReadDir(modelsFS, "models") + if err != nil { + return nil, fmt.Errorf("reading embedded pricing models: %w", err) + } + + t := &Table{index: make(map[string]int)} + for _, e := range entries { + if e.IsDir() { + continue + } + data, err := modelsFS.ReadFile(path.Join("models", e.Name())) + if err != nil { + return nil, fmt.Errorf("reading embedded pricing model %s: %w", e.Name(), err) + } + var schema fileSchema + if err := json.Unmarshal(data, &schema); err != nil { + return nil, fmt.Errorf("parsing pricing model %s: %w", e.Name(), err) + } + for i := range schema.Models { + m := schema.Models[i] + if err := validateRate(m); err != nil { + return nil, fmt.Errorf("invalid pricing entry in %s: %w", e.Name(), err) + } + t.add(m) + } + } + + for i := range overrides { + o := overrides[i] + if err := validateRate(o); err != nil { + return nil, fmt.Errorf("invalid pricing override: %w", err) + } + t.add(o) + } + + return t, nil +} + +// ValidOverrides returns the subset of overrides that pass validation, dropping +// and logging each invalid entry rather than failing on it. It mirrors +// LoadRemoteEntries so a single malformed user override in .entire/settings.json +// cannot sink the whole pricing table: LoadTable hard-errors on any invalid +// override, and LoadPricingTable would then discard the entire table and disable +// cost estimation. Callers pre-filter user overrides through this before handing +// them to LoadTable. Order is preserved so LoadTable's last-writer-wins layering +// is unaffected. A nil/empty input yields nil. +func ValidOverrides(ctx context.Context, overrides []ModelRate) []ModelRate { + if len(overrides) == 0 { + return nil + } + valid := make([]ModelRate, 0, len(overrides)) + for i := range overrides { + o := overrides[i] + if err := validateRate(o); err != nil { + logging.Warn(ctx, "pricing: dropping invalid pricing override entry", + slog.String("id", o.ID), slog.String("error", err.Error())) + continue + } + valid = append(valid, o) + } + return valid +} + +// add inserts m, replacing any existing entry that shares its id. Ids are +// canonicalized (trimmed, lower-cased) for the index key so that an override +// whose id differs only in case or surrounding whitespace replaces the builtin +// entry instead of appending a phantom duplicate. The stored id is trimmed. When +// an override omits aliases, provider, or cache rates, it inherits them from +// the replaced entry so a minimal price override (id + input/output only) +// keeps resolving the dated and provider-prefixed spellings the embedded +// entry covered and keeps the provider's cache economics. +func (t *Table) add(m ModelRate) { + m.ID = strings.TrimSpace(m.ID) + key := strings.ToLower(m.ID) + if idx, ok := t.index[key]; ok { + prev := t.models[idx] + if len(m.Aliases) == 0 { + m.Aliases = prev.Aliases + } + if strings.TrimSpace(m.Provider) == "" { + m.Provider = prev.Provider + } + if m.CacheReadPerMTok == nil { + m.CacheReadPerMTok = prev.CacheReadPerMTok + } + if m.CacheWritePerMTok == nil { + m.CacheWritePerMTok = prev.CacheWritePerMTok + } + t.models[idx] = m + return + } + t.index[key] = len(t.models) + t.models = append(t.models, m) +} + +// Lookup resolves a model identifier to its rate. It first tries an exact match +// against a known id, then falls back to normalized (lower-cased, trimmed) +// matching against ids and aliases, where an alias may be a shell-style glob. +// It performs no implicit model-family fallback; an unknown model returns +// (ModelRate{}, false). +// +// Matching is attempted against the normalized query and, when the query carries +// a bracketed long-context suffix such as "[1m]", the same query with that +// suffix removed. path.Match treats "[...]" as a character class, so a literal +// "[1m]" alias would never match; normalizing on the query side sidesteps that +// trap and lets a bare id like "claude-fable-5[1m]" (the form Claude Code emits) +// resolve to "claude-fable-5", which bills at the same base rate. +func (t *Table) Lookup(model string) (ModelRate, bool) { + norm := strings.ToLower(strings.TrimSpace(model)) + if norm == "" { + return ModelRate{}, false + } + + // The index is keyed by canonical (trimmed, lower-cased) id, so the fast + // path must probe with the same canonicalization the query was normalized to. + if idx, ok := t.index[norm]; ok { + return t.models[idx], true + } + + candidates := []string{norm} + if base := stripLongContextSuffix(norm); base != norm && base != "" { + candidates = append(candidates, base) + } + + for i := range t.models { + r := t.models[i] + id := strings.ToLower(r.ID) + for _, q := range candidates { + if id == q { + return r, true + } + } + for _, alias := range r.Aliases { + na := strings.ToLower(strings.TrimSpace(alias)) + if na == "" { + continue + } + isGlob := strings.ContainsAny(na, "*?[") + for _, q := range candidates { + if isGlob { + if ok, err := path.Match(na, q); err == nil && ok { + return r, true + } + continue + } + if na == q { + return r, true + } + } + } + } + + return ModelRate{}, false +} + +// stripLongContextSuffix removes a trailing bracketed suffix such as the "[1m]" +// long-context marker that Claude Code appends to model ids (e.g. +// "claude-fable-5[1m]"). It returns s unchanged when there is no such suffix. +func stripLongContextSuffix(s string) string { + if strings.HasSuffix(s, "]") { + if i := strings.IndexByte(s, '['); i >= 0 { + return s[:i] + } + } + return s +} + +// Estimate returns the USD cost of u under rate r. Fresh input, cache-read, +// cache-write, and output tokens are each priced separately and summed. +// +// The nil-cache-rate fallbacks are provider-aware. The +// AnthropicCacheReadMultiplier (0.1x input) and AnthropicCacheWriteMultiplier +// (1.25x input) encode Anthropic economics, so they apply only when Provider is +// "anthropic" (case-insensitive). For any other provider a missing cache-read or +// cache-write rate falls back to the full input rate (1.0x) — billing cached +// tokens as normal input, which never undercharges. In practice the non-Anthropic +// embedded tables carry explicit cache rates, so this fallback is a safety net. +func Estimate(r ModelRate, u types.TokenUsage) float64 { + isAnthropic := strings.EqualFold(strings.TrimSpace(r.Provider), "anthropic") + + crRate := r.InputPerMTok + if isAnthropic { + crRate = AnthropicCacheReadMultiplier * r.InputPerMTok + } + if r.CacheReadPerMTok != nil { + crRate = *r.CacheReadPerMTok + } + cwRate := r.InputPerMTok + if isAnthropic { + cwRate = AnthropicCacheWriteMultiplier * r.InputPerMTok + } + if r.CacheWritePerMTok != nil { + cwRate = *r.CacheWritePerMTok + } + + input := float64(u.InputTokens) * r.InputPerMTok + cacheRead := float64(u.CacheReadTokens) * crRate + cacheWrite := float64(u.CacheCreationTokens) * cwRate + output := float64(u.OutputTokens) * r.OutputPerMTok + + return (input + cacheRead + cacheWrite + output) / 1e6 +} + +// validateRate rejects entries that would make an estimate meaningless: a +// missing id or a non-positive input/output rate. Explicit cache rates, when +// present, must not be negative. +func validateRate(r ModelRate) error { + if strings.TrimSpace(r.ID) == "" { + return errors.New("model id must not be empty") + } + if r.InputPerMTok <= 0 { + return fmt.Errorf("model %q: input_per_mtok must be positive, got %v", r.ID, r.InputPerMTok) + } + if r.OutputPerMTok <= 0 { + return fmt.Errorf("model %q: output_per_mtok must be positive, got %v", r.ID, r.OutputPerMTok) + } + if r.CacheReadPerMTok != nil && *r.CacheReadPerMTok < 0 { + return fmt.Errorf("model %q: cache_read_per_mtok must not be negative, got %v", r.ID, *r.CacheReadPerMTok) + } + if r.CacheWritePerMTok != nil && *r.CacheWritePerMTok < 0 { + return fmt.Errorf("model %q: cache_write_per_mtok must not be negative, got %v", r.ID, *r.CacheWritePerMTok) + } + for _, alias := range r.Aliases { + if strings.TrimSpace(alias) == "" { + return fmt.Errorf("model %q: alias must not be empty or whitespace", r.ID) + } + // A glob alias with malformed metacharacters (e.g. an unterminated "[") + // would silently never match at Lookup time (path.Match swallows + // ErrBadPattern there), so reject it up front where it is a config error. + if strings.ContainsAny(alias, "*?[") { + if _, err := path.Match(alias, "probe"); errors.Is(err, path.ErrBadPattern) { + return fmt.Errorf("model %q: invalid alias glob %q: %w", r.ID, alias, err) + } + } + } + return nil +} diff --git a/cmd/entire/cli/pricing/table_test.go b/cmd/entire/cli/pricing/table_test.go new file mode 100644 index 0000000000..f0d3b0b5b5 --- /dev/null +++ b/cmd/entire/cli/pricing/table_test.go @@ -0,0 +1,475 @@ +package pricing + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + modelOpus48 = "claude-opus-4-8" + modelGPT55 = "gpt-5.5" +) + +func floatPtr(f float64) *float64 { return &f } + +func TestLoadTable_EmbeddedParsesAndValidates(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + require.NotNil(t, tbl) + + // A representative model from each embedded provider file resolves. + for _, id := range []string{modelOpus48, modelGPT55, "gemini-3-pro", "composer-2"} { + _, ok := tbl.Lookup(id) + assert.Truef(t, ok, "expected embedded model %q to resolve", id) + } +} + +func TestTable_Lookup(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + tests := []struct { + name string + query string + wantID string + wantOK bool + }{ + {name: "exact id", query: modelOpus48, wantID: modelOpus48, wantOK: true}, + {name: "alias glob dated", query: "claude-opus-4-8-20260624", wantID: modelOpus48, wantOK: true}, + {name: "alias provider-prefixed", query: "anthropic/claude-opus-4-8", wantID: modelOpus48, wantOK: true}, + {name: "case insensitive id", query: "CLAUDE-OPUS-4-8", wantID: modelOpus48, wantOK: true}, + {name: "case insensitive with whitespace", query: " Claude-Opus-4-8 ", wantID: modelOpus48, wantOK: true}, + {name: "openai exact", query: modelGPT55, wantID: modelGPT55, wantOK: true}, + {name: "openai alias glob", query: "gpt-5.5-2026-01-01", wantID: modelGPT55, wantOK: true}, + {name: "miss returns false", query: "totally-unknown-model", wantOK: false}, + {name: "empty returns false", query: "", wantOK: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, ok := tbl.Lookup(tc.query) + assert.Equal(t, tc.wantOK, ok) + if tc.wantOK { + assert.Equal(t, tc.wantID, got.ID) + } + }) + } +} + +func TestTable_Lookup_AnthropicIDForms(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + tests := []struct { + name string + query string + wantID string + wantOK bool + }{ + {name: "bare id", query: "claude-opus-4-8", wantID: "claude-opus-4-8", wantOK: true}, + {name: "anthropic dated", query: "claude-haiku-4-5-20251001", wantID: "claude-haiku-4-5", wantOK: true}, + {name: "long-context 1m suffix on bare id", query: "claude-fable-5[1m]", wantID: "claude-fable-5", wantOK: true}, + {name: "long-context 1m suffix on dated id", query: "claude-opus-4-8-20260624[1m]", wantID: "claude-opus-4-8", wantOK: true}, + {name: "long-context 1m suffix uppercase", query: "CLAUDE-SONNET-5[1M]", wantID: "claude-sonnet-5", wantOK: true}, + {name: "bedrock bare", query: "anthropic.claude-opus-4-8", wantID: "claude-opus-4-8", wantOK: true}, + {name: "bedrock versioned regional", query: "us.anthropic.claude-opus-4-8-v1:0", wantID: "claude-opus-4-8", wantOK: true}, + {name: "bedrock legacy dated versioned", query: "anthropic.claude-3-5-sonnet-20241022-v2:0", wantID: "claude-3-5-sonnet", wantOK: true}, + {name: "bedrock global inference profile", query: "global.anthropic.claude-opus-4-8-v1:0", wantID: "claude-opus-4-8", wantOK: true}, + {name: "vertex dated", query: "claude-opus-4-5@20251101", wantID: "claude-opus-4-5", wantOK: true}, + {name: "vertex legacy versioned", query: "claude-3-5-sonnet-v2@20241022", wantID: "claude-3-5-sonnet", wantOK: true}, + {name: "slash prefixed", query: "anthropic/claude-sonnet-5", wantID: "claude-sonnet-5", wantOK: true}, + {name: "slash prefixed dated", query: "anthropic/claude-3-5-sonnet-20241022", wantID: "claude-3-5-sonnet", wantOK: true}, + {name: "legacy dated sonnet", query: "claude-3-5-sonnet-20241022", wantID: "claude-3-5-sonnet", wantOK: true}, + {name: "legacy opus canonical alias", query: "claude-3-opus", wantID: "claude-opus-3", wantOK: true}, + {name: "legacy opus dated alias", query: "claude-3-opus-20240229", wantID: "claude-opus-3", wantOK: true}, + {name: "legacy haiku dated", query: "claude-3-haiku-20240307", wantID: "claude-3-haiku", wantOK: true}, + {name: "gpt-5.6 sol resolves", query: "gpt-5.6-sol", wantID: "gpt-5.6-sol", wantOK: true}, + {name: "gpt-5.6 sol dated resolves", query: "gpt-5.6-sol-20260709", wantID: "gpt-5.6-sol", wantOK: true}, + {name: "gpt-5.6-sol-pro must NOT resolve to sol", query: "gpt-5.6-sol-pro", wantOK: false}, + {name: "prefixed gpt-5.6-sol-pro must NOT resolve to sol", query: "openai/gpt-5.6-sol-pro", wantOK: false}, + {name: "global-prefixed gpt-5.6-sol-pro must NOT resolve to sol", query: "global.openai.gpt-5.6-sol-pro", wantOK: false}, + {name: "prefixed gpt-5.6-sol exact resolves", query: "openai/gpt-5.6-sol", wantID: "gpt-5.6-sol", wantOK: true}, + {name: "prefixed gpt-5.6-sol dated resolves", query: "openai/gpt-5.6-sol-20260709", wantID: "gpt-5.6-sol", wantOK: true}, + {name: "non-claude id misses", query: "meta-llama-3-70b", wantOK: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, ok := tbl.Lookup(tc.query) + assert.Equal(t, tc.wantOK, ok) + if tc.wantOK { + assert.Equal(t, tc.wantID, got.ID) + } + }) + } +} + +func TestTable_Lookup_LongContextSharesBaseRate(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + // The "[1m]" long-context id resolves to the base model and bills at the + // base rate (no long-context premium). + got, ok := tbl.Lookup("claude-fable-5[1m]") + require.True(t, ok) + assert.Equal(t, "claude-fable-5", got.ID) + assert.InDelta(t, 10.0, got.InputPerMTok, 1e-9) + assert.InDelta(t, 50.0, got.OutputPerMTok, 1e-9) +} + +func TestTable_Round2CorrectedAndNewRates(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + // Corrected OpenAI rates: the table previously carried the legacy gpt-5 rate + // (1.25/10) for both gpt-5.5 and gpt-5.4. + gpt55, ok := tbl.Lookup("gpt-5.5") + require.True(t, ok) + assert.InDelta(t, 5.0, gpt55.InputPerMTok, 1e-9) + assert.InDelta(t, 30.0, gpt55.OutputPerMTok, 1e-9) + + gpt54, ok := tbl.Lookup("gpt-5.4") + require.True(t, ok) + assert.InDelta(t, 2.5, gpt54.InputPerMTok, 1e-9) + assert.InDelta(t, 15.0, gpt54.OutputPerMTok, 1e-9) + + // gpt-5 keeps its genuine legacy rate. + gpt5, ok := tbl.Lookup("gpt-5") + require.True(t, ok) + assert.InDelta(t, 1.25, gpt5.InputPerMTok, 1e-9) + assert.InDelta(t, 10.0, gpt5.OutputPerMTok, 1e-9) + + // New entries: gpt-5.3-codex, gpt-5.5-priority, and the cursor composer family. + codex, ok := tbl.Lookup("gpt-5.3-codex") + require.True(t, ok) + assert.InDelta(t, 1.75, codex.InputPerMTok, 1e-9) + assert.InDelta(t, 14.0, codex.OutputPerMTok, 1e-9) + + prio, ok := tbl.Lookup("gpt-5.5-priority") + require.True(t, ok) + assert.InDelta(t, 12.5, prio.InputPerMTok, 1e-9) + assert.InDelta(t, 75.0, prio.OutputPerMTok, 1e-9) + + composer, ok := tbl.Lookup("composer-2.5-fast") + require.True(t, ok) + assert.InDelta(t, 3.0, composer.InputPerMTok, 1e-9) + assert.InDelta(t, 15.0, composer.OutputPerMTok, 1e-9) +} + +func TestTable_GPT55PriorityDistinctFromBase(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + // Both are exact ids resolving to their own entries: the base gpt-5.5 glob + // must not swallow the priority id, nor vice versa. + base, ok := tbl.Lookup("gpt-5.5") + require.True(t, ok) + assert.Equal(t, "gpt-5.5", base.ID) + assert.InDelta(t, 5.0, base.InputPerMTok, 1e-9) + + prio, ok := tbl.Lookup("gpt-5.5-priority") + require.True(t, ok) + assert.Equal(t, "gpt-5.5-priority", prio.ID) + assert.InDelta(t, 12.5, prio.InputPerMTok, 1e-9) + + // A dated priority spelling resolves to the priority entry — it is not + // glob-captured by the base gpt-5.5 alias (tightened to "gpt-5.5-2*"). + datedPrio, ok := tbl.Lookup("gpt-5.5-priority-20260710") + require.True(t, ok) + assert.Equal(t, "gpt-5.5-priority", datedPrio.ID) + + // A dated base spelling still resolves to the base entry. + datedBase, ok := tbl.Lookup("gpt-5.5-20260710") + require.True(t, ok) + assert.Equal(t, "gpt-5.5", datedBase.ID) +} + +func TestTable_OpusFastVariant(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable(nil) + require.NoError(t, err) + + fast, ok := tbl.Lookup("claude-opus-4-8-fast") + require.True(t, ok) + assert.Equal(t, "claude-opus-4-8-fast", fast.ID) + assert.InDelta(t, 10.0, fast.InputPerMTok, 1e-9) + assert.InDelta(t, 50.0, fast.OutputPerMTok, 1e-9) + + // The fast entry carries no explicit cache rates, so Estimate derives them + // from the fast input rate via the Anthropic multipliers: cache-read 0.1x + // ($1.00 per MTok) and cache-write 1.25x ($12.50 per MTok). + require.Nil(t, fast.CacheReadPerMTok) + require.Nil(t, fast.CacheWritePerMTok) + assert.InDelta(t, 1.0, Estimate(fast, types.TokenUsage{CacheReadTokens: 1_000_000}), 1e-9) + assert.InDelta(t, 12.5, Estimate(fast, types.TokenUsage{CacheCreationTokens: 1_000_000}), 1e-9) + + // The dated fast spelling resolves to the fast entry, not the base + // claude-opus-4-8 (whose "claude-opus-4-8-2*" glob would otherwise capture + // it — the fast entry is ordered first so its "-fast" glob wins). + dated, ok := tbl.Lookup("claude-opus-4-8-20260624-fast") + require.True(t, ok) + assert.Equal(t, "claude-opus-4-8-fast", dated.ID) + + // A dated base spelling (no -fast) still resolves to the base entry. + datedBase, ok := tbl.Lookup("claude-opus-4-8-20260624") + require.True(t, ok) + assert.Equal(t, "claude-opus-4-8", datedBase.ID) + + // The OpenRouter dotted fast spelling resolves too. + dotted, ok := tbl.Lookup("anthropic/claude-opus-4.8-fast") + require.True(t, ok) + assert.Equal(t, "claude-opus-4-8-fast", dotted.ID) +} + +func TestEstimate_DefaultMultipliers(t *testing.T) { + t.Parallel() + + // claude-opus-4-8: $5 input / $25 output per MTok, no explicit cache rates, + // so cache-read defaults to 0.1x and cache-write to 1.25x of input. + r := ModelRate{ + ID: modelOpus48, + Provider: "anthropic", + InputPerMTok: 5, + OutputPerMTok: 25, + } + u := types.TokenUsage{ + InputTokens: 1_000_000, // 1M fresh input -> 5 * 1 = 5.00 + CacheReadTokens: 1_000_000, // 1M cache read -> 0.5 * 1 = 0.50 + CacheCreationTokens: 1_000_000, // 1M cache write -> 6.25 * 1 = 6.25 + OutputTokens: 1_000_000, // 1M output -> 25 * 1 = 25.00 + } + + assert.InDelta(t, 36.75, Estimate(r, u), 1e-9) +} + +func TestEstimate_ExplicitCacheRatesOverrideMultipliers(t *testing.T) { + t.Parallel() + + // Explicit cache rates must be used verbatim instead of the multipliers. + r := ModelRate{ + ID: modelGPT55, + Provider: "openai", + InputPerMTok: 10, + OutputPerMTok: 10, + CacheReadPerMTok: floatPtr(3), + CacheWritePerMTok: floatPtr(2), + } + u := types.TokenUsage{ + InputTokens: 1_000_000, // 10 * 1 = 10 + CacheReadTokens: 1_000_000, // explicit 3 * 1 = 3 (multiplier would give 0.1*10 = 1) + CacheCreationTokens: 1_000_000, // explicit 2 * 1 = 2 (multiplier would give 1.25*10 = 12.5) + OutputTokens: 1_000_000, // 10 * 1 = 10 + } + + // With multipliers this would be 10 + 1 + 12.5 + 10 = 33.5; explicit rates give 25. + assert.InDelta(t, 25.0, Estimate(r, u), 1e-9) +} + +func TestLoadTable_OverrideReplacesByID(t *testing.T) { + t.Parallel() + + // Deliberately minimal: no provider, no aliases, no cache rates — the + // natural shape of a price-bump override. Everything else must inherit. + override := ModelRate{ + ID: modelOpus48, + InputPerMTok: 99, + OutputPerMTok: 199, + } + + tbl, err := LoadTable([]ModelRate{override}) + require.NoError(t, err) + + got, ok := tbl.Lookup(modelOpus48) + require.True(t, ok) + assert.InDelta(t, 99.0, got.InputPerMTok, 1e-9) + assert.InDelta(t, 199.0, got.OutputPerMTok, 1e-9) + + // The override carried no aliases, so it must inherit the embedded entry's + // aliases: a dated spelling still resolves, and to the overridden price. + dated, ok := tbl.Lookup("claude-opus-4-8-20260624") + require.True(t, ok, "dated alias must still resolve after an alias-less override") + assert.Equal(t, modelOpus48, dated.ID) + assert.InDelta(t, 99.0, dated.InputPerMTok, 1e-9) + + // Replacement must not duplicate the id in the table. + base, err := LoadTable(nil) + require.NoError(t, err) + assert.Len(t, tbl.models, len(base.models)) + + // The override omitted provider, so it must inherit the embedded entry's + // provider and keep Anthropic cache economics: 1M cache-read tokens at the + // overridden $99 input rate bill at the 0.1x multiplier, not full input. + assert.Equal(t, "anthropic", got.Provider) + cacheOnly := types.TokenUsage{CacheReadTokens: 1_000_000} + assert.InDelta(t, 9.9, Estimate(got, cacheOnly), 1e-9) +} + +func TestLoadTable_OverrideInheritsCacheRates(t *testing.T) { + t.Parallel() + + // gpt-5.5 carries explicit cache rates in the embedded table. A minimal + // provider-less override must inherit provider AND both cache-rate + // pointers, so cache tokens keep billing at the explicit OpenAI rates + // rather than falling back to full input rate. + base, err := LoadTable(nil) + require.NoError(t, err) + embedded, ok := base.Lookup(modelGPT55) + require.True(t, ok) + require.NotNil(t, embedded.CacheReadPerMTok) + + tbl, err := LoadTable([]ModelRate{{ID: modelGPT55, InputPerMTok: 2, OutputPerMTok: 20}}) + require.NoError(t, err) + got, ok := tbl.Lookup(modelGPT55) + require.True(t, ok) + assert.Equal(t, embedded.Provider, got.Provider) + require.NotNil(t, got.CacheReadPerMTok) + assert.InDelta(t, *embedded.CacheReadPerMTok, *got.CacheReadPerMTok, 1e-9) + require.NotNil(t, got.CacheWritePerMTok) + assert.InDelta(t, *embedded.CacheWritePerMTok, *got.CacheWritePerMTok, 1e-9) +} + +func TestLoadTable_OverrideCanonicalizesID(t *testing.T) { + t.Parallel() + + base, err := LoadTable(nil) + require.NoError(t, err) + + // An override whose id differs from the builtin only in case or surrounding + // whitespace must REPLACE the builtin, not append a phantom duplicate. + for _, variant := range []string{"Claude-Opus-4-8", "claude-opus-4-8 ", " CLAUDE-OPUS-4-8 "} { + t.Run(variant, func(t *testing.T) { + t.Parallel() + + tbl, err := LoadTable([]ModelRate{{ + ID: variant, + Provider: "anthropic", + InputPerMTok: 42, + OutputPerMTok: 84, + }}) + require.NoError(t, err) + + got, ok := tbl.Lookup(modelOpus48) + require.True(t, ok) + assert.InDelta(t, 42.0, got.InputPerMTok, 1e-9, "override must take effect") + + // No phantom entry: table size is unchanged from the embedded defaults. + assert.Len(t, tbl.models, len(base.models)) + }) + } +} + +func TestEstimate_NonAnthropicCacheFallbackBillsAsInput(t *testing.T) { + t.Parallel() + + // A non-anthropic rate with no explicit cache rates: both cache-read and + // cache-write fall back to the full input rate (1.0x), NOT the Anthropic + // 0.1x / 1.25x multipliers. Provider casing is ignored. + r := ModelRate{ + ID: "some-openai-model", + Provider: "OpenAI", + InputPerMTok: 10, + OutputPerMTok: 20, + } + u := types.TokenUsage{ + InputTokens: 1_000_000, // 10 * 1 = 10 + CacheReadTokens: 1_000_000, // 1.0x input = 10 (anthropic would be 1) + CacheCreationTokens: 1_000_000, // 1.0x input = 10 (anthropic would be 12.5) + OutputTokens: 1_000_000, // 20 * 1 = 20 + } + + // 10 + 10 + 10 + 20 = 50 (anthropic multipliers would give 10+1+12.5+20 = 43.5). + assert.InDelta(t, 50.0, Estimate(r, u), 1e-9) +} + +func TestLoadTable_RejectsBadAliasGlob(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + alias string + }{ + {name: "unterminated character class", alias: "claude-[bad"}, + {name: "whitespace-only alias", alias: " "}, + {name: "empty alias", alias: ""}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + _, err := LoadTable([]ModelRate{{ + ID: "custom-model", + Provider: "inhouse", + Aliases: []string{tc.alias}, + InputPerMTok: 1, + OutputPerMTok: 2, + }}) + require.Error(t, err) + }) + } +} + +func TestLoadTable_OverrideAppendsNewID(t *testing.T) { + t.Parallel() + + override := ModelRate{ + ID: "custom-inhouse-model", + Provider: "inhouse", + InputPerMTok: 7, + OutputPerMTok: 8, + } + + base, err := LoadTable(nil) + require.NoError(t, err) + + tbl, err := LoadTable([]ModelRate{override}) + require.NoError(t, err) + + got, ok := tbl.Lookup("custom-inhouse-model") + require.True(t, ok) + assert.InDelta(t, 7.0, got.InputPerMTok, 1e-9) + assert.Len(t, tbl.models, len(base.models)+1) +} + +func TestLoadTable_RejectsInvalidOverride(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + override ModelRate + }{ + {name: "empty id", override: ModelRate{InputPerMTok: 1, OutputPerMTok: 2}}, + {name: "non-positive input", override: ModelRate{ID: "x", InputPerMTok: 0, OutputPerMTok: 2}}, + {name: "non-positive output", override: ModelRate{ID: "x", InputPerMTok: 1, OutputPerMTok: 0}}, + {name: "negative cache read", override: ModelRate{ID: "x", InputPerMTok: 1, OutputPerMTok: 2, CacheReadPerMTok: floatPtr(-1)}}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + _, err := LoadTable([]ModelRate{tc.override}) + require.Error(t, err) + }) + } +} diff --git a/cmd/entire/cli/pricing_refresh.go b/cmd/entire/cli/pricing_refresh.go new file mode 100644 index 0000000000..6188358d1c --- /dev/null +++ b/cmd/entire/cli/pricing_refresh.go @@ -0,0 +1,83 @@ +package cli + +import ( + "context" + "log/slog" + "os" + + "github.com/entireio/cli/cmd/entire/cli/execx" + "github.com/entireio/cli/cmd/entire/cli/logging" + "github.com/entireio/cli/cmd/entire/cli/pricing" + "github.com/entireio/cli/cmd/entire/cli/settings" + "github.com/spf13/cobra" +) + +// refreshPricingCmdName is the hidden worker subcommand the detached spawner +// invokes to run the daily remote-pricing refresh out of the foreground path. +const refreshPricingCmdName = "__refresh_pricing" + +// detachedSpawn is the detached spawner used for the pricing-refresh worker, +// held as a package var so tests can capture the working directory the worker is +// asked to run in without spawning a real process. It re-execs the current +// binary itself (execx.SpawnDetached) and is a no-op under `go test`. +var detachedSpawn = execx.SpawnDetached + +// spawnDetachedRefresh starts the hidden pricing-refresh worker as a detached +// background process rooted at the current working directory — the project the +// hook ran in — so the worker's own IsRemoteEnabled check resolves the project's +// .entire/settings.json. Running at the platform default (analytics' empty dir) +// would hide a project-only pricing.remote opt-in and silently no-op the refresh +// forever. It is a package var so tests can replace it to assert the turn-end / +// post-run triggers only ever *spawn* — they never fetch inline. +var spawnDetachedRefresh = func() { + // Root the worker at the process working directory so its own settings load + // resolves the project's .entire/settings.json. os.Getwd (not paths.RepoRoot) + // is deliberate: this captures the literal cwd to hand a detached child, not a + // git-relative path — the worker re-resolves the repo via git from there. An + // empty dir (Getwd failure) falls back to the platform default in SpawnDetached. + dir, err := os.Getwd() //nolint:forbidigo // hand cwd to a detached child; not a git-relative path + if err != nil { + dir = "" + } + detachedSpawn(dir, refreshPricingCmdName) +} + +// maybeSpawnPricingRefresh spawns the detached remote-pricing refresh worker +// when remote pricing merging is enabled (opt-in) and the local cache is stale +// (>24h) or absent. It never blocks and never performs the network fetch inline +// — the fetch happens in the detached __refresh_pricing worker. Safe to call +// from any command's post-run and from the turn-end hook; the cheap +// ShouldRefresh stat means the common (fresh or disabled) case does no work. +func maybeSpawnPricingRefresh(ctx context.Context) { + if !settings.IsRemoteEnabled(ctx) { + return + } + if !pricing.ShouldRefresh() { + return + } + spawnDetachedRefresh() +} + +// newRefreshPricingCmd builds the hidden worker the detached spawner invokes. It +// runs the actual network refresh, gated on the pricing.remote opt-in. Being +// Hidden, the PersistentPostRun parent-chain guard skips it, so it can never +// itself trigger another refresh spawn (no fork bomb) and never runs telemetry +// or the version check. +func newRefreshPricingCmd() *cobra.Command { + return &cobra.Command{ + Use: refreshPricingCmdName, + Hidden: true, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + ctx := cmd.Context() + if !settings.IsRemoteEnabled(ctx) { + return nil + } + if err := pricing.RefreshRemoteCache(ctx); err != nil { + logging.Debug(ctx, "pricing: remote refresh failed", + slog.String("error", err.Error())) + } + return nil + }, + } +} diff --git a/cmd/entire/cli/pricing_refresh_test.go b/cmd/entire/cli/pricing_refresh_test.go new file mode 100644 index 0000000000..58599dbf10 --- /dev/null +++ b/cmd/entire/cli/pricing_refresh_test.go @@ -0,0 +1,154 @@ +package cli + +import ( + "context" + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/spf13/cobra" +) + +// setupRemoteRepo creates an isolated repo dir with the given settings and an +// empty (throwaway) cache, then chdirs into it. Returns the cache-file path so a +// test can assert whether an inline fetch wrote it. +func setupRemoteRepo(t *testing.T, settingsJSON string) (cacheFile string) { + t.Helper() + cacheHome := t.TempDir() + t.Setenv("XDG_CACHE_HOME", cacheHome) + dir := t.TempDir() + if err := os.MkdirAll(filepath.Join(dir, ".entire"), 0o755); err != nil { + t.Fatalf("mkdir .entire: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, ".entire", "settings.json"), []byte(settingsJSON), 0o644); err != nil { + t.Fatalf("write settings: %v", err) + } + t.Chdir(dir) + return filepath.Join(cacheHome, "entire", "pricing_remote.json") +} + +// stubSpawn replaces the detached spawn hook with a counter for the test. +func stubSpawn(t *testing.T) *int { + t.Helper() + var calls int + orig := spawnDetachedRefresh + spawnDetachedRefresh = func() { calls++ } + t.Cleanup(func() { spawnDetachedRefresh = orig }) + return &calls +} + +// TestSpawnDetachedRefreshUsesProjectDir is the regression for the defect where +// the detached worker ran with cwd "/" (SpawnDetached's old hardcoded dir) and so +// could not see a project-level pricing.remote opt-in — the auto-refresh silently +// no-oped forever. The worker must be spawned rooted at the process working +// directory (the project the hook ran in). +func TestSpawnDetachedRefreshUsesProjectDir(t *testing.T) { + // Not parallel: uses t.Chdir and swaps a package var. + project := t.TempDir() + resolved, err := filepath.EvalSymlinks(project) + if err != nil { + t.Fatalf("evalsymlinks(project): %v", err) + } + t.Chdir(resolved) + + var gotDir string + spawned := false + orig := detachedSpawn + detachedSpawn = func(dir string, _ ...string) { + gotDir = dir + spawned = true + } + t.Cleanup(func() { detachedSpawn = orig }) + + spawnDetachedRefresh() + + if !spawned { + t.Fatal("spawnDetachedRefresh did not spawn the worker") + } + if gotDir == "/" || gotDir == "" { + t.Fatalf("worker dir = %q, want the project dir (a %q/empty dir hides project settings)", gotDir, "/") + } + gotResolved, err := filepath.EvalSymlinks(gotDir) + if err != nil { + t.Fatalf("evalsymlinks(gotDir=%q): %v", gotDir, err) + } + if gotResolved != resolved { + t.Fatalf("worker dir = %q (resolved %q), want project dir %q", gotDir, gotResolved, resolved) + } +} + +func TestNewRefreshPricingCmd_Hidden(t *testing.T) { + t.Parallel() + cmd := newRefreshPricingCmd() + if cmd.Use != refreshPricingCmdName { + t.Errorf("Use = %q, want %q", cmd.Use, refreshPricingCmdName) + } + if !cmd.Hidden { + t.Error("refresh-pricing worker command must be Hidden") + } +} + +func TestRefreshPricingCmd_RegisteredAndHidden(t *testing.T) { + t.Parallel() + root := NewRootCmd() + var found *cobra.Command + for _, c := range root.Commands() { + if c.Use == refreshPricingCmdName { + found = c + break + } + } + if found == nil { + t.Fatalf("%q not registered on root", refreshPricingCmdName) + } + if !found.Hidden { + t.Errorf("%q must be Hidden so PersistentPostRun skips it", refreshPricingCmdName) + } +} + +func TestMaybeSpawnPricingRefresh_SpawnsWhenEnabledAndStale(t *testing.T) { + cacheFile := setupRemoteRepo(t, `{"enabled":true,"pricing":{"remote":true}}`) + calls := stubSpawn(t) + + maybeSpawnPricingRefresh(context.Background()) + + if *calls != 1 { + t.Errorf("spawn calls = %d, want 1 (enabled + no cache = stale)", *calls) + } + // The trigger must never fetch inline: no cache file should have been written. + if _, err := os.Stat(cacheFile); !os.IsNotExist(err) { + t.Errorf("cache file exists (%v); the trigger must only spawn, never fetch inline", err) + } +} + +func TestMaybeSpawnPricingRefresh_SkipsWhenDisabled(t *testing.T) { + setupRemoteRepo(t, `{"enabled":true}`) + calls := stubSpawn(t) + + maybeSpawnPricingRefresh(context.Background()) + + if *calls != 0 { + t.Errorf("spawn calls = %d, want 0 (remote disabled)", *calls) + } +} + +func TestMaybeSpawnPricingRefresh_SkipsWhenFresh(t *testing.T) { + cacheFile := setupRemoteRepo(t, `{"enabled":true,"pricing":{"remote":true}}`) + // Seed a fresh cache so ShouldRefresh is false. + if err := os.MkdirAll(filepath.Dir(cacheFile), 0o755); err != nil { + t.Fatalf("mkdir cache: %v", err) + } + fresh := fmt.Sprintf(`{"fetched_at":%q}`, time.Now().UTC().Format(time.RFC3339)) + if err := os.WriteFile(cacheFile, []byte(fresh), 0o644); err != nil { + t.Fatalf("write cache: %v", err) + } + calls := stubSpawn(t) + + maybeSpawnPricingRefresh(context.Background()) + + if *calls != 0 { + t.Errorf("spawn calls = %d, want 0 (cache is fresh)", *calls) + } +} diff --git a/cmd/entire/cli/root.go b/cmd/entire/cli/root.go index de7af21e4b..a2fc3537e8 100644 --- a/cmd/entire/cli/root.go +++ b/cmd/entire/cli/root.go @@ -91,6 +91,12 @@ func NewRootCmd() *cobra.Command { // stdout is piped into jq or captured by scripts — a notice on stdout // corrupts that output while staying invisible in the caller's logs. versioncheck.CheckAndNotify(cmd.Context(), cmd.ErrOrStderr(), versioninfo.Version) + + // Daily remote pricing refresh (opt-in, detached, non-blocking). The + // Hidden-command guard above means hooks, MCP, and the __refresh_pricing + // worker itself never reach here, so this only fires for real top-level + // commands and can't fork-bomb. + maybeSpawnPricingRefresh(cmd.Context()) }, RunE: func(cmd *cobra.Command, _ []string) error { ctx := cmd.Context() @@ -164,6 +170,7 @@ func NewRootCmd() *cobra.Command { cmd.AddCommand(newHooksCmd()) cmd.AddCommand(newTrailCmd()) cmd.AddCommand(newSendAnalyticsCmd()) + cmd.AddCommand(newRefreshPricingCmd()) // hidden worker for the detached daily pricing refresh cmd.AddCommand(newCurlBashPostInstallCmd()) cmd.AddCommand(newRefreshTrailEnablementCmd()) diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index 1eed68ebfb..0665ec2133 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -298,6 +298,14 @@ type State struct { // This is checkpoint-scoped; TokenUsage remains the session-wide total. CheckpointTokenUsage *agent.TokenUsage `json:"checkpoint_token_usage,omitempty"` + // ModelUsage tracks per-model token usage since the last condensation, keyed by + // model. Checkpoint-scoped like CheckpointTokenUsage (reset together at every + // condensation boundary); it is the per-model breakdown persisted into + // per-checkpoint metadata and used as the fallback when the freshly-extracted + // transcript yields no per-model buckets (e.g. Cursor, whose usage comes only + // from stop-hook payloads). + ModelUsage map[string]*agent.TokenUsage `json:"model_usage,omitempty"` + // SubagentTokensBaseline is a snapshot of TokenUsage.SubagentTokens captured // at the last condensation reset. Subagent token usage is always re-read // from the start of each subagent transcript (agent IDs are discovered from diff --git a/cmd/entire/cli/session_adopt.go b/cmd/entire/cli/session_adopt.go index 92a70bbd92..9f932d4f74 100644 --- a/cmd/entire/cli/session_adopt.go +++ b/cmd/entire/cli/session_adopt.go @@ -480,6 +480,7 @@ func buildAdoptedSessionState(ctx context.Context, source *session.State) (*sess adopted.LastCheckpointID = id.EmptyCheckpointID adopted.LastCheckpointCommitHash = "" adopted.CheckpointTokenUsage = nil + adopted.ModelUsage = nil // Re-baseline the subagent cumulative for the fresh target-local window. The // cloned TokenUsage carries the SOURCE session's full cumulative subagent // total; without re-baselining here, the first post-adopt checkpoint would diff --git a/cmd/entire/cli/session_tokens.go b/cmd/entire/cli/session_tokens.go index 7de7418f22..999fca4ecc 100644 --- a/cmd/entire/cli/session_tokens.go +++ b/cmd/entire/cli/session_tokens.go @@ -9,6 +9,8 @@ import ( "strings" "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/pricing" "github.com/entireio/cli/cmd/entire/cli/strategy" "github.com/spf13/cobra" ) @@ -27,13 +29,15 @@ type sessionTokensReport struct { } type sessionTokensUsage struct { - Total int `json:"total"` - Input int `json:"input"` - CacheRead int `json:"cache_read"` - CacheWrite int `json:"cache_write"` - Output int `json:"output"` - APICalls int `json:"api_calls"` - SubagentTotal int `json:"subagent_total,omitempty"` + Total int `json:"total"` + Input int `json:"input"` + CacheRead int `json:"cache_read"` + CacheWrite int `json:"cache_write"` + Output int `json:"output"` + APICalls int `json:"api_calls"` + SubagentTotal int `json:"subagent_total,omitempty"` + CostUSD *float64 `json:"cost_usd,omitempty"` + CostSource string `json:"cost_source,omitempty"` } type sessionTokensContext struct { @@ -141,7 +145,7 @@ func runSessionTokens(ctx context.Context, cmd *cobra.Command, sessionID string, return NewSilentError(fmt.Errorf("session not found: %s", sessionID)) } - report := buildSessionTokensReport(state, sessionPhaseLabel(state)) + report := buildSessionTokensReport(state, sessionPhaseLabel(state), loadDisplayPricingTable(ctx)) if jsonOutput { return printJSON(cmd.OutOrStdout(), report) } @@ -167,7 +171,7 @@ func tokenCommandError(err error) error { return err } -func buildSessionTokensReport(state *strategy.SessionState, status string) sessionTokensReport { +func buildSessionTokensReport(state *strategy.SessionState, status string, table *pricing.Table) sessionTokensReport { agentLabel := string(state.AgentType) if agentLabel == "" { agentLabel = unknownPlaceholder @@ -183,6 +187,16 @@ func buildSessionTokensReport(state *strategy.SessionState, status string) sessi if tokens := buildSessionTokensUsage(state.TokenUsage); tokens != nil { report.Tokens = tokens + // Estimate cost from the SESSION-WIDE flat usage (state.TokenUsage) — the + // same scope as the displayed token total. state.ModelUsage is deliberately + // NOT passed: it is CHECKPOINT-scoped (reset at every condensation, see + // session.State.ModelUsage), so pricing it while the displayed total is the + // whole session would silently understate cost after any condensation. + // Session state carries no session-cumulative per-model breakdown, so the + // flat total (EstimateCost flattens the subagent subtree into it, matching + // totalTokens) is the only scope-consistent basis; it is priced under the + // session's current model as an explicit local estimate. + applyLocalCostEstimate(tokens, state.TokenUsage, nil, state.ModelName, table) if tokens.SubagentTotal > 0 { report.Contributors = append(report.Contributors, sessionTokensContributor{ Kind: "subagents", @@ -239,6 +253,9 @@ func buildSessionTokensUsage(usage *agent.TokenUsage) *sessionTokensUsage { if total == 0 && usage.APICallCount == 0 { return nil } + // Cost is intentionally NOT copied from usage: the CLI no longer persists + // cost, so any persisted value is nil. Callers apply a local cost estimate + // via applyLocalCostEstimate instead (entire-api owns the authoritative cost). return &sessionTokensUsage{ Total: total, Input: usage.InputTokens, @@ -250,6 +267,31 @@ func buildSessionTokensUsage(usage *agent.TokenUsage) *sessionTokensUsage { } } +// loadDisplayPricingTable returns the pricing table the token commands use to +// compute a LOCAL cost estimate for display, or nil when estimation is disabled +// (so no cost is shown). Never fails: on load error the table is nil. +func loadDisplayPricingTable(ctx context.Context) *pricing.Table { + table, disableEstimation := LoadPricingTable(ctx) + if disableEstimation { + return nil + } + return table +} + +// applyLocalCostEstimate replaces a rendered usage's cost with a LOCAL, on-the-fly +// estimate computed from its token breakdown at current rates. The CLI no longer +// persists cost — entire-api prices server-side — so token commands show this as +// an explicitly-labeled local estimate (see costSourceSuffix / localCostEstimateNote). +// A nil result (nothing priceable) leaves the cost unset so no cost line renders. +func applyLocalCostEstimate(tokens *sessionTokensUsage, usage *agent.TokenUsage, models []types.ModelUsage, fallbackModel string, table *pricing.Table) { + if tokens == nil { + return + } + cost, source := agent.EstimateCost(usage, models, fallbackModel, table) + tokens.CostUSD = cost + tokens.CostSource = source +} + func topLevelSessionTokenTotal(tokens *sessionTokensUsage) int { if tokens == nil { return 0 @@ -409,6 +451,72 @@ func formatPercent(percent float64) string { return formatted + "%" } +// localCostEstimateNote explains that a rendered cost is a CLI-side estimate and +// that the authoritative value lives on the platform. The CLI no longer persists +// cost (entire-api prices server-side), so every cost the token commands show is +// recomputed locally from the persisted token breakdown. +const localCostEstimateNote = "Cost is a local estimate at current rates; the authoritative cost is computed on the Entire platform." + +// formatCostUSD renders a locally-estimated cost for display. The value is a +// current-rate estimate the token command recomputed from the persisted token +// breakdown — the CLI no longer persists cost. +// +// - v == nil -> "" (unknown/unpriceable cost: caller omits the line, never $0.00) +// - 0 < v < $0.005 -> "<$0.01" (sub-cent, would round to $0.00) +// - otherwise -> "$X.XX" (2dp) +// +// A non-empty source appends a provenance suffix: " (estimated locally)", or +// " (estimated locally, partial)" when only some tokens were priceable. An +// empty/unknown source adds no suffix. +func formatCostUSD(v *float64, source string) string { + if v == nil { + return "" + } + out := formatCostAmount(*v) + if suffix := costSourceSuffix(source); suffix != "" { + out += " " + suffix + } + return out +} + +// formatCostAmount formats a bare dollar amount ("$X.XX"), collapsing a positive +// sub-cent value to "<$0.01" so it never displays as "$0.00". +func formatCostAmount(v float64) string { + if v > 0 && v < 0.005 { + return "<$0.01" + } + return fmt.Sprintf("$%.2f", v) +} + +// costSourceSuffix maps a CostSource provenance to its parenthetical display +// suffix. Cost is a local estimate now (the CLI no longer persists cost), so the +// estimated/mixed cases say so explicitly; an empty/unknown source yields "". +func costSourceSuffix(source string) string { + switch source { + case types.CostSourceReported: + return "(reported)" + case types.CostSourceEstimated: + return "(estimated locally)" + case types.CostSourceMixed: + return "(estimated locally, partial)" + default: + return "" + } +} + +// writeTokenCostLine emits the "Cost:" line inside a token usage block when a +// cost is known. Shared by the session and checkpoint text writers. A nil usage +// or nil cost prints nothing (unknown cost is never rendered as $0.00). +func writeTokenCostLine(w io.Writer, tokens *sessionTokensUsage) { + if tokens == nil { + return + } + if line := formatCostUSD(tokens.CostUSD, tokens.CostSource); line != "" { + fmt.Fprintf(w, "Cost: %s\n", line) + fmt.Fprintf(w, " %s\n", localCostEstimateNote) + } +} + func skillEventLabels(events []agent.SkillEvent) []string { seen := make(map[string]struct{}, len(events)) labels := make([]string, 0, len(events)) @@ -443,6 +551,7 @@ func writeSessionTokensText(w io.Writer, report sessionTokensReport) { fmt.Fprintf(w, "Status: %s\n", report.Status) writeTokenUsageSection(w, report.Tokens) + writeTokenCostLine(w, report.Tokens) if len(report.Recommendations) > 0 { writeTokenRecommendations(w, report.Recommendations) } @@ -474,15 +583,21 @@ func agentBriefUsageLine(tokens *sessionTokensUsage) string { if tokens == nil { return "Token usage: unavailable." } + var line string if tokens.CacheRead > 0 { - return fmt.Sprintf( + line = fmt.Sprintf( "Token usage: %s total; %s cache/context replay; %s.", formatTokenCount(tokens.Total), formatPercent(tokenPercent(tokens.CacheRead, topLevelSessionTokenTotal(tokens))), formatAPICalls(tokens.APICalls), ) + } else { + line = fmt.Sprintf("Token usage: %s total; %s.", formatTokenCount(tokens.Total), formatAPICalls(tokens.APICalls)) + } + if cost := formatCostUSD(tokens.CostUSD, tokens.CostSource); cost != "" { + line += " Cost: " + cost + "." } - return fmt.Sprintf("Token usage: %s total; %s.", formatTokenCount(tokens.Total), formatAPICalls(tokens.APICalls)) + return line } func formatAPICalls(count int) string { diff --git a/cmd/entire/cli/session_tokens_cost_test.go b/cmd/entire/cli/session_tokens_cost_test.go new file mode 100644 index 0000000000..1ac9d66008 --- /dev/null +++ b/cmd/entire/cli/session_tokens_cost_test.go @@ -0,0 +1,196 @@ +package cli + +import ( + "bytes" + "encoding/json" + "strings" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/session" +) + +// costPtr and testDisplayPricingTable are defined in checkpoint_tokens_cost_test.go (same package). + +func TestFormatCostUSD_Matrix(t *testing.T) { + t.Parallel() + tests := []struct { + name string + v *float64 + source string + want string + }{ + {"nil is empty", nil, types.CostSourceEstimated, ""}, + {"known zero", costPtr(0), "", "$0.00"}, + {"two dp estimated is local", costPtr(0.42), types.CostSourceEstimated, "$0.42 (estimated locally)"}, + {"trailing zero padded reported", costPtr(1.5), types.CostSourceReported, "$1.50 (reported)"}, + {"whole dollars mixed is partial local", costPtr(3), types.CostSourceMixed, "$3.00 (estimated locally, partial)"}, + {"sub-cent estimated is local", costPtr(0.004), types.CostSourceEstimated, "<$0.01 (estimated locally)"}, + {"sub-cent no source", costPtr(0.001), "", "<$0.01"}, + {"cent boundary rounds to a cent", costPtr(0.005), types.CostSourceReported, "$0.01 (reported)"}, + {"unknown source has no suffix", costPtr(0.42), "weird", "$0.42"}, + {"empty source has no suffix", costPtr(0.42), "", "$0.42"}, + {"large amount estimated is local", costPtr(1234.5), types.CostSourceEstimated, "$1234.50 (estimated locally)"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := formatCostUSD(tc.v, tc.source); got != tc.want { + t.Errorf("formatCostUSD(%v, %q) = %q, want %q", tc.v, tc.source, got, tc.want) + } + }) + } +} + +// buildSessionTokensReport must recompute a LOCAL cost estimate from the +// session's token breakdown (the CLI no longer persists cost) and label it as a +// local estimate. The per-model breakdown carries 420000 priceable tokens under +// test-model ($1/MTok) => $0.42. +func TestBuildSessionTokensReport_RecomputesLocalCostEstimate(t *testing.T) { + t.Parallel() + + state := makeSessionState("cost-session", session.PhaseActive) + state.AgentType = testAgentClaude + state.ModelName = "test-model" + state.TokenUsage = &agent.TokenUsage{ + InputTokens: 400000, + OutputTokens: 20000, + APICallCount: 3, + } + state.ModelUsage = map[string]*agent.TokenUsage{ + "test-model": {InputTokens: 400000, OutputTokens: 20000}, + } + + report := buildSessionTokensReport(state, "active", testDisplayPricingTable(t)) + if report.Tokens == nil || report.Tokens.CostUSD == nil || *report.Tokens.CostUSD != 0.42 { + t.Fatalf("expected recomputed cost 0.42, got %+v", report.Tokens) + } + if report.Tokens.CostSource != types.CostSourceEstimated { + t.Fatalf("cost source = %q, want estimated", report.Tokens.CostSource) + } + + var buf bytes.Buffer + writeSessionTokensText(&buf, report) + out := buf.String() + if !strings.Contains(out, "Cost: $0.42 (estimated locally)") { + t.Fatalf("expected local-estimate cost line, got:\n%s", out) + } + if !strings.Contains(out, localCostEstimateNote) { + t.Fatalf("expected local-estimate note, got:\n%s", out) + } + + data, err := json.Marshal(report) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if !strings.Contains(string(data), `"cost_usd":0.42`) { + t.Fatalf("expected cost_usd in JSON, got: %s", data) + } + if !strings.Contains(string(data), `"cost_source":"estimated"`) { + t.Fatalf("expected cost_source in JSON, got: %s", data) + } +} + +// When there is no pricing table (estimation disabled) or no priceable model, no +// cost is shown — never $0. +func TestBuildSessionTokensReport_OmitsCostWhenUnpriceable(t *testing.T) { + t.Parallel() + + state := makeSessionState("no-cost-session", session.PhaseActive) + state.AgentType = testAgentClaude + state.TokenUsage = &agent.TokenUsage{ + InputTokens: 1000, + OutputTokens: 100, + APICallCount: 3, + } + + // nil table => estimation disabled => no cost. + report := buildSessionTokensReport(state, "active", nil) + if report.Tokens == nil { + t.Fatal("expected token usage") + } + if report.Tokens.CostUSD != nil { + t.Fatalf("expected nil cost, got %v", report.Tokens.CostUSD) + } + + var buf bytes.Buffer + writeSessionTokensText(&buf, report) + if strings.Contains(buf.String(), "Cost:") { + t.Fatalf("expected no cost line, got:\n%s", buf.String()) + } + + data, err := json.Marshal(report) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if strings.Contains(string(data), "cost_usd") { + t.Fatalf("expected cost_usd omitted, got: %s", data) + } +} + +func TestAgentBriefUsageLine_AppendsLocalEstimateWhenKnown(t *testing.T) { + t.Parallel() + + withCost := agentBriefUsageLine(&sessionTokensUsage{ + Total: 1000, + APICalls: 3, + CostUSD: costPtr(1.5), + CostSource: types.CostSourceEstimated, + }) + if !strings.Contains(withCost, "Cost: $1.50 (estimated locally).") { + t.Fatalf("expected local-estimate cost clause, got: %q", withCost) + } + + withoutCost := agentBriefUsageLine(&sessionTokensUsage{Total: 1000, APICalls: 3}) + if strings.Contains(withoutCost, "Cost:") { + t.Fatalf("expected no cost clause, got: %q", withoutCost) + } +} + +// After a condensation, state.ModelUsage is reset to only the most recent +// checkpoint's per-model buckets, while state.TokenUsage keeps the session-wide +// cumulative total that drives the displayed token counts. The local cost +// estimate must be priced from the SAME (session-wide) scope as the displayed +// token total — not the checkpoint-scoped per-model map — or the displayed cost +// silently understates the displayed token total. Mutation check: pricing the +// checkpoint-scoped ModelUsage (the pre-fix behavior) yields $0.42 for a session +// whose displayed total is 840000 tokens; pricing the session-wide flat total +// yields the scope-consistent $0.84. +func TestBuildSessionTokensReport_CostMatchesSessionScopeAfterCondensation(t *testing.T) { + t.Parallel() + + state := makeSessionState("post-condensation-session", session.PhaseActive) + state.AgentType = testAgentClaude + state.ModelName = "test-model" + // Session-wide cumulative total spans the whole session (840000 priceable tokens). + state.TokenUsage = &agent.TokenUsage{ + InputTokens: 800000, + OutputTokens: 40000, + APICallCount: 6, + } + // ModelUsage is checkpoint-scoped: reset at the last condensation, it covers + // only the most recent checkpoint (420000 tokens) — half the session. + state.ModelUsage = map[string]*agent.TokenUsage{ + "test-model": {InputTokens: 400000, OutputTokens: 20000}, + } + + report := buildSessionTokensReport(state, "active", testDisplayPricingTable(t)) + if report.Tokens == nil { + t.Fatal("expected token usage") + } + // Displayed token total is session-wide. + if report.Tokens.Total != 840000 { + t.Fatalf("displayed total = %d, want 840000 (session-wide)", report.Tokens.Total) + } + // Cost must be priced from the same session-wide scope: 840000 * $1/MTok = $0.84. + if report.Tokens.CostUSD == nil { + t.Fatal("cost = nil, want 0.84 (session-wide scope)") + } + if *report.Tokens.CostUSD != 0.84 { + t.Fatalf("cost = %v, want 0.84 (session-wide scope); pricing the checkpoint-scoped ModelUsage understates to 0.42", *report.Tokens.CostUSD) + } + if report.Tokens.CostSource != types.CostSourceEstimated { + t.Fatalf("cost source = %q, want estimated", report.Tokens.CostSource) + } +} diff --git a/cmd/entire/cli/sessions_test.go b/cmd/entire/cli/sessions_test.go index e4f920834e..2cee2807f3 100644 --- a/cmd/entire/cli/sessions_test.go +++ b/cmd/entire/cli/sessions_test.go @@ -2253,6 +2253,7 @@ func TestCheckpointTokensReport_UsesRootSummaryWhenSessionMetadataIncomplete(t * }, }, 1, + nil, ) if report.Tokens == nil { @@ -2336,6 +2337,7 @@ func TestCheckpointTokensReport_UsesRootSummaryWhenNoSessionMetadataReadable(t * }, nil, 2, + nil, ) if report.Tokens == nil { diff --git a/cmd/entire/cli/settings/pricing_model_test.go b/cmd/entire/cli/settings/pricing_model_test.go new file mode 100644 index 0000000000..0b65846a52 --- /dev/null +++ b/cmd/entire/cli/settings/pricing_model_test.go @@ -0,0 +1,101 @@ +package settings + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent/types" +) + +const ( + testAgentCodex types.AgentType = "Codex" // mirrors agent.AgentTypeCodex + testAgentClaude types.AgentType = "Claude Code" // mirrors agent.AgentTypeClaudeCode + + testModelGPT55 = "gpt-5.5" + testModelGPT55Priority = "gpt-5.5-priority" +) + +// TestPricingModelForAgent covers the tier-suffix seam directly on the settings +// method: only a Codex session with the "priority" knob gets the "-priority" +// pricing variant, and never a double suffix. +func TestPricingModelForAgent(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + agentType types.AgentType + tier string + model string + want string + }{ + {"codex priority appends variant", testAgentCodex, "priority", testModelGPT55, testModelGPT55Priority}, + {"codex priority case-insensitive", testAgentCodex, "Priority", testModelGPT55, testModelGPT55Priority}, + {"codex empty tier unchanged", testAgentCodex, "", testModelGPT55, testModelGPT55}, + {"codex explicit standard unchanged", testAgentCodex, "standard", testModelGPT55, testModelGPT55}, + {"non-codex priority ignored", testAgentClaude, "priority", "claude-opus-4-8", "claude-opus-4-8"}, + {"empty model unchanged", testAgentCodex, "priority", "", ""}, + {"already-priority not doubled", testAgentCodex, "priority", testModelGPT55Priority, testModelGPT55Priority}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + s := &EntireSettings{Pricing: &PricingSettings{CodexServiceTier: tc.tier}} + if got := s.PricingModelForAgent(tc.agentType, tc.model); got != tc.want { + t.Errorf("PricingModelForAgent(%q, %q) tier=%q = %q, want %q", tc.agentType, tc.model, tc.tier, got, tc.want) + } + }) + } +} + +// TestPricingModelForAgent_NilSafe confirms the method never panics on a nil +// receiver or unset pricing settings (both read as the default empty tier). +func TestPricingModelForAgent_NilSafe(t *testing.T) { + t.Parallel() + var nilSettings *EntireSettings + if got := nilSettings.PricingModelForAgent(testAgentCodex, testModelGPT55); got != testModelGPT55 { + t.Errorf("nil settings = %q, want %q (no tier)", got, testModelGPT55) + } + if got := (&EntireSettings{}).PricingModelForAgent(testAgentCodex, testModelGPT55); got != testModelGPT55 { + t.Errorf("empty settings = %q, want %q (no tier)", got, testModelGPT55) + } +} + +// TestPricingModelForAgentCtx exercises the context-based loader: it reads the +// project settings for the tier, short-circuits non-Codex agents without a load, +// and returns the model unchanged when settings cannot be loaded. +func TestPricingModelForAgentCtx(t *testing.T) { + t.Parallel() + + writeTier := func(t *testing.T, tier string) context.Context { + t.Helper() + dir := t.TempDir() + entireDir := filepath.Join(dir, ".entire") + if err := os.MkdirAll(entireDir, 0o755); err != nil { + t.Fatalf("mkdir .entire: %v", err) + } + body := `{"enabled": true, "pricing": {"codex_service_tier": "` + tier + `"}}` + if err := os.WriteFile(filepath.Join(entireDir, "settings.json"), []byte(body), 0o644); err != nil { + t.Fatalf("write settings.json: %v", err) + } + return WithWorktreeRoot(context.Background(), dir) + } + + // Codex + project priority knob -> suffixed. + if got := PricingModelForAgent(writeTier(t, "priority"), testAgentCodex, testModelGPT55); got != testModelGPT55Priority { + t.Errorf("codex priority via ctx = %q, want %q", got, testModelGPT55Priority) + } + // Codex + no knob -> unchanged. + if got := PricingModelForAgent(writeTier(t, ""), testAgentCodex, testModelGPT55); got != testModelGPT55 { + t.Errorf("codex standard via ctx = %q, want %q", got, testModelGPT55) + } + // Non-Codex short-circuits before any load, even with a priority file present. + if got := PricingModelForAgent(writeTier(t, "priority"), testAgentClaude, "claude-opus-4-8"); got != "claude-opus-4-8" { + t.Errorf("non-codex via ctx = %q, want claude-opus-4-8", got) + } + // Empty model is returned as-is regardless of settings. + if got := PricingModelForAgent(writeTier(t, "priority"), testAgentCodex, ""); got != "" { + t.Errorf("empty model via ctx = %q, want empty", got) + } +} diff --git a/cmd/entire/cli/settings/settings.go b/cmd/entire/cli/settings/settings.go index ee36271fa2..4d06da1028 100644 --- a/cmd/entire/cli/settings/settings.go +++ b/cmd/entire/cli/settings/settings.go @@ -18,10 +18,12 @@ import ( "strings" "time" + "github.com/entireio/cli/cmd/entire/cli/agent/types" "github.com/entireio/cli/cmd/entire/cli/internal/flock" "github.com/entireio/cli/cmd/entire/cli/jsonutil" "github.com/entireio/cli/cmd/entire/cli/logging" "github.com/entireio/cli/cmd/entire/cli/paths" + "github.com/entireio/cli/cmd/entire/cli/pricing" "github.com/entireio/cli/cmd/entire/cli/session" "github.com/entireio/cli/redact" ) @@ -150,6 +152,10 @@ type EntireSettings struct { // settings loader (DisallowUnknownFields) accepts a "checkpoints" key. Checkpoints *CheckpointsConfig `json:"checkpoints,omitempty"` + // Pricing configures token-cost estimation: per-model rate overrides layered + // on top of the embedded defaults, plus a switch to disable estimation. + Pricing *PricingSettings `json:"pricing,omitempty"` + // Deprecated: no longer used. Exists to tolerate old settings files // that still contain "strategy": "auto-commit" or similar. Strategy string `json:"strategy,omitempty"` @@ -266,6 +272,167 @@ type RedactionSettings struct { ExternalizeImages bool `json:"externalize_images,omitempty"` } +// PricingSettings configures token-cost estimation. Models are per-model rate +// overrides layered on top of the embedded pricing defaults (an override whose +// id matches a built-in entry replaces it; otherwise it is appended). +// DisableEstimation, when true, turns off derived cost estimates entirely — +// agent-reported costs still flow through, but unpriced usage keeps a nil cost +// rather than an estimate. +type PricingSettings struct { + Models []pricing.ModelRate `json:"models,omitempty"` + DisableEstimation *bool `json:"disable_estimation,omitempty"` + // CodexServiceTier opts a Codex session into a non-default pricing tier. + // "priority" prices Codex turns at the priority-tier premium via the + // model's "-priority" variant; empty (the default) prices at standard rates. + CodexServiceTier string `json:"codex_service_tier,omitempty"` + // Remote, when true, merges a locally cached remote pricing table (refreshed + // in the background, never fetched inline) on top of the embedded defaults + // and beneath any pricing.models overrides. Opt-in via a pointer so the + // default (nil/false) keeps estimation on embedded defaults + local + // overrides only, with no network activity. + Remote *bool `json:"remote,omitempty"` +} + +// PricingOverrides returns the configured per-model rate overrides, or nil when +// none are set. +func (s *EntireSettings) PricingOverrides() []pricing.ModelRate { + if s == nil || s.Pricing == nil { + return nil + } + return s.Pricing.Models +} + +// IsEstimationDisabled reports whether cost estimation is turned off. Defaults +// to false (estimation enabled) when unset. +func (s *EntireSettings) IsEstimationDisabled() bool { + if s == nil || s.Pricing == nil || s.Pricing.DisableEstimation == nil { + return false + } + return *s.Pricing.DisableEstimation +} + +// CodexServiceTier returns the configured Codex pricing service tier (e.g. +// "priority"), or "" when unset. It is nil-safe. Codex priority-tier turns bill +// at a premium that is priced via a "-priority" model variant; the default +// empty tier prices at standard rates. +func (s *EntireSettings) CodexServiceTier() string { + if s == nil || s.Pricing == nil { + return "" + } + return s.Pricing.CodexServiceTier +} + +// codexAgentType mirrors agent.AgentTypeCodex. The settings package cannot +// import the agent package (the agent plugin tree transitively imports settings), +// so the Codex display-name value is duplicated here. It must stay in sync with +// agent.AgentTypeCodex — the value is stored in metadata/trailers and is stable. +const codexAgentType types.AgentType = "Codex" + +// PricingModelForAgent maps a turn's model to the id it should be priced under, +// applying this session's opt-in service tier. A Codex session whose +// pricing.codex_service_tier is "priority" bills at a premium the pricing table +// carries as a "-priority" model variant; every other agent or tier prices under +// the model unchanged. An empty model, or a model already ending in "-priority", +// is returned as-is so the id never double-suffixes. +// +// The returned id is a pricing-lookup concern only: callers keep the STORED model +// name (state.ModelName / checkpoint metadata) raw and apply this solely at the +// point of pricing so both turn-end and condensation price the same tier. +func (s *EntireSettings) PricingModelForAgent(agentType types.AgentType, model string) string { + if model == "" { + return model + } + if agentType == codexAgentType && strings.EqualFold(s.CodexServiceTier(), "priority") && !strings.HasSuffix(model, "-priority") { + return model + "-priority" + } + return model +} + +// PricingModelForAgent loads settings and applies the service-tier suffix via the +// EntireSettings method of the same name. Non-Codex agents short-circuit before +// any settings load, preserving the "only Codex reads settings" property on the +// pricing path. A settings-load failure returns the model unchanged, matching the +// estimation defaults so a load error can never silently switch pricing tiers. +func PricingModelForAgent(ctx context.Context, agentType types.AgentType, model string) string { + if model == "" || agentType != codexAgentType { + return model + } + s, err := Load(ctx) + if err != nil { + return model + } + return s.PricingModelForAgent(agentType, model) +} + +// IsRemoteEnabled reports whether remote pricing merging is enabled. It is +// nil-safe and defaults to false (embedded defaults + local overrides only, no +// network activity) when the setting is unset. +func (s *EntireSettings) IsRemoteEnabled() bool { + if s == nil || s.Pricing == nil || s.Pricing.Remote == nil { + return false + } + return *s.Pricing.Remote +} + +// IsRemoteEnabled loads settings and reports whether remote pricing merging is +// enabled. A settings-load failure returns false, matching the opt-in default so +// a load error can never switch the background network refresh on. +func IsRemoteEnabled(ctx context.Context) bool { + s, err := Load(ctx) + if err != nil { + return false + } + return s.IsRemoteEnabled() +} + +// LoadPricingTable builds the model pricing table from settings — the embedded +// defaults with any configured pricing.models overrides layered on top — and +// reports whether cost estimation is disabled. It never fails the caller: a +// settings-load or table-build error is logged and yields a nil table, which +// callers treat as "cost unavailable" (they must never fail a hook path over +// missing pricing). +func LoadPricingTable(ctx context.Context) (*pricing.Table, bool) { + s, err := Load(ctx) + if err != nil { + logging.Debug(ctx, "pricing: settings load failed; cost estimation unavailable", + slog.String("error", err.Error())) + return nil, false + } + disable := s.IsEstimationDisabled() + + // Remote entries (opt-in) sit beneath the user's pricing.models: append the + // cached remote table first, then the user overrides, so LoadTable's + // last-writer-wins keeps a user override authoritative over the same id from + // the remote table, which in turn overrides the embedded default. + // + // Validate the user overrides per-entry and drop only the invalid ones + // (mirroring LoadRemoteEntries): a single malformed pricing.models entry must + // not make LoadTable hard-error and disable cost estimation for the whole + // table. + overrides := pricing.ValidOverrides(ctx, s.PricingOverrides()) + if s.IsRemoteEnabled() { + remote := pricing.LoadRemoteEntries(ctx) + fetchedAt := pricing.RemoteFetchedAt() + var staleness time.Duration + if !fetchedAt.IsZero() { + staleness = time.Since(fetchedAt) + } + logging.Debug(ctx, "pricing: merging remote pricing entries", + slog.Int("remote_entries", len(remote)), + slog.Time("fetched_at", fetchedAt), + slog.Duration("staleness", staleness)) + overrides = append(remote, overrides...) + } + + table, err := pricing.LoadTable(overrides) + if err != nil { + logging.Warn(ctx, "pricing: table load failed; cost estimation unavailable", + slog.String("error", err.Error())) + return nil, disable + } + return table, disable +} + // PIISettings configures PII detection categories. // When Enabled is true, email and phone default to true; address defaults to false. type PIISettings struct { @@ -941,6 +1108,55 @@ func mergeJSON(settings *EntireSettings, data []byte) error { return err } + // Merge pricing sub-fields if present (field-level, not wholesale replace). + if pricingRaw, ok := raw["pricing"]; ok { + if settings.Pricing == nil { + settings.Pricing = &PricingSettings{} + } + if err := mergePricing(settings.Pricing, pricingRaw); err != nil { + return fmt.Errorf("parsing pricing field: %w", err) + } + } + + return nil +} + +// mergePricing merges pricing overrides into existing PricingSettings. Only +// fields present in the override JSON are applied; models are replaced wholesale +// so a local override file fully controls the override set. +func mergePricing(dst *PricingSettings, data json.RawMessage) error { + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return fmt.Errorf("parsing pricing: %w", err) + } + if modelsRaw, ok := raw["models"]; ok { + var models []pricing.ModelRate + if err := unmarshalField("pricing.models", modelsRaw, &models); err != nil { + return err + } + dst.Models = models + } + if disableRaw, ok := raw["disable_estimation"]; ok { + var b bool + if err := unmarshalField("pricing.disable_estimation", disableRaw, &b); err != nil { + return err + } + dst.DisableEstimation = &b + } + if tierRaw, ok := raw["codex_service_tier"]; ok { + var tier string + if err := unmarshalField("pricing.codex_service_tier", tierRaw, &tier); err != nil { + return err + } + dst.CodexServiceTier = tier + } + if remoteRaw, ok := raw["remote"]; ok { + var b bool + if err := unmarshalField("pricing.remote", remoteRaw, &b); err != nil { + return err + } + dst.Remote = &b + } return nil } diff --git a/cmd/entire/cli/settings/settings_pricing_remote_test.go b/cmd/entire/cli/settings/settings_pricing_remote_test.go new file mode 100644 index 0000000000..5619bd3d2c --- /dev/null +++ b/cmd/entire/cli/settings/settings_pricing_remote_test.go @@ -0,0 +1,144 @@ +package settings + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/internal/entireclient/userdirs" +) + +// writeRemotePricingCache writes a raw remote pricing cache document to the +// isolated cache dir. Call t.Setenv("XDG_CACHE_HOME", ...) first so it lands in +// a throwaway location, never the developer's real ~/.cache/entire. +func writeRemotePricingCache(t *testing.T, jsonBody string) { + t.Helper() + dir := userdirs.Cache() + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("mkdir cache dir: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "pricing_remote.json"), []byte(jsonBody), 0o644); err != nil { + t.Fatalf("write remote pricing cache: %v", err) + } +} + +// remoteGPT55At999 is a well-formed remote cache doc that reprices gpt-5.5 far +// away from its embedded rate (5) so tests can tell which layer won. +const remoteGPT55At999 = `{"fetched_at":"2026-07-10T00:00:00Z","doc":{"schema_version":1,` + + `"models":[{"id":"gpt-5.5","provider":"openai","input_per_mtok":999,"output_per_mtok":1000}]}}` + +func lookupInput(t *testing.T, settingsJSON, cacheJSON string) float64 { + t.Helper() + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + setupSettingsDir(t, settingsJSON, "") + if cacheJSON != "" { + writeRemotePricingCache(t, cacheJSON) + } + table, _ := LoadPricingTable(context.Background()) + if table == nil { + t.Fatal("LoadPricingTable returned nil table") + } + rate, ok := table.Lookup("gpt-5.5") + if !ok { + t.Fatal("gpt-5.5 did not resolve") + } + return rate.InputPerMTok +} + +func TestLoadPricingTable_RemoteDisabled_EmbeddedOnly(t *testing.T) { + // Remote setting absent: the cached doc must be ignored entirely. + got := lookupInput(t, `{"enabled":true}`, remoteGPT55At999) + if got != 5 { + t.Errorf("gpt-5.5 input = %v, want 5 (embedded); remote must not merge when disabled", got) + } +} + +func TestLoadPricingTable_RemoteEnabled_RemoteOverridesEmbedded(t *testing.T) { + got := lookupInput(t, `{"enabled":true,"pricing":{"remote":true}}`, remoteGPT55At999) + if got != 999 { + t.Errorf("gpt-5.5 input = %v, want 999 (remote overrides embedded)", got) + } +} + +func TestLoadPricingTable_RemoteEnabled_UserOverrideWins(t *testing.T) { + settingsJSON := `{"enabled":true,"pricing":{"remote":true,` + + `"models":[{"id":"gpt-5.5","provider":"openai","input_per_mtok":111,"output_per_mtok":222}]}}` + got := lookupInput(t, settingsJSON, remoteGPT55At999) + if got != 111 { + t.Errorf("gpt-5.5 input = %v, want 111 (user override wins over remote)", got) + } +} + +func TestLoadPricingTable_RemoteEnabled_CorruptCache_EmbeddedOnly(t *testing.T) { + got := lookupInput(t, `{"enabled":true,"pricing":{"remote":true}}`, `{not valid json`) + if got != 5 { + t.Errorf("gpt-5.5 input = %v, want 5 (corrupt cache falls back to embedded)", got) + } +} + +func TestLoadPricingTable_RemoteEnabled_SchemaVersion99Ignored(t *testing.T) { + cache := `{"fetched_at":"2026-07-10T00:00:00Z","doc":{"schema_version":99,` + + `"models":[{"id":"gpt-5.5","provider":"openai","input_per_mtok":999,"output_per_mtok":1000}]}}` + got := lookupInput(t, `{"enabled":true,"pricing":{"remote":true}}`, cache) + if got != 5 { + t.Errorf("gpt-5.5 input = %v, want 5 (unsupported schema_version ignored)", got) + } +} + +// A single malformed pricing.models override must not sink the whole table. The +// invalid entry (input_per_mtok 0) is dropped and logged; the valid one still +// applies and cost estimation stays enabled. Mutation check: passing the raw +// overrides straight to LoadTable (the pre-fix behavior) hard-errors on the bad +// entry, so LoadPricingTable returns a nil table and disables estimation. +func TestLoadPricingTable_DropsInvalidUserOverride_KeepsValid(t *testing.T) { + // Not parallel: setupSettingsDir chdirs and LoadPricingTable reads cwd settings. + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + settingsJSON := `{"enabled":true,"pricing":{"models":[` + + `{"id":"gpt-5.5","provider":"openai","input_per_mtok":111,"output_per_mtok":222},` + + `{"id":"bad-model","provider":"openai","input_per_mtok":0,"output_per_mtok":10}]}}` + setupSettingsDir(t, settingsJSON, "") + + table, disable := LoadPricingTable(context.Background()) + if table == nil { + t.Fatal("expected a usable table; one bad override must not disable estimation") + } + if disable { + t.Fatal("estimation must stay enabled when only one override entry is invalid") + } + rate, ok := table.Lookup("gpt-5.5") + if !ok { + t.Fatal("gpt-5.5 did not resolve") + } + if rate.InputPerMTok != 111 { + t.Errorf("gpt-5.5 input = %v, want 111 (valid override applied despite a sibling bad entry)", rate.InputPerMTok) + } + // The invalid entry must have been dropped, not added to the table. + if _, ok := table.Lookup("bad-model"); ok { + t.Error("bad-model must have been dropped, not added to the table") + } +} + +func TestIsRemoteEnabled_Accessor(t *testing.T) { + t.Parallel() + + var nilSettings *EntireSettings + if nilSettings.IsRemoteEnabled() { + t.Error("nil settings IsRemoteEnabled() = true, want false") + } + if (&EntireSettings{}).IsRemoteEnabled() { + t.Error("empty settings IsRemoteEnabled() = true, want false") + } + if (&EntireSettings{Pricing: &PricingSettings{}}).IsRemoteEnabled() { + t.Error("pricing without remote IsRemoteEnabled() = true, want false") + } + + enabled := true + if !(&EntireSettings{Pricing: &PricingSettings{Remote: &enabled}}).IsRemoteEnabled() { + t.Error("remote=true IsRemoteEnabled() = false, want true") + } + disabled := false + if (&EntireSettings{Pricing: &PricingSettings{Remote: &disabled}}).IsRemoteEnabled() { + t.Error("remote=false IsRemoteEnabled() = true, want false") + } +} diff --git a/cmd/entire/cli/settings/settings_pricing_test.go b/cmd/entire/cli/settings/settings_pricing_test.go new file mode 100644 index 0000000000..592979770d --- /dev/null +++ b/cmd/entire/cli/settings/settings_pricing_test.go @@ -0,0 +1,132 @@ +package settings + +import ( + "context" + "encoding/json" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/pricing" +) + +func TestPricingAccessors_Defaults(t *testing.T) { + t.Parallel() + + // nil receiver and empty settings both report the safe defaults. + var nilSettings *EntireSettings + if nilSettings.PricingOverrides() != nil { + t.Error("nil settings PricingOverrides() != nil") + } + if nilSettings.IsEstimationDisabled() { + t.Error("nil settings IsEstimationDisabled() = true, want false") + } + + empty := &EntireSettings{} + if empty.PricingOverrides() != nil { + t.Error("empty settings PricingOverrides() != nil") + } + if empty.IsEstimationDisabled() { + t.Error("empty settings IsEstimationDisabled() = true, want false") + } +} + +func TestPricingSettings_LoadFromBytes(t *testing.T) { + t.Parallel() + + data := []byte(`{ + "enabled": true, + "pricing": { + "models": [ + {"id": "custom-x", "provider": "test", "input_per_mtok": 3, "output_per_mtok": 9} + ], + "disable_estimation": true + } + }`) + + s, err := LoadFromBytes(data) + if err != nil { + t.Fatalf("LoadFromBytes: %v", err) + } + overrides := s.PricingOverrides() + if len(overrides) != 1 { + t.Fatalf("PricingOverrides() len = %d, want 1", len(overrides)) + } + if overrides[0].ID != "custom-x" || overrides[0].InputPerMTok != 3 || overrides[0].OutputPerMTok != 9 { + t.Errorf("override = %+v, want custom-x in=3 out=9", overrides[0]) + } + if !s.IsEstimationDisabled() { + t.Error("IsEstimationDisabled() = false, want true") + } +} + +func TestPricingSettings_JSONRoundTrip(t *testing.T) { + t.Parallel() + + disable := true + orig := &EntireSettings{ + Enabled: true, + Pricing: &PricingSettings{ + DisableEstimation: &disable, + Models: []pricing.ModelRate{ + {ID: "rt-model", Provider: "test", InputPerMTok: 4, OutputPerMTok: 8}, + }, + }, + } + + raw, err := json.Marshal(orig) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + got, err := LoadFromBytes(raw) + if err != nil { + t.Fatalf("LoadFromBytes(round-trip): %v", err) + } + if !got.IsEstimationDisabled() { + t.Error("round-trip lost disable_estimation") + } + if len(got.PricingOverrides()) != 1 || got.PricingOverrides()[0].ID != "rt-model" { + t.Errorf("round-trip overrides = %+v, want single rt-model", got.PricingOverrides()) + } +} + +func TestPricingSettings_LocalOverrideMerges(t *testing.T) { + base := `{ + "enabled": true, + "pricing": { + "models": [{"id": "base-model", "provider": "test", "input_per_mtok": 1, "output_per_mtok": 1}], + "disable_estimation": false + } + }` + local := `{"pricing": {"disable_estimation": true}}` + setupSettingsDir(t, base, local) + + s, err := Load(context.Background()) + if err != nil { + t.Fatalf("Load: %v", err) + } + // Local override flips disable_estimation on; base models survive (models key + // absent from the override, so it is not replaced). + if !s.IsEstimationDisabled() { + t.Error("local override did not enable disable_estimation") + } + if len(s.PricingOverrides()) != 1 || s.PricingOverrides()[0].ID != "base-model" { + t.Errorf("overrides = %+v, want base-model preserved", s.PricingOverrides()) + } +} + +func TestPricingSettings_LocalOverrideReplacesModels(t *testing.T) { + base := `{ + "enabled": true, + "pricing": {"models": [{"id": "base-model", "provider": "test", "input_per_mtok": 1, "output_per_mtok": 1}]} + }` + local := `{"pricing": {"models": [{"id": "local-model", "provider": "test", "input_per_mtok": 2, "output_per_mtok": 2}]}}` + setupSettingsDir(t, base, local) + + s, err := Load(context.Background()) + if err != nil { + t.Fatalf("Load: %v", err) + } + if len(s.PricingOverrides()) != 1 || s.PricingOverrides()[0].ID != "local-model" { + t.Errorf("overrides = %+v, want models replaced by local-model", s.PricingOverrides()) + } +} diff --git a/cmd/entire/cli/settings/settings_test.go b/cmd/entire/cli/settings/settings_test.go index 90a10f68f3..cd14b82bb6 100644 --- a/cmd/entire/cli/settings/settings_test.go +++ b/cmd/entire/cli/settings/settings_test.go @@ -465,6 +465,53 @@ func TestMergeJSON_CommitLinking(t *testing.T) { } } +func TestPricingSettings_CodexServiceTier(t *testing.T) { + // Accessor defaults to "" and is nil-safe. + var nilSettings *EntireSettings + if got := nilSettings.CodexServiceTier(); got != "" { + t.Errorf("nil settings tier = %q, want empty", got) + } + if got := (&EntireSettings{}).CodexServiceTier(); got != "" { + t.Errorf("empty settings tier = %q, want empty", got) + } + if got := (&EntireSettings{Pricing: &PricingSettings{}}).CodexServiceTier(); got != "" { + t.Errorf("empty pricing tier = %q, want empty", got) + } + + // omitempty: an unset tier does not serialize. + data, err := json.Marshal(&PricingSettings{}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if strings.Contains(string(data), "codex_service_tier") { + t.Errorf("empty PricingSettings serialized the tier: %s", data) + } + + // Load merges pricing.codex_service_tier through the strict (unknown-field- + // rejecting) decoder and the map-based merge, and the accessor reads it back. + tmpDir := t.TempDir() + entireDir := filepath.Join(tmpDir, ".entire") + if err := os.MkdirAll(entireDir, 0o755); err != nil { + t.Fatalf("mkdir .entire: %v", err) + } + if err := os.WriteFile(filepath.Join(entireDir, "settings.json"), + []byte(`{"enabled": true, "pricing": {"codex_service_tier": "priority"}}`), 0o644); err != nil { + t.Fatalf("write settings.json: %v", err) + } + if err := os.MkdirAll(filepath.Join(tmpDir, ".git"), 0o755); err != nil { + t.Fatalf("mkdir .git: %v", err) + } + t.Chdir(tmpDir) + + s, err := Load(context.Background()) + if err != nil { + t.Fatalf("Load: %v", err) + } + if got := s.CodexServiceTier(); got != "priority" { + t.Errorf("CodexServiceTier() = %q, want priority", got) + } +} + func TestExternalAgents_DefaultsFalse(t *testing.T) { s := &EntireSettings{} if s.ExternalAgents { diff --git a/cmd/entire/cli/strategy/accumulate_model_usage_test.go b/cmd/entire/cli/strategy/accumulate_model_usage_test.go new file mode 100644 index 0000000000..9a00160215 --- /dev/null +++ b/cmd/entire/cli/strategy/accumulate_model_usage_test.go @@ -0,0 +1,144 @@ +package strategy + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" +) + +// mu builds a single-model bucket for terse test construction. +func mu(model string, u agent.TokenUsage) agent.ModelUsage { + return agent.ModelUsage{Model: model, TokenUsage: u} +} + +// accumulateModelUsage must fold buckets into a per-model map, unioning model +// keys across steps and summing token counts + costs for repeated models. +func TestAccumulateModelUsage_MultiStepUnionAndSum(t *testing.T) { + t.Parallel() + + // Step 1: two models. + m := accumulateModelUsage(nil, []agent.ModelUsage{ + mu("opus", agent.TokenUsage{InputTokens: 10, OutputTokens: 2, CostUSD: costPtr(0.5), CostSource: types.CostSourceEstimated}), + mu("haiku", agent.TokenUsage{InputTokens: 4, CostUSD: costPtr(0.1), CostSource: types.CostSourceEstimated}), + }) + // Step 2: opus again (must sum) plus a new model sonnet (union). + m = accumulateModelUsage(m, []agent.ModelUsage{ + mu("opus", agent.TokenUsage{InputTokens: 5, OutputTokens: 3, CostUSD: costPtr(0.25), CostSource: types.CostSourceEstimated}), + mu("sonnet", agent.TokenUsage{InputTokens: 7, CostUSD: costPtr(0.2), CostSource: types.CostSourceReported}), + }) + + if len(m) != 3 { + t.Fatalf("model count = %d, want 3 (opus, haiku, sonnet)", len(m)) + } + + opus := m["opus"] + if opus == nil { + t.Fatal("opus bucket missing") + } + if opus.InputTokens != 15 || opus.OutputTokens != 5 { + t.Fatalf("opus tokens = in %d out %d, want in 15 out 5", opus.InputTokens, opus.OutputTokens) + } + if opus.CostUSD == nil || *opus.CostUSD != 0.75 { + t.Fatalf("opus cost = %v, want 0.75", opus.CostUSD) + } + if opus.CostSource != types.CostSourceEstimated { + t.Fatalf("opus source = %q, want estimated", opus.CostSource) + } + if m["haiku"].InputTokens != 4 { + t.Fatalf("haiku input = %d, want 4", m["haiku"].InputTokens) + } + if m["sonnet"].InputTokens != 7 || m["sonnet"].CostSource != types.CostSourceReported { + t.Fatalf("sonnet bucket = %+v, want in 7 reported", *m["sonnet"]) + } +} + +// A repeated model whose steps carry differing cost sources folds to mixed. +func TestAccumulateModelUsage_MixedCostSource(t *testing.T) { + t.Parallel() + m := accumulateModelUsage(nil, []agent.ModelUsage{ + mu("opus", agent.TokenUsage{InputTokens: 1, CostUSD: costPtr(1.0), CostSource: types.CostSourceReported}), + }) + m = accumulateModelUsage(m, []agent.ModelUsage{ + mu("opus", agent.TokenUsage{InputTokens: 1, CostUSD: costPtr(0.5), CostSource: types.CostSourceEstimated}), + }) + if m["opus"].CostUSD == nil || *m["opus"].CostUSD != 1.5 { + t.Fatalf("opus cost = %v, want 1.5", m["opus"].CostUSD) + } + if m["opus"].CostSource != types.CostSourceMixed { + t.Fatalf("opus source = %q, want mixed", m["opus"].CostSource) + } +} + +// The incoming buckets must not be aliased: mutating the accumulated map must not +// reach back into the caller's bucket slice. +func TestAccumulateModelUsage_CopiesBuckets(t *testing.T) { + t.Parallel() + incoming := []agent.ModelUsage{mu("opus", agent.TokenUsage{InputTokens: 10, CostUSD: costPtr(0.5)})} + m := accumulateModelUsage(nil, incoming) + m["opus"].InputTokens = 999 + if incoming[0].TokenUsage.InputTokens != 10 { + t.Fatalf("mutating accumulated map mutated incoming bucket: %d", incoming[0].TokenUsage.InputTokens) + } + if m["opus"].CostUSD == incoming[0].TokenUsage.CostUSD { + t.Fatal("accumulated cost pointer aliases incoming bucket") + } +} + +// Empty incoming leaves the map untouched (nil stays nil). +func TestAccumulateModelUsage_EmptyIncoming(t *testing.T) { + t.Parallel() + if got := accumulateModelUsage(nil, nil); got != nil { + t.Fatalf("nil+nil = %v, want nil", got) + } + existing := map[string]*agent.TokenUsage{"opus": {InputTokens: 3}} + if got := accumulateModelUsage(existing, nil); len(got) != 1 || got["opus"].InputTokens != 3 { + t.Fatalf("existing+nil mutated map: %v", got) + } +} + +// Reset-on-checkpoint boundary: after condensation sets state.ModelUsage = nil, +// the next step must start a fresh map rather than re-adding to prior totals. +func TestAccumulateModelUsage_ResetBoundary(t *testing.T) { + t.Parallel() + m := accumulateModelUsage(nil, []agent.ModelUsage{mu("opus", agent.TokenUsage{InputTokens: 100})}) + if m["opus"].InputTokens != 100 { + t.Fatalf("pre-reset opus = %d, want 100", m["opus"].InputTokens) + } + // Condensation boundary reset. + m = nil + // Next checkpoint window. + m = accumulateModelUsage(m, []agent.ModelUsage{mu("opus", agent.TokenUsage{InputTokens: 7})}) + if m["opus"].InputTokens != 7 { + t.Fatalf("post-reset opus = %d, want 7 (prior total must not carry over)", m["opus"].InputTokens) + } +} + +// sortedModelUsage flattens the map into a slice ordered by model for +// deterministic serialization; an empty map yields nil (omitempty drops it). +func TestSortedModelUsage_DeterministicOrder(t *testing.T) { + t.Parallel() + if got := sortedModelUsage(nil); got != nil { + t.Fatalf("nil map = %v, want nil", got) + } + m := map[string]*agent.TokenUsage{ + "sonnet": {InputTokens: 1}, + "opus": {InputTokens: 2}, + "haiku": {InputTokens: 3}, + } + got := sortedModelUsage(m) + want := []string{"haiku", "opus", "sonnet"} + if len(got) != len(want) { + t.Fatalf("len = %d, want %d", len(got), len(want)) + } + for i, model := range want { + if got[i].Model != model { + t.Fatalf("got[%d].Model = %q, want %q", i, got[i].Model, model) + } + } + // Value is a copy, not a shared pointer. + got[1].TokenUsage.InputTokens = 999 + if m["opus"].InputTokens != 2 { + t.Fatalf("sortedModelUsage aliased the map's pointer: %d", m["opus"].InputTokens) + } +} diff --git a/cmd/entire/cli/strategy/accumulate_token_cost_test.go b/cmd/entire/cli/strategy/accumulate_token_cost_test.go new file mode 100644 index 0000000000..e80838ebd9 --- /dev/null +++ b/cmd/entire/cli/strategy/accumulate_token_cost_test.go @@ -0,0 +1,122 @@ +package strategy + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" +) + +func costPtr(v float64) *float64 { return &v } + +// accumulateTokenUsage must carry CostUSD/CostSource forward, copy (not alias) +// the cost pointer when existing is nil, and recurse cost into SubagentTokens. +func TestAccumulateTokenUsage_CarriesCost_ExistingNil(t *testing.T) { + t.Parallel() + incoming := &agent.TokenUsage{ + InputTokens: 10, + CostUSD: costPtr(0.5), + CostSource: types.CostSourceReported, + } + got := accumulateTokenUsage(nil, incoming) + if got.CostUSD == nil || *got.CostUSD != 0.5 { + t.Fatalf("CostUSD = %v, want 0.5", got.CostUSD) + } + if got.CostSource != types.CostSourceReported { + t.Fatalf("CostSource = %q, want reported", got.CostSource) + } + // Must not alias incoming.CostUSD. + if got.CostUSD == incoming.CostUSD { + t.Fatal("result aliased incoming.CostUSD") + } + *got.CostUSD = 99 + if *incoming.CostUSD != 0.5 { + t.Fatalf("mutating result mutated incoming: %v", *incoming.CostUSD) + } +} + +func TestAccumulateTokenUsage_SumsCost_MixedSource(t *testing.T) { + t.Parallel() + existing := &agent.TokenUsage{InputTokens: 5, CostUSD: costPtr(1.0), CostSource: types.CostSourceReported} + incoming := &agent.TokenUsage{InputTokens: 3, CostUSD: costPtr(0.25), CostSource: types.CostSourceEstimated} + got := accumulateTokenUsage(existing, incoming) + if got.CostUSD == nil || *got.CostUSD != 1.25 { + t.Fatalf("CostUSD = %v, want 1.25", got.CostUSD) + } + if got.CostSource != types.CostSourceMixed { + t.Fatalf("CostSource = %q, want mixed", got.CostSource) + } +} + +// Two independent aggregates (mirroring state.TokenUsage and +// state.CheckpointTokenUsage in manual_commit_git) that fold the SAME step must +// not share the step's SubagentTokens object; otherwise mutating one aggregate's +// subtree would cross-contaminate the other and the shared step. Subagent tokens +// are "latest snapshot wins" (replace, not sum) because each step already carries +// the cumulative-since-session-start total, so folding the same 100-token step +// twice leaves each aggregate at exactly 100 — not 200 — with each holding its +// own copy of the subtree. +func TestAccumulateTokenUsage_SubagentNoSharedPointerAcrossAggregates(t *testing.T) { + t.Parallel() + step := &agent.TokenUsage{ + InputTokens: 10, + SubagentTokens: &agent.TokenUsage{InputTokens: 100}, + } + + var aggA, aggB *agent.TokenUsage + // Step 1: seed both aggregates from the same step. + aggA = accumulateTokenUsage(aggA, step) + aggB = accumulateTokenUsage(aggB, step) + // Step 2: fold the same step again into each. + aggA = accumulateTokenUsage(aggA, step) + aggB = accumulateTokenUsage(aggB, step) + + if aggA.SubagentTokens.InputTokens != 100 { + t.Fatalf("aggA subagent input = %d, want 100 (latest snapshot wins, not summed)", aggA.SubagentTokens.InputTokens) + } + if aggB.SubagentTokens.InputTokens != 100 { + t.Fatalf("aggB subagent input = %d, want 100 (latest snapshot wins, not summed)", aggB.SubagentTokens.InputTokens) + } + // The incoming step must never be mutated by accumulation. + if step.SubagentTokens.InputTokens != 100 { + t.Fatalf("step subagent mutated: %d, want 100", step.SubagentTokens.InputTokens) + } + // The two aggregates must not share the same SubagentTokens object. + if aggA.SubagentTokens == aggB.SubagentTokens { + t.Fatal("aggregates share the same SubagentTokens pointer") + } +} + +// Subagent usage is "latest snapshot wins" (replace, not sum): each step already +// carries the cumulative-since-session-start subagent total, so the incoming +// step's SubagentTokens (and its cost) supersede whatever was recorded before. +// The snapshot is deep-copied, so the aggregate never aliases the step's subtree. +func TestAccumulateTokenUsage_SubagentReplacedByLatestSnapshot(t *testing.T) { + t.Parallel() + existing := &agent.TokenUsage{ + InputTokens: 5, + SubagentTokens: &agent.TokenUsage{InputTokens: 1, CostUSD: costPtr(0.25), CostSource: types.CostSourceReported}, + } + incoming := &agent.TokenUsage{ + InputTokens: 3, + SubagentTokens: &agent.TokenUsage{InputTokens: 2, CostUSD: costPtr(0.5), CostSource: types.CostSourceReported}, + } + got := accumulateTokenUsage(existing, incoming) + if got.SubagentTokens == nil { + t.Fatal("expected SubagentTokens carried") + } + // Replace, not sum: the incoming snapshot's subagent tokens and cost win. + if got.SubagentTokens.InputTokens != 2 { + t.Fatalf("subagent input = %d, want 2 (latest snapshot wins)", got.SubagentTokens.InputTokens) + } + if got.SubagentTokens.CostUSD == nil || *got.SubagentTokens.CostUSD != 0.5 { + t.Fatalf("subagent CostUSD = %v, want 0.5 (latest snapshot wins)", got.SubagentTokens.CostUSD) + } + if got.SubagentTokens.CostSource != types.CostSourceReported { + t.Fatalf("subagent CostSource = %q, want reported", got.SubagentTokens.CostSource) + } + // Deep copy: the aggregate must not alias the incoming step's subtree. + if got.SubagentTokens == incoming.SubagentTokens { + t.Fatal("result aliased incoming.SubagentTokens") + } +} diff --git a/cmd/entire/cli/strategy/condensation_codex_tier_test.go b/cmd/entire/cli/strategy/condensation_codex_tier_test.go new file mode 100644 index 0000000000..c37065867a --- /dev/null +++ b/cmd/entire/cli/strategy/condensation_codex_tier_test.go @@ -0,0 +1,146 @@ +package strategy + +import ( + "context" + "math" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/settings" +) + +// codexTierTranscript is a minimal Codex rollout: a turn_context naming the +// serving model (as every real rollout carries before its token_counts) +// followed by a single cumulative token_count of 35000 total input of which +// 20000 cached, and 8000 output. Over an empty baseline that yields +// fresh=15000, cache_read=20000, output=8000 (Codex never reports +// cache-creation). The turn_context keeps the fixture valid for both pricing +// paths: without a ModelUsageCalculator the flat usage falls back to a single +// bucket under the caller's fallback model; with one, the snapshot is +// attributed to gpt-5.5 and applyTierVariant retargets that bucket onto the +// tier-variant id. +const codexTierTranscript = `{"type":"turn_context","payload":{"cwd":"/tmp/repo","model":"gpt-5.5","effort":"high","summary":"auto"}} +{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":35000,"cached_input_tokens":20000,"output_tokens":8000,"total_tokens":63000}}}}` + +// Expected gpt-5.5 costs for the fixture above, computed from the embedded table: +// +// standard (5/30, cache_read 0.5): 15000*5 + 20000*0.5 + 8000*30 = 325000 -> $0.325 +// priority (12.5/75, cache_read 1.25): x2.5 -> $0.8125 +// +// The priority table is exactly 2.5x standard, so priority == 2.5*standard is the +// invariant that proves the "-priority" suffix took effect at condensation. This +// mirrors the live regression, where a priority Codex session was persisted at the +// discarded standard $0.154513 instead of the expected priority $0.3862825 +// (= 2.5 x $0.154513). +const ( + codexTierStandardCost = 0.325 + codexTierPriorityCost = 0.8125 +) + +// ctxWithCodexTier returns a context whose settings resolve to the given Codex +// service tier via an explicit worktree root (no chdir, parallel-safe). An empty +// tier writes no settings file, so estimation runs on the embedded defaults. +func ctxWithCodexTier(t *testing.T, tier string) context.Context { + t.Helper() + dir := t.TempDir() + if tier != "" { + entireDir := filepath.Join(dir, ".entire") + if err := os.MkdirAll(entireDir, 0o755); err != nil { + t.Fatalf("mkdir .entire: %v", err) + } + body := `{"enabled": true, "pricing": {"codex_service_tier": "` + tier + `"}}` + if err := os.WriteFile(filepath.Join(entireDir, "settings.json"), []byte(body), 0o644); err != nil { + t.Fatalf("write settings.json: %v", err) + } + } + return settings.WithWorktreeRoot(context.Background(), dir) +} + +// TestCondensationRepricesCodexPriorityTier is the regression for the defect where +// condensation re-derived the committed per-model buckets from the raw model name, +// discarding the turn-end priority-tier pricing. tokenUsageWithCost is the shared +// choke point every condensation/extraction pricing pass funnels through; with the +// "priority" knob it must price the Codex fallback bucket under gpt-5.5-priority. +func TestCondensationRepricesCodexPriorityTier(t *testing.T) { + t.Parallel() + + ag, err := agent.GetByAgentType(agent.AgentTypeCodex) + if err != nil || ag == nil { + t.Fatalf("GetByAgentType(codex) = %v, %v", ag, err) + } + transcript := []byte(codexTierTranscript) + + // Priority: fallback bucket priced under gpt-5.5-priority. + prioUsage, prioBuckets := tokenUsageWithCost(ctxWithCodexTier(t, "priority"), ag, transcript, 0, "gpt-5.5", agent.AgentTypeCodex) + assertCost(t, "priority flat", prioUsage.CostUSD, codexTierPriorityCost) + assertSingleBucket(t, "priority", prioBuckets, "gpt-5.5-priority", codexTierPriorityCost) + + // Standard (no knob): fallback bucket priced under gpt-5.5. + stdUsage, stdBuckets := tokenUsageWithCost(ctxWithCodexTier(t, ""), ag, transcript, 0, "gpt-5.5", agent.AgentTypeCodex) + assertCost(t, "standard flat", stdUsage.CostUSD, codexTierStandardCost) + assertSingleBucket(t, "standard", stdBuckets, "gpt-5.5", codexTierStandardCost) + + // The suffix multiplies cost by exactly 2.5 (priority rates are 2.5x standard). + if math.Abs(*prioUsage.CostUSD-2.5*(*stdUsage.CostUSD)) > 1e-9 { + t.Errorf("priority %.10f != 2.5 x standard %.10f", *prioUsage.CostUSD, *stdUsage.CostUSD) + } + + // A non-Codex agent type is unaffected even with the priority knob present. + claudeUsage, claudeBuckets := tokenUsageWithCost(ctxWithCodexTier(t, "priority"), ag, transcript, 0, "gpt-5.5", agent.AgentTypeClaudeCode) + assertCost(t, "non-codex flat", claudeUsage.CostUSD, codexTierStandardCost) + assertSingleBucket(t, "non-codex", claudeBuckets, "gpt-5.5", codexTierStandardCost) +} + +// TestCondensationDowngradesUnpriceableTierVariant covers the tier knob on a +// Codex model whose "-priority" variant has no table entry (only gpt-5.5 has +// one): the session must price at the standard rate under the raw id — an +// honest undercount — rather than going entirely unpriced under a +// premium-claiming id. +func TestCondensationDowngradesUnpriceableTierVariant(t *testing.T) { + t.Parallel() + + ag, err := agent.GetByAgentType(agent.AgentTypeCodex) + if err != nil || ag == nil { + t.Fatalf("GetByAgentType(codex) = %v, %v", ag, err) + } + // gpt-5.3-codex over the shared fixture counts (fresh=15000, + // cache_read=20000, output=8000) at 1.75/14 with cache_read 0.175: + // 15000*1.75 + 20000*0.175 + 8000*14 = 141750 -> $0.14175. + transcript := []byte(`{"type":"turn_context","payload":{"cwd":"/tmp/repo","model":"gpt-5.3-codex","effort":"high","summary":"auto"}} +{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":35000,"cached_input_tokens":20000,"output_tokens":8000,"total_tokens":63000}}}}`) + const standardCost = 0.14175 + + usage, buckets := tokenUsageWithCost(ctxWithCodexTier(t, "priority"), ag, transcript, 0, "gpt-5.3-codex", agent.AgentTypeCodex) + assertCost(t, "downgraded flat", usage.CostUSD, standardCost) + assertSingleBucket(t, "downgraded", buckets, "gpt-5.3-codex", standardCost) +} + +func assertCost(t *testing.T, label string, got *float64, want float64) { + t.Helper() + if got == nil { + t.Fatalf("%s cost = nil, want %.10f", label, want) + } + if math.Abs(*got-want) > 1e-9 { + t.Fatalf("%s cost = %.10f, want %.10f", label, *got, want) + } +} + +func assertSingleBucket(t *testing.T, label string, buckets []types.ModelUsage, wantModel string, wantCost float64) { + t.Helper() + if len(buckets) != 1 { + t.Fatalf("%s: got %d buckets, want 1", label, len(buckets)) + } + b := buckets[0] + if b.Model != wantModel { + t.Errorf("%s: bucket model = %q, want %q", label, b.Model, wantModel) + } + if b.TokenUsage.CostUSD == nil || math.Abs(*b.TokenUsage.CostUSD-wantCost) > 1e-9 { + t.Errorf("%s: bucket cost = %v, want %.10f", label, b.TokenUsage.CostUSD, wantCost) + } + if b.TokenUsage.CostSource != types.CostSourceEstimated { + t.Errorf("%s: bucket source = %q, want estimated", label, b.TokenUsage.CostSource) + } +} diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index e8ee29eba8..7ea1ce9d6c 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -275,15 +275,19 @@ func (s *ManualCommitStrategy) CondenseSession(ctx context.Context, repo *git.Re sessionData.TokenUsage = accumulateTokenUsage(nil, state.CheckpointTokenUsage) } + // Mirror the TokenUsage backfill for per-model usage: when the freshly-extracted + // transcript yields no per-model buckets (e.g. Cursor, whose usage arrives only + // via stop-hook payloads accumulated in state), fall back to the checkpoint-scoped + // per-model map so the checkpoint metadata still carries a per-model breakdown. + if !hasModelUsageData(sessionData.ModelUsage) && len(state.ModelUsage) > 0 { + sessionData.ModelUsage = sortedModelUsage(state.ModelUsage) + } + // Backfill the model from the transcript for agents that don't report it via // hooks (e.g., Pi records message.model but its hook events carry no model // field). Only fills when the model is otherwise unknown — hook-reported // models take precedence. - if state.ModelName == "" { - if model := sessionStateBackfillModel(ctx, ag, sessionData.Transcript); model != "" { - state.ModelName = model - } - } + backfillModelAndReprice(ctx, ag, sessionData, state) // Skip gate: if there is no transcript AND no files touched, there is nothing // meaningful to condense. Return early to avoid writing metadata-only stubs. @@ -382,6 +386,7 @@ func (s *ManualCommitStrategy) CondenseSession(ctx context.Context, repo *git.Re TranscriptIdentifierAtStart: state.TranscriptIdentifierAtStart, CheckpointTranscriptStart: state.CheckpointTranscriptStart, TokenUsage: sessionData.TokenUsage, + ModelUsage: sessionData.ModelUsage, SkillEvents: skillEvents, SessionMetrics: buildSessionMetrics(state), Attribution: attribution, @@ -556,7 +561,7 @@ func (s *ManualCommitStrategy) extractOrCreateSessionData(ctx context.Context, r case hasShadowBranch: // Shadow branch exists (from SaveStep commits) — extract transcript and // metadata from the branch tree, preferring the live transcript if fresher. - data, err := s.extractSessionData(ctx, repo, shadowHash, state.SessionID, state.FilesTouched, state.AgentType, state.TranscriptPath, state.CheckpointTranscriptStart, state.Phase.IsActive()) + data, err := s.extractSessionData(ctx, repo, shadowHash, state.SessionID, state.FilesTouched, state.AgentType, state.TranscriptPath, state.CheckpointTranscriptStart, state.ModelName, state.Phase.IsActive()) if err != nil { return nil, fmt.Errorf("failed to extract session data: %w", err) } @@ -736,6 +741,128 @@ func hasTokenUsageData(usage *agent.TokenUsage) bool { return hasTokenUsageData(usage.SubagentTokens) } +// hasModelUsageData reports whether any per-model bucket carries real token data. +// It mirrors hasTokenUsageData so the ModelUsage backfill triggers on exactly the +// same "no usable data" condition (an empty slice, or buckets with only zero +// counts, both fall back to the accumulated state map). +func hasModelUsageData(buckets []agent.ModelUsage) bool { + for i := range buckets { + bucket := buckets[i].TokenUsage + if hasTokenUsageData(&bucket) { + return true + } + } + return false +} + +// sessionDataUnpricedFromTranscript reports whether transcript-extracted usage +// came out unpriced in a way that a just-recovered model can fix: the flat usage +// carries data but no cost, or a bucket keyed under the empty model still holds +// token data (the signature of a pricing pass that ran before the model was +// known). It first requires the flat usage to carry data so the reprice never +// fires in the state.ModelUsage mirror-fallback case — where the transcript +// produced no usage and re-pricing would clobber the mirrored per-model buckets. +func sessionDataUnpricedFromTranscript(d *ExtractedSessionData) bool { + if !hasTokenUsageData(d.TokenUsage) { + return false + } + if d.TokenUsage.CostUSD == nil { + return true + } + for i := range d.ModelUsage { + if d.ModelUsage[i].Model == "" { + bucket := d.ModelUsage[i].TokenUsage + if hasTokenUsageData(&bucket) { + return true + } + } + } + return false +} + +// tokenUsageWithCost computes token usage from a transcript slice and attributes +// cost using the configured pricing table. It mirrors agent.CalculateTokenUsage's +// nil-on-error contract (errors are debug-logged, never fatal to condensation) +// and returns the per-model buckets alongside the flat usage. fallbackModel +// prices agents that don't attribute usage per model; an empty value leaves such +// usage unpriced (agents that implement ModelUsageCalculator, e.g. pi and gemini, +// are priced from their own per-model buckets regardless). Condensation never +// passes a subagents dir (see the TODOs at the call sites), so none is taken. +// +// agentType selects the pricing service tier: a Codex session opted into the +// "priority" tier (pricing.codex_service_tier) prices the fallback bucket under +// the model's "-priority" variant, matching the turn-end hook so a committed +// checkpoint carries the same tier the turn was priced at. The suffix is a +// pricing-lookup concern only — the stored model name stays raw. +func tokenUsageWithCost(ctx context.Context, ag agent.Agent, transcript []byte, fromOffset int, fallbackModel string, agentType types.AgentType) (*agent.TokenUsage, []agent.ModelUsage) { + table, disableEstimation := settings.LoadPricingTable(ctx) + pricingModel := settings.PricingModelForAgent(ctx, agentType, fallbackModel) + usage, buckets, err := agent.CalculateUsageWithCost(ag, transcript, fromOffset, "", table, pricingModel, disableEstimation) + if err != nil { + logging.Debug(ctx, "failed usage-with-cost extraction", + slog.String("error", err.Error())) + return nil, nil + } + return usage, buckets +} + +// backfillModelAndReprice recovers the model from the transcript for agents +// that don't report it via hooks (e.g., Pi records message.model but its hook +// events carry no model field) and re-prices the extracted usage with it. Only +// fills when the model is otherwise unknown — hook-reported models take +// precedence. +// +// The pricing pass in extractOrCreateSessionData ran while the model was still +// unknown, so it used an empty fallbackModel and left the usage unpriced under +// an empty-model bucket. Re-pricing the SAME extracted slice with the recovered +// model is equivalent to the original pass (both extraction paths price from +// state.CheckpointTranscriptStart with no subagents dir). The swap requires the +// re-extraction to carry real tokens — a zero slice (e.g. a Pi fork abandoning +// the branch that consumed the tokens) must not clobber usage restored from +// checkpoint state. When the state total aliases the pre-reprice usage (Pi's +// backfill returns the checkpoint usage as-is), it is repointed to the priced +// copy so the session-state diagnostic matches the persisted checkpoint. +// +// Copilot CLI is a second, independent case: sessionStateBackfillTokenUsage (run +// earlier, before this function) prices the full-session transcript with +// whatever model was known at that point, so when the model was still unknown +// state.TokenUsage ends up as its OWN unpriced usage — a distinct object from +// sessionData.TokenUsage's checkpoint-scoped slice, so the alias check above +// never fires for it. Reprice that full-session total too, guarded the same way +// (recovered model, real token data required) so a zero-data recovery window can +// never clobber a good value. +func backfillModelAndReprice(ctx context.Context, ag agent.Agent, sessionData *ExtractedSessionData, state *SessionState) { + if state.ModelName != "" { + return + } + model := sessionStateBackfillModel(ctx, ag, sessionData.Transcript) + if model == "" { + return + } + state.ModelName = model + + if state.AgentType == agent.AgentTypeCopilotCLI && state.TokenUsage != sessionData.TokenUsage && + hasTokenUsageData(state.TokenUsage) && state.TokenUsage.CostUSD == nil { + fullUsage, _ := tokenUsageWithCost(ctx, ag, sessionData.Transcript, 0, state.ModelName, state.AgentType) + if hasTokenUsageData(fullUsage) { + state.TokenUsage = fullUsage + } + } + + if !sessionDataUnpricedFromTranscript(sessionData) { + return + } + usage, buckets := tokenUsageWithCost(ctx, ag, sessionData.Transcript, state.CheckpointTranscriptStart, state.ModelName, state.AgentType) + if !hasTokenUsageData(usage) { + return + } + if state.TokenUsage == sessionData.TokenUsage { + state.TokenUsage = usage + } + sessionData.TokenUsage = usage + sessionData.ModelUsage = buckets +} + // applyBackfilledSessionTokenUsage overwrites state.TokenUsage with the // transcript-recomputed session total (see sessionStateBackfillTokenUsage) when // one is available, preserving the cumulative subagent total across the backfill. @@ -750,7 +877,7 @@ func hasTokenUsageData(usage *agent.TokenUsage) bool { // is never mixed into checkpointUsage, which is the checkpoint-scoped value // written to metadata. func applyBackfilledSessionTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte, checkpointUsage *agent.TokenUsage) { - backfillUsage := sessionStateBackfillTokenUsage(ctx, ag, state.AgentType, transcript, checkpointUsage) + backfillUsage := sessionStateBackfillTokenUsage(ctx, ag, state.AgentType, transcript, state.ModelName, checkpointUsage) if backfillUsage == nil { return } @@ -768,9 +895,11 @@ func applyBackfilledSessionTokenUsage(ctx context.Context, ag agent.Agent, state // sessionStateBackfillTokenUsage returns the best session-level token usage to // persist in session state after condensation. -func sessionStateBackfillTokenUsage(ctx context.Context, ag agent.Agent, agentType types.AgentType, transcript []byte, checkpointUsage *agent.TokenUsage) *agent.TokenUsage { +func sessionStateBackfillTokenUsage(ctx context.Context, ag agent.Agent, agentType types.AgentType, transcript []byte, sessionModel string, checkpointUsage *agent.TokenUsage) *agent.TokenUsage { if agentType == agent.AgentTypeCopilotCLI && len(transcript) > 0 { - fullSessionUsage := agent.CalculateTokenUsage(ctx, ag, transcript, 0, "") + // Session-wide backfill (state.TokenUsage); per-model buckets are + // checkpoint-scoped and not persisted here, so discard them. + fullSessionUsage, _ := tokenUsageWithCost(ctx, ag, transcript, 0, sessionModel, agentType) if hasTokenUsageData(fullSessionUsage) { return fullSessionUsage } @@ -986,7 +1115,7 @@ func committedFilesExcludingMetadata(committedFiles map[string]struct{}) []strin // This handles the case where SaveStep was skipped (no code changes) but the transcript // continued growing — the shadow branch copy would be stale. // checkpointTranscriptStart is the line offset (Claude) or message index (Gemini) where the current checkpoint began. -func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git.Repository, shadowRef plumbing.Hash, sessionID string, filesTouched []string, agentType types.AgentType, liveTranscriptPath string, checkpointTranscriptStart int, isActive bool) (*ExtractedSessionData, error) { +func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git.Repository, shadowRef plumbing.Hash, sessionID string, filesTouched []string, agentType types.AgentType, liveTranscriptPath string, checkpointTranscriptStart int, sessionModel string, isActive bool) (*ExtractedSessionData, error) { ag, _ := agent.GetByAgentType(agentType) //nolint:errcheck // ag may be nil for unknown agent types; callers use type assertions so nil is safe commit, err := repo.CommitObject(shadowRef) if err != nil { @@ -1054,7 +1183,9 @@ func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git // extract them from offset 0; consumers can filter by checkpoint_transcript_start // if they only render the checkpoint-scoped slice. if len(data.Transcript) > 0 { - data.TokenUsage = agent.CalculateTokenUsage(ctx, ag, data.Transcript, checkpointTranscriptStart, "") //TODO: why do we not use here subagents dir? + // sessionModel prices the fallback and remainder buckets (e.g. Claude + // Code subagent shortfall) that CalculateModelUsage cannot attribute. + data.TokenUsage, data.ModelUsage = tokenUsageWithCost(ctx, ag, data.Transcript, checkpointTranscriptStart, sessionModel, agentType) //TODO: why do we not use here subagents dir? data.SkillEvents = agent.ExtractSkillEvents(ctx, ag, data.Transcript, 0) } @@ -1096,7 +1227,7 @@ func (s *ManualCommitStrategy) extractSessionDataFromLiveTranscript(ctx context. // extract them from offset 0; consumers can filter by checkpoint_transcript_start // if they only render the checkpoint-scoped slice. if len(data.Transcript) > 0 { - data.TokenUsage = agent.CalculateTokenUsage(ctx, ag, data.Transcript, state.CheckpointTranscriptStart, "") //TODO: why do we not use here subagents dir? + data.TokenUsage, data.ModelUsage = tokenUsageWithCost(ctx, ag, data.Transcript, state.CheckpointTranscriptStart, state.ModelName, state.AgentType) //TODO: why do we not use here subagents dir? data.SkillEvents = agent.ExtractSkillEvents(ctx, ag, data.Transcript, 0) } diff --git a/cmd/entire/cli/strategy/manual_commit_condensation_test.go b/cmd/entire/cli/strategy/manual_commit_condensation_test.go index 438d255fc7..ee33895b40 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation_test.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation_test.go @@ -264,7 +264,7 @@ func TestSessionStateBackfillTokenUsage_CopilotUsesZeroInputSessionAggregate(t * require.Zero(t, checkpointUsage.InputTokens) require.Equal(t, 25, checkpointUsage.OutputTokens) - backfillUsage := sessionStateBackfillTokenUsage(context.Background(), ag, agent.AgentTypeCopilotCLI, transcript, checkpointUsage) + backfillUsage := sessionStateBackfillTokenUsage(context.Background(), ag, agent.AgentTypeCopilotCLI, transcript, "", checkpointUsage) require.NotNil(t, backfillUsage) require.Zero(t, backfillUsage.InputTokens) require.Equal(t, 50, backfillUsage.OutputTokens) @@ -521,6 +521,207 @@ func TestCondenseSession_TagsCheckpointSummaryWithHasInvestigation(t *testing.T) require.Equal(t, "Why is checkout flaky?", meta.InvestigateTopic, "per-session InvestigateTopic") } +// TestCondenseSession_RepricesTranscriptModelAfterBackfill verifies the +// reprice-after-backfill fix: Pi reports no model through hooks, so +// state.ModelName is "" at condensation and the extraction pricing pass ran with +// an empty fallback model — leaving the checkpoint usage unpriced under an +// empty-model bucket. Once CondenseSession backfills the model from the +// transcript it must re-price the same slice, so the persisted checkpoint carries +// a real cost bucketed under the model instead of Model="gpt-5.5" with a nil cost +// and ModelUsage=[{Model:"",…}]. +// +// Uses t.Chdir for CWD-based git resolution, so it cannot run in parallel. +func TestCondenseSession_RepricesTranscriptModelAfterBackfill(t *testing.T) { + dir := setupGitRepo(t) + t.Chdir(dir) + + repo, err := git.PlainOpen(dir) + require.NoError(t, err) + + s := &ManualCommitStrategy{} + sessionID := "2026-06-01-pi-reprice" + + metadataDir := ".entire/metadata/" + sessionID + metadataDirAbs := filepath.Join(dir, metadataDir) + require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) + + // Pi JSONL transcript: the assistant message carries model=gpt-5.5 and usage, + // but no model is reported through hooks (the whole point of ModelExtractor). + transcript := strings.Join([]string{ + `{"type":"session","version":3,"id":"pi-uuid","cwd":"/tmp"}`, + `{"type":"message","id":"m1","parentId":null,"message":{"role":"user","content":[{"type":"text","text":"Hi"}]}}`, + `{"type":"message","id":"m2","parentId":"m1","message":{"role":"assistant","content":[{"type":"text","text":"Hello"}],"model":"gpt-5.5","provider":"openai-codex","usage":{"input":1000000,"output":500000,"cacheRead":0,"cacheWrite":0}}}`, + }, "\n") + "\n" + require.NoError(t, os.WriteFile(filepath.Join(metadataDirAbs, paths.TranscriptFileName), []byte(transcript), 0o644)) + + // Modify a tracked file so SaveStep produces a non-empty session + shadow branch. + require.NoError(t, os.WriteFile(filepath.Join(dir, "test.txt"), []byte("agent-modified content"), 0o644)) + + require.NoError(t, s.SaveStep(context.Background(), StepContext{ + SessionID: sessionID, + ModifiedFiles: []string{"test.txt"}, + MetadataDir: metadataDir, + MetadataDirAbs: metadataDirAbs, + CommitMessage: "Pi checkpoint 1", + AuthorName: "Test", + AuthorEmail: "test@test.com", + })) + + state, err := s.loadSessionState(context.Background(), sessionID) + require.NoError(t, err) + + // The bug scenario: Pi agent, model unknown at condensation, first checkpoint + // (offset 0 so the whole transcript is in scope). + state.AgentType = agent.AgentTypePi + state.ModelName = "" + state.CheckpointTranscriptStart = 0 + require.NoError(t, SaveSessionState(context.Background(), state)) + + checkpointID := id.MustCheckpointID("abcdef012345") + result, err := s.CondenseSession(context.Background(), repo, checkpointID, state, nil) + require.NoError(t, err) + require.False(t, result.Skipped, "condensation must not skip when files are touched") + + // The model was backfilled from the transcript in-memory. + require.Equal(t, "gpt-5.5", state.ModelName) + + // Read persisted per-session metadata off the metadata branch. + ref, err := repo.Reference(plumbing.NewBranchReferenceName(paths.MetadataBranchName), true) + require.NoError(t, err) + commit, err := repo.CommitObject(ref.Hash()) + require.NoError(t, err) + tree, err := commit.Tree() + require.NoError(t, err) + checkpointTree, err := tree.Tree(checkpointID.Path()) + require.NoError(t, err) + + sessionMeta, err := checkpointTree.File(checkpointID.Path() + "/0/" + paths.MetadataFileName) + if err != nil { + subtree, subErr := checkpointTree.Tree("0") + require.NoError(t, subErr) + sessionMeta, err = subtree.File(paths.MetadataFileName) + require.NoError(t, err) + } + sessionBytes, err := sessionMeta.Contents() + require.NoError(t, err) + var meta checkpoint.Metadata + require.NoError(t, json.Unmarshal([]byte(sessionBytes), &meta)) + + // The persisted checkpoint carries the backfilled model AND the token + // breakdown, but NOT cost: the CLI no longer persists cost (entire-api prices + // server-side from these tokens). The reprice/rebucket after model backfill is + // what makes the persisted per-model breakdown correct — that is preserved; + // only the cost stamp is dropped on write. + require.Equal(t, "gpt-5.5", meta.Model, "persisted model") + require.NotNil(t, meta.TokenUsage, "token usage must be persisted") + require.Equal(t, 1000000, meta.TokenUsage.InputTokens, "input tokens from transcript") + require.Nil(t, meta.TokenUsage.CostUSD, "cost must NOT be persisted") + require.Empty(t, meta.TokenUsage.CostSource, "cost source must NOT be persisted") + + // The per-model breakdown must be bucketed under the real model, never "", + // with its four token fields intact (the platform prices from these). + require.NotEmpty(t, meta.ModelUsage, "per-model breakdown must be present") + var gpt55 *types.TokenUsage + for i := range meta.ModelUsage { + require.NotEmpty(t, meta.ModelUsage[i].Model, "no bucket may be keyed under the empty model") + require.Nil(t, meta.ModelUsage[i].TokenUsage.CostUSD, "per-model cost must NOT be persisted") + require.Empty(t, meta.ModelUsage[i].TokenUsage.CostSource, "per-model cost source must NOT be persisted") + if meta.ModelUsage[i].Model == "gpt-5.5" { + u := meta.ModelUsage[i].TokenUsage + gpt55 = &u + } + } + require.NotNil(t, gpt55, "usage must be bucketed under gpt-5.5") + require.Equal(t, 1000000, gpt55.InputTokens, "gpt-5.5 bucket carries the token counts") + + // The session-state total (local, not platform-facing) still carries the + // repriced cost: the compute/reprice pipeline is intact; only the checkpoint + // write drops cost. + require.NotNil(t, state.TokenUsage, "state total present") + require.NotNil(t, state.TokenUsage.CostUSD, "state total still carries the repriced cost (local only)") +} + +// A zero-token re-extraction (e.g. a Pi fork abandoning the branch that +// consumed the tokens) must not clobber usage restored from checkpoint state: +// the reprice swap requires real tokens in the re-extracted slice. +func TestCondenseSession_RepriceDoesNotClobberStateRestoredUsage(t *testing.T) { + dir := setupGitRepo(t) + t.Chdir(dir) + + repo, err := git.PlainOpen(dir) + require.NoError(t, err) + + s := &ManualCommitStrategy{} + sessionID := "2026-06-01-pi-fork" + + metadataDir := ".entire/metadata/" + sessionID + metadataDirAbs := filepath.Join(dir, metadataDir) + require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) + + // The assistant message (with the model) sits BEFORE the checkpoint window: + // the model backfill scans the full transcript (offset 0) and recovers + // gpt-5.5, while the usage extraction starts at CheckpointTranscriptStart + // and sees nothing — the zero-data re-extraction that must not clobber the + // state-restored usage. + transcript := strings.Join([]string{ + `{"type":"session","version":3,"id":"pi-uuid","cwd":"/tmp"}`, + `{"type":"message","id":"m1","parentId":null,"message":{"role":"user","content":[{"type":"text","text":"Hi"}]}}`, + `{"type":"message","id":"m2","parentId":"m1","message":{"role":"assistant","content":[{"type":"text","text":"Hello"}],"model":"gpt-5.5","provider":"openai-codex","usage":{"input":5,"output":5,"cacheRead":0,"cacheWrite":0}}}`, + }, "\n") + "\n" + require.NoError(t, os.WriteFile(filepath.Join(metadataDirAbs, paths.TranscriptFileName), []byte(transcript), 0o644)) + + require.NoError(t, os.WriteFile(filepath.Join(dir, "test.txt"), []byte("fork content"), 0o644)) + require.NoError(t, s.SaveStep(context.Background(), StepContext{ + SessionID: sessionID, + ModifiedFiles: []string{"test.txt"}, + MetadataDir: metadataDir, + MetadataDirAbs: metadataDirAbs, + CommitMessage: "Pi fork checkpoint", + AuthorName: "Test", + AuthorEmail: "test@test.com", + })) + + state, err := s.loadSessionState(context.Background(), sessionID) + require.NoError(t, err) + state.AgentType = agent.AgentTypePi + state.ModelName = "" + state.CheckpointTranscriptStart = 3 // window starts past m2: zero in-window usage + // Real usage accumulated in checkpoint state from earlier in the session, + // bucketed under the empty model (it was unknown at SaveStep time). + state.CheckpointTokenUsage = &agent.TokenUsage{InputTokens: 777000, OutputTokens: 111} + state.ModelUsage = map[string]*agent.TokenUsage{"": {InputTokens: 777000, OutputTokens: 111}} + require.NoError(t, SaveSessionState(context.Background(), state)) + + checkpointID := id.MustCheckpointID("fedcba987654") + result, err := s.CondenseSession(context.Background(), repo, checkpointID, state, nil) + require.NoError(t, err) + require.False(t, result.Skipped) + + ref, err := repo.Reference(plumbing.NewBranchReferenceName(paths.MetadataBranchName), true) + require.NoError(t, err) + commit, err := repo.CommitObject(ref.Hash()) + require.NoError(t, err) + tree, err := commit.Tree() + require.NoError(t, err) + checkpointTree, err := tree.Tree(checkpointID.Path()) + require.NoError(t, err) + sessionMeta, err := checkpointTree.File(checkpointID.Path() + "/0/" + paths.MetadataFileName) + if err != nil { + subtree, subErr := checkpointTree.Tree("0") + require.NoError(t, subErr) + sessionMeta, err = subtree.File(paths.MetadataFileName) + require.NoError(t, err) + } + sessionBytes, err := sessionMeta.Contents() + require.NoError(t, err) + var meta checkpoint.Metadata + require.NoError(t, json.Unmarshal([]byte(sessionBytes), &meta)) + + // The state-restored usage must survive: 777000 input tokens, not zero. + require.NotNil(t, meta.TokenUsage, "restored usage must be persisted") + require.Equal(t, 777000, meta.TokenUsage.InputTokens, "zero re-extraction must not clobber state-restored tokens") +} + // TestCheckpointStepCount covers the prompt-window math that produces the // displayed "steps" count: SessionTurnCount - PromptWindowBase, floored at 1. func TestCheckpointStepCount(t *testing.T) { @@ -551,3 +752,88 @@ func TestCheckpointStepCount(t *testing.T) { }) } } + +// TestBackfillModelAndReprice_CopilotSessionStateRepricedOnModelRecovery proves +// the session-state diagnostic total (state.TokenUsage) is repriced alongside +// the checkpoint-scoped sessionData.TokenUsage once the model is recovered, +// for the Copilot CLI full-session-backfill case where state.TokenUsage is a +// distinct object from sessionData.TokenUsage (so the pre-existing +// "state.TokenUsage == sessionData.TokenUsage" alias check never fires for it). +// The test uses Pi's real ModelExtractor/TokenCalculator implementation to +// drive model recovery, while forcing state.AgentType to Copilot CLI to +// exercise that branch in isolation. +func TestBackfillModelAndReprice_CopilotSessionStateRepricedOnModelRecovery(t *testing.T) { + t.Parallel() + + transcript := []byte(strings.Join([]string{ + `{"type":"session","version":3,"id":"pi-uuid","cwd":"/tmp"}`, + `{"type":"message","id":"m1","parentId":null,"message":{"role":"user","content":[{"type":"text","text":"Hi"}]}}`, + `{"type":"message","id":"m2","parentId":"m1","message":{"role":"assistant","content":[{"type":"text","text":"Hello"}],"model":"gpt-5.5","provider":"openai-codex","usage":{"input":100,"output":50,"cacheRead":0,"cacheWrite":0}}}`, + }, "\n") + "\n") + + ag, err := agent.GetByAgentType(agent.AgentTypePi) + require.NoError(t, err) + + sessionData := &ExtractedSessionData{ + Transcript: transcript, + TokenUsage: &agent.TokenUsage{InputTokens: 100, OutputTokens: 50}, + } + // A distinct object from sessionData.TokenUsage, standing in for the + // full-session total sessionStateBackfillTokenUsage would have produced + // earlier with an empty fallback model (unpriced). + sessionStateTotal := &agent.TokenUsage{InputTokens: 100, OutputTokens: 50} + state := &session.State{ + AgentType: agent.AgentTypeCopilotCLI, + TokenUsage: sessionStateTotal, + } + + backfillModelAndReprice(context.Background(), ag, sessionData, state) + + require.Equal(t, "gpt-5.5", state.ModelName) + + require.NotSame(t, sessionStateTotal, state.TokenUsage) + require.NotNil(t, state.TokenUsage.CostUSD, "session-state diagnostic total was not repriced") + require.Equal(t, types.CostSourceEstimated, state.TokenUsage.CostSource) + + require.NotNil(t, sessionData.TokenUsage.CostUSD, "checkpoint-scoped sessionData total was not repriced") + require.Equal(t, types.CostSourceEstimated, sessionData.TokenUsage.CostSource) +} + +// TestBackfillModelAndReprice_ZeroDataRecoveryDoesNotClobberSessionState proves +// the Copilot session-state reprice guard: if the recovered-model pass yields +// no token data (e.g. a truncated re-extraction), the pre-existing +// state.TokenUsage is left untouched rather than being clobbered with a +// zero/unpriced value. +func TestBackfillModelAndReprice_ZeroDataRecoveryDoesNotClobberSessionState(t *testing.T) { + t.Parallel() + + // Transcript carries a model on the assistant message, but no usage data + // (so tokenUsageWithCost's recovered-model pass yields zero tokens). + transcript := []byte(strings.Join([]string{ + `{"type":"session","version":3,"id":"pi-uuid","cwd":"/tmp"}`, + `{"type":"message","id":"m1","parentId":null,"message":{"role":"user","content":[{"type":"text","text":"Hi"}]}}`, + `{"type":"message","id":"m2","parentId":"m1","message":{"role":"assistant","content":[{"type":"text","text":"Hello"}],"model":"gpt-5.5","provider":"openai-codex"}}`, + }, "\n") + "\n") + + ag, err := agent.GetByAgentType(agent.AgentTypePi) + require.NoError(t, err) + + // No sessionData.TokenUsage of its own, so the checkpoint-scoped reprice + // branch is a no-op here and the test isolates the session-state branch. + sessionData := &ExtractedSessionData{Transcript: transcript} + unpricedUsage := &agent.TokenUsage{InputTokens: 100, OutputTokens: 50} + state := &session.State{ + AgentType: agent.AgentTypeCopilotCLI, + TokenUsage: unpricedUsage, + } + + backfillModelAndReprice(context.Background(), ag, sessionData, state) + + // The model is still recovered (unguarded by token data)... + require.Equal(t, "gpt-5.5", state.ModelName) + // ...but the recovered-model transcript pass yields zero token data (no + // usage on the assistant message), so state.TokenUsage must survive + // untouched rather than being clobbered with an empty priced value. + require.Same(t, unpricedUsage, state.TokenUsage) + require.Nil(t, state.TokenUsage.CostUSD) +} diff --git a/cmd/entire/cli/strategy/manual_commit_git.go b/cmd/entire/cli/strategy/manual_commit_git.go index 03803941fc..142ce6f8fd 100644 --- a/cmd/entire/cli/strategy/manual_commit_git.go +++ b/cmd/entire/cli/strategy/manual_commit_git.go @@ -151,6 +151,9 @@ func (s *ManualCommitStrategy) SaveStep(ctx context.Context, step StepContext) e state.TokenUsage.SubagentTokens, state.SubagentTokensBaseline) } } + if len(step.ModelUsage) > 0 { + state.ModelUsage = accumulateModelUsage(state.ModelUsage, step.ModelUsage) + } if !branchExisted { logging.Info(logging.WithComponent(ctx, "checkpoint"), "created shadow branch and committed changes", @@ -354,18 +357,29 @@ func accumulateTokenUsage(existing, incoming *agent.TokenUsage) *agent.TokenUsag return existing } if existing == nil { - // Return a copy to avoid sharing the pointer + // Return a copy to avoid sharing the pointer. SubagentTokens must be + // deep-copied too: aliasing incoming's subtree lets a later in-place + // accumulation mutate incoming (and cross-contaminate any other aggregate + // that folded the same step), inflating subagent token counts. return &agent.TokenUsage{ InputTokens: incoming.InputTokens, CacheCreationTokens: incoming.CacheCreationTokens, CacheReadTokens: incoming.CacheReadTokens, OutputTokens: incoming.OutputTokens, APICallCount: incoming.APICallCount, - SubagentTokens: incoming.SubagentTokens, + SubagentTokens: accumulateTokenUsage(nil, incoming.SubagentTokens), + // AddCostUSD(nil, ...) returns a fresh pointer, never aliasing incoming.CostUSD. + CostUSD: types.AddCostUSD(nil, incoming.CostUSD), + CostSource: types.MergeCostSource("", incoming.CostSource, nil, incoming.CostUSD), } } - // Accumulate values + // Accumulate values. Merge the cost source BEFORE summing CostUSD and the + // token fields, so the merge still sees the pre-sum existing usage (its + // possibly-nil cost and its token counts). MergeCostSourceUsages folds a + // priced side with an unpriced-but-token-bearing side to mixed. + existing.CostSource = types.MergeCostSourceUsages(existing, incoming) + existing.CostUSD = types.AddCostUSD(existing.CostUSD, incoming.CostUSD) existing.InputTokens += incoming.InputTokens existing.CacheCreationTokens += incoming.CacheCreationTokens existing.CacheReadTokens += incoming.CacheReadTokens @@ -374,26 +388,71 @@ func accumulateTokenUsage(existing, incoming *agent.TokenUsage) *agent.TokenUsag // Replace (not add) subagent tokens: incoming.SubagentTokens is already // the cumulative total as of this step, so the latest snapshot supersedes - // whatever was recorded before rather than stacking on top of it. + // whatever was recorded before rather than stacking on top of it. Deep-copy + // the snapshot (not alias it) so the two aggregates that fold the same step — + // state.TokenUsage and state.CheckpointTokenUsage — never share, and thus can + // never cross-contaminate, the step's SubagentTokens object. if incoming.SubagentTokens != nil { - existing.SubagentTokens = incoming.SubagentTokens + existing.SubagentTokens = accumulateTokenUsage(nil, incoming.SubagentTokens) } return existing } +// accumulateModelUsage merges per-model token buckets into an existing +// checkpoint-scoped map keyed by model. Each bucket is folded into its model's +// entry with accumulateTokenUsage (token sums plus AddCostUSD/MergeCostSource +// cost folding). Returns existing unchanged when incoming is empty; lazily +// allocates the map on first use. The map is the per-model mirror of +// CheckpointTokenUsage and is reset together with it at every condensation. +func accumulateModelUsage(existing map[string]*agent.TokenUsage, incoming []agent.ModelUsage) map[string]*agent.TokenUsage { + if len(incoming) == 0 { + return existing + } + if existing == nil { + existing = make(map[string]*agent.TokenUsage, len(incoming)) + } + for i := range incoming { + bucket := incoming[i].TokenUsage + existing[incoming[i].Model] = accumulateTokenUsage(existing[incoming[i].Model], &bucket) + } + return existing +} + +// sortedModelUsage flattens a per-model usage map into a slice sorted by model, +// for deterministic serialization into checkpoint metadata. Returns nil for an +// empty map so the omitempty JSON tag drops the field entirely. +func sortedModelUsage(m map[string]*agent.TokenUsage) []agent.ModelUsage { + if len(m) == 0 { + return nil + } + models := make([]string, 0, len(m)) + for model := range m { + models = append(models, model) + } + sort.Strings(models) + out := make([]agent.ModelUsage, 0, len(models)) + for _, model := range models { + out = append(out, agent.ModelUsage{Model: model, TokenUsage: *m[model]}) + } + return out +} + // resetCheckpointWindow resets the per-checkpoint accumulation window after a // condensation reset. It zeroes the step count, clears the checkpoint-scoped -// token usage, and snapshots the cumulative subagent total into -// SubagentTokensBaseline so the next window's CheckpointTokenUsage.SubagentTokens -// can be rescoped to "since this condensation" rather than re-reporting the full -// cumulative subagent total (see accumulateTokenUsage and the SaveStep rescoping -// in this file, plus SessionState.SubagentTokensBaseline). Shared by all three -// condensation reset sites (CondenseSessionByID, CondenseAndMarkFullyCondensed, -// condenseAndUpdateState) so the baseline capture cannot drift between them. +// token usage and its per-model mirror (ModelUsage), and snapshots the +// cumulative subagent total into SubagentTokensBaseline so the next window's +// CheckpointTokenUsage.SubagentTokens can be rescoped to "since this +// condensation" rather than re-reporting the full cumulative subagent total +// (see accumulateTokenUsage and the SaveStep rescoping in this file, plus +// SessionState.SubagentTokensBaseline). ModelUsage is the per-model mirror of +// CheckpointTokenUsage and is reset with it. Shared by all three condensation +// reset sites (CondenseSessionByID, CondenseAndMarkFullyCondensed, +// condenseAndUpdateState) so the reset cannot drift between them. func resetCheckpointWindow(state *SessionState) { state.StepCount = 0 state.CheckpointTokenUsage = nil + state.ModelUsage = nil state.RebaselineSubagentTokens() } diff --git a/cmd/entire/cli/strategy/manual_commit_types.go b/cmd/entire/cli/strategy/manual_commit_types.go index 8ac6bd31da..37805299e7 100644 --- a/cmd/entire/cli/strategy/manual_commit_types.go +++ b/cmd/entire/cli/strategy/manual_commit_types.go @@ -63,5 +63,6 @@ type ExtractedSessionData struct { Prompts []string // User prompts from the current checkpoint portion FilesTouched []string TokenUsage *agent.TokenUsage // Token usage calculated from transcript (since CheckpointTranscriptStart) + ModelUsage []agent.ModelUsage // Per-model breakdown of TokenUsage (since CheckpointTranscriptStart) SkillEvents []agent.SkillEvent // Skill events detected from transcript data } diff --git a/cmd/entire/cli/strategy/strategy.go b/cmd/entire/cli/strategy/strategy.go index 05b4071ee5..983ab91328 100644 --- a/cmd/entire/cli/strategy/strategy.go +++ b/cmd/entire/cli/strategy/strategy.go @@ -163,6 +163,10 @@ type StepContext struct { // TokenUsage contains the token usage for this checkpoint TokenUsage *agent.TokenUsage + + // ModelUsage is the per-model breakdown of TokenUsage for this checkpoint, + // carried alongside TokenUsage so the strategy can accumulate per-model usage. + ModelUsage []agent.ModelUsage } // TaskStepContext contains all information needed for saving a task step checkpoint. diff --git a/cmd/entire/cli/tokens_pricing_refresh.go b/cmd/entire/cli/tokens_pricing_refresh.go new file mode 100644 index 0000000000..c72affaf75 --- /dev/null +++ b/cmd/entire/cli/tokens_pricing_refresh.go @@ -0,0 +1,79 @@ +package cli + +import ( + "context" + "fmt" + "time" + + "github.com/entireio/cli/cmd/entire/cli/pricing" + "github.com/entireio/cli/cmd/entire/cli/settings" + "github.com/spf13/cobra" +) + +// pricingRefreshReport is the --json shape of the pricing-refresh outcome. +type pricingRefreshReport struct { + SourceURL string `json:"source_url"` + Outcome string `json:"outcome"` + Entries int `json:"entries"` + FetchedAt time.Time `json:"fetched_at"` + StalenessSecs int `json:"staleness_seconds"` + RemoteEnabled bool `json:"remote_enabled"` +} + +func newTokensPricingRefreshCmd() *cobra.Command { + var jsonFlag bool + cmd := &cobra.Command{ + Use: "pricing-refresh", + Short: "Force-refresh the cached remote pricing table", + Long: `Force-refresh the locally cached remote pricing table. + +Fetches the remote pricing table now, bypassing the daily throttle, and reports +the outcome. This refreshes the cache regardless of settings so you can inspect +what would be merged; remote rates are only merged into cost estimates when the +"pricing.remote" setting is enabled.`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + return runTokensPricingRefresh(cmd.Context(), cmd, jsonFlag) + }, + } + cmd.Flags().BoolVar(&jsonFlag, "json", false, "Output as JSON") + return cmd +} + +func runTokensPricingRefresh(ctx context.Context, cmd *cobra.Command, jsonOutput bool) error { + result, err := pricing.RefreshRemoteCacheForce(ctx) + if err != nil { + return fmt.Errorf("refresh remote pricing cache: %w", err) + } + enabled := settings.IsRemoteEnabled(ctx) + + if jsonOutput { + return printJSON(cmd.OutOrStdout(), pricingRefreshReport{ + SourceURL: result.SourceURL, + Outcome: result.Outcome, + Entries: result.Entries, + FetchedAt: result.FetchedAt, + StalenessSecs: int(result.Staleness().Seconds()), + RemoteEnabled: enabled, + }) + } + + w := cmd.OutOrStdout() + fmt.Fprintln(w, "Remote pricing refresh") + fmt.Fprintf(w, "Source: %s\n", result.SourceURL) + fmt.Fprintf(w, "Outcome: %s\n", result.Outcome) + fmt.Fprintf(w, "Entries: %d\n", result.Entries) + if result.FetchedAt.IsZero() { + fmt.Fprintln(w, "Fetched at: never") + } else { + fmt.Fprintf(w, "Fetched at: %s (%s ago)\n", + result.FetchedAt.UTC().Format(time.RFC3339), + result.Staleness().Round(time.Second)) + } + if enabled { + fmt.Fprintln(w, "Remote merge: enabled") + } else { + fmt.Fprintln(w, `Remote merge: disabled — set "pricing": {"remote": true} in .entire/settings.json to merge remote rates into cost estimates`) + } + return nil +} diff --git a/cmd/entire/cli/tokens_profile.go b/cmd/entire/cli/tokens_profile.go index b069a6add9..6804f925cd 100644 --- a/cmd/entire/cli/tokens_profile.go +++ b/cmd/entire/cli/tokens_profile.go @@ -7,8 +7,10 @@ import ( "io" "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" "github.com/entireio/cli/cmd/entire/cli/checkpoint" "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" + "github.com/entireio/cli/cmd/entire/cli/pricing" "github.com/spf13/cobra" ) @@ -18,9 +20,12 @@ type tokensProfileReport struct { CheckpointsAvailable int `json:"checkpoints_available"` CheckpointsAnalyzed int `json:"checkpoints_analyzed"` CheckpointsWithTokenData int `json:"checkpoints_with_token_data"` + CheckpointsWithCostData int `json:"checkpoints_with_cost_data,omitempty"` MissingTokenData int `json:"missing_token_data"` MetadataReadWarnings int `json:"metadata_read_warnings,omitempty"` Tokens *sessionTokensUsage `json:"tokens,omitempty"` + CostUSD *float64 `json:"cost_usd,omitempty"` + CostSource string `json:"cost_source,omitempty"` Signals []tokensProfileSignal `json:"signals,omitempty"` Recommendations []sessionTokensRecommendation `json:"recommendations,omitempty"` Limitations []string `json:"limitations,omitempty"` @@ -56,17 +61,20 @@ func newTokensGroupCmd() *cobra.Command { Long: `Analyze token usage across sessions and checkpoints. Commands: - profile Aggregate token usage across committed checkpoints + profile Aggregate token usage across committed checkpoints + pricing-refresh Force-refresh the cached remote pricing table Examples: entire tokens profile - entire tokens profile --json`, + entire tokens profile --json + entire tokens pricing-refresh`, RunE: func(cmd *cobra.Command, _ []string) error { return cmd.Help() }, } cmd.AddCommand(newTokensProfileCmd()) + cmd.AddCommand(newTokensPricingRefreshCmd()) return cmd } @@ -122,7 +130,7 @@ func runTokensProfile(ctx context.Context, cmd *cobra.Command, jsonOutput bool, return fmt.Errorf("failed to list checkpoints: %w", err) } - report, err := buildTokensProfileReport(ctx, store, infos, limit) + report, err := buildTokensProfileReport(ctx, store, infos, limit, loadDisplayPricingTable(ctx)) if err != nil { return err } @@ -134,7 +142,7 @@ func runTokensProfile(ctx context.Context, cmd *cobra.Command, jsonOutput bool, return nil } -func buildTokensProfileReport(ctx context.Context, store checkpoint.PersistentStore, infos []checkpoint.CheckpointInfo, limit int) (tokensProfileReport, error) { +func buildTokensProfileReport(ctx context.Context, store checkpoint.PersistentStore, infos []checkpoint.CheckpointInfo, limit int, table *pricing.Table) (tokensProfileReport, error) { checkpointsAvailable := len(infos) infos = limitTokensProfileCheckpoints(infos, limit) report := tokensProfileReport{ @@ -161,7 +169,7 @@ func buildTokensProfileReport(ctx context.Context, store checkpoint.PersistentSt continue } - usage, metadataReadWarning, err := tokensProfileCheckpointUsage(ctx, store, info.CheckpointID, summary) + usage, buckets, model, metadataReadWarning, err := tokensProfileCheckpointUsage(ctx, store, info.CheckpointID, summary) if err != nil { return tokensProfileReport{}, err } @@ -175,12 +183,32 @@ func buildTokensProfileReport(ctx context.Context, store checkpoint.PersistentSt continue } + // Local cost estimate for this checkpoint (the CLI no longer persists + // cost). Stamped onto usage so addCheckpointTokenUsage folds the + // per-checkpoint estimates into the aggregate total/source. + if usage != nil { + cost, source := agent.EstimateCost(usage, buckets, model, table) + usage.CostUSD = cost + usage.CostSource = source + } + report.CheckpointsWithTokenData++ + if usage != nil && usage.CostUSD != nil { + report.CheckpointsWithCostData++ + } aggregate = addCheckpointTokenUsage(aggregate, usage) addTokensProfileTokenSignals(signals, info.CheckpointID, tokens, report.CheckpointsAnalyzed) } report.Tokens = buildSessionTokensUsage(aggregate) + if report.Tokens != nil && aggregate != nil { + // buildSessionTokensUsage no longer carries cost; surface the folded + // local estimate from the aggregate directly. + report.Tokens.CostUSD = aggregate.CostUSD + report.Tokens.CostSource = aggregate.CostSource + report.CostUSD = aggregate.CostUSD + report.CostSource = aggregate.CostSource + } report.Signals = orderedTokensProfileSignals(signals) report.Recommendations = tokensProfileRecommendations(report) report.Limitations = tokensProfileLimitations(report) @@ -194,9 +222,9 @@ func limitTokensProfileCheckpoints(infos []checkpoint.CheckpointInfo, limit int) return infos[:limit] } -func tokensProfileCheckpointUsage(ctx context.Context, store checkpoint.PersistentStore, checkpointID id.CheckpointID, summary *checkpoint.CheckpointSummary) (*agent.TokenUsage, bool, error) { +func tokensProfileCheckpointUsage(ctx context.Context, store checkpoint.PersistentStore, checkpointID id.CheckpointID, summary *checkpoint.CheckpointSummary) (*agent.TokenUsage, []types.ModelUsage, string, bool, error) { if summary == nil { - return nil, false, nil + return nil, nil, "", false, nil } metas := make([]*checkpoint.Metadata, 0, len(summary.Sessions)) @@ -205,14 +233,28 @@ func tokensProfileCheckpointUsage(ctx context.Context, store checkpoint.Persiste meta, err := store.ReadSessionMetadata(ctx, checkpointID, i) if err != nil { if ctxErr := ctx.Err(); ctxErr != nil { - return nil, false, ctxErr //nolint:wrapcheck // Propagating context cancellation. + return nil, nil, "", false, ctxErr //nolint:wrapcheck // Propagating context cancellation. } metadataReadWarning = true continue } metas = append(metas, meta) } - return checkpointTokenUsage(summary, metas, metadataReadWarning), metadataReadWarning, nil + usage := checkpointTokenUsage(summary, metas, metadataReadWarning) + buckets := checkpointModelUsage(summary, metas, metadataReadWarning) + return usage, buckets, firstCheckpointModelName(metas), metadataReadWarning, nil +} + +// firstCheckpointModelName returns the first non-empty model id across the +// checkpoint's session metadata, used as the fallback pricing model when a +// checkpoint carries no per-model buckets. Empty when no model is recorded. +func firstCheckpointModelName(metas []*checkpoint.Metadata) string { + for _, meta := range metas { + if meta != nil && meta.Model != "" { + return meta.Model + } + } + return "" } func addTokensProfileTokenSignals(signals map[string]*tokensProfileSignal, checkpointID id.CheckpointID, tokens *sessionTokensUsage, denominator int) { @@ -379,6 +421,12 @@ func writeTokensProfileText(w io.Writer, report tokensProfileReport) { } writeTokenUsageSectionWithTitle(w, "Checkpoint-observed token usage", report.Tokens) + if report.CostUSD != nil { + fmt.Fprintf(w, "Total cost: %s across %d of %d checkpoints\n", + formatCostUSD(report.CostUSD, report.CostSource), + report.CheckpointsWithCostData, report.CheckpointsAnalyzed) + fmt.Fprintf(w, " %s\n", localCostEstimateNote) + } writeTokensProfileSignals(w, report.Signals) if len(report.Recommendations) > 0 { writeTokenRecommendations(w, report.Recommendations) diff --git a/cmd/entire/cli/tokens_profile_test.go b/cmd/entire/cli/tokens_profile_test.go index a50cb56b93..b6e4eee963 100644 --- a/cmd/entire/cli/tokens_profile_test.go +++ b/cmd/entire/cli/tokens_profile_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" "github.com/entireio/cli/cmd/entire/cli/checkpoint" "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" "github.com/entireio/cli/cmd/entire/cli/strategy" @@ -296,6 +297,117 @@ func TestTokensProfileCmd_EmptyHistory(t *testing.T) { } } +// The profile total is a LOCAL estimate summed across checkpoints (the CLI no +// longer persists cost). Two checkpoints under test-model ($1/MTok) price to +// $0.42 + $0.58 = $1.00; a third with token data under an unpriceable model +// contributes tokens but no cost, folding the aggregate source to "mixed" +// (partial coverage) and counting toward M, not N. +func TestTokensProfileCmd_TotalCostAcrossCheckpoints(t *testing.T) { + repo, _ := runExplainAutoTestRepo(t) + ctx := context.Background() + store := checkpoint.NewGitStore(repo, checkpoint.DefaultV1Refs()) + + writeProfileTokenCheckpointModel(ctx, t, store, "700ddd000001", "profile-cost-a", "test-model", &agent.TokenUsage{ + InputTokens: 400000, OutputTokens: 20000, APICallCount: 2, // $0.42 + }) + writeProfileTokenCheckpointModel(ctx, t, store, "700ddd000002", "profile-cost-b", "test-model", &agent.TokenUsage{ + InputTokens: 560000, OutputTokens: 20000, APICallCount: 3, // $0.58 + }) + writeProfileTokenCheckpointModel(ctx, t, store, "700ddd000003", "profile-cost-none", "unpriceable-model", &agent.TokenUsage{ + InputTokens: 500, OutputTokens: 50, APICallCount: 1, + }) + + infos, err := store.List(ctx) + if err != nil { + t.Fatalf("List: %v", err) + } + report, err := buildTokensProfileReport(ctx, store, infos, 0, testDisplayPricingTable(t)) + if err != nil { + t.Fatalf("buildTokensProfileReport: %v", err) + } + + if report.CostUSD == nil || *report.CostUSD != 1.00 { + t.Fatalf("total cost = %v, want 1.00", report.CostUSD) + } + if report.CostSource != types.CostSourceMixed { + t.Fatalf("cost source = %q, want mixed (partial coverage)", report.CostSource) + } + if report.CheckpointsWithCostData != 2 { + t.Fatalf("checkpoints_with_cost_data = %d, want 2", report.CheckpointsWithCostData) + } + + var buf bytes.Buffer + writeTokensProfileText(&buf, report) + out := buf.String() + if !strings.Contains(out, "Total cost: $1.00 (estimated locally, partial) across 2 of 3 checkpoints") { + t.Fatalf("expected local-estimate total cost line, got:\n%s", out) + } + if !strings.Contains(out, localCostEstimateNote) { + t.Fatalf("expected local-estimate note, got:\n%s", out) + } +} + +func TestTokensProfileCmd_JSONCost(t *testing.T) { + repo, _ := runExplainAutoTestRepo(t) + ctx := context.Background() + store := checkpoint.NewGitStore(repo, checkpoint.DefaultV1Refs()) + + writeProfileTokenCheckpointModel(ctx, t, store, "710eee000001", "profile-json-cost", "test-model", &agent.TokenUsage{ + InputTokens: 400000, OutputTokens: 20000, APICallCount: 1, // $0.42 + }) + + infos, err := store.List(ctx) + if err != nil { + t.Fatalf("List: %v", err) + } + report, err := buildTokensProfileReport(ctx, store, infos, 0, testDisplayPricingTable(t)) + if err != nil { + t.Fatalf("buildTokensProfileReport: %v", err) + } + + data, err := json.Marshal(report) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var result tokensProfileReport + if err := json.Unmarshal(data, &result); err != nil { + t.Fatalf("expected valid JSON, got parse error: %v\noutput: %s", err, data) + } + if result.CostUSD == nil || *result.CostUSD != 0.42 { + t.Fatalf("cost_usd = %v, want 0.42 (local estimate)", result.CostUSD) + } + if result.CostSource != types.CostSourceEstimated { + t.Fatalf("cost_source = %q, want estimated", result.CostSource) + } + if result.CheckpointsWithCostData != 1 { + t.Fatalf("checkpoints_with_cost_data = %d, want 1", result.CheckpointsWithCostData) + } +} + +func TestTokensProfileCmd_JSONOmitsCostWhenAbsent(t *testing.T) { + repo, _ := runExplainAutoTestRepo(t) + ctx := context.Background() + store := checkpoint.NewGitStore(repo, checkpoint.DefaultV1Refs()) + + writeProfileTokenCheckpoint(ctx, t, store, "720fff000001", "profile-json-nocost", &agent.TokenUsage{ + InputTokens: 1000, + APICallCount: 1, + }) + + cmd := newTokensGroupCmd() + var stdout bytes.Buffer + cmd.SetOut(&stdout) + cmd.SetArgs([]string{"profile", "--json"}) + + if err := cmd.ExecuteContext(ctx); err != nil { + t.Fatalf("expected no error, got: %v", err) + } + + if strings.Contains(stdout.String(), "cost_usd") { + t.Fatalf("expected cost_usd omitted when no cost tracked, got: %s", stdout.String()) + } +} + func signalCount(signals []tokensProfileSignal, id string) int { for _, signal := range signals { if signal.ID == id { @@ -307,6 +419,13 @@ func signalCount(signals []tokensProfileSignal, id string) int { func writeProfileTokenCheckpoint(ctx context.Context, t *testing.T, store *checkpoint.GitStore, checkpointID string, sessionID string, usage *agent.TokenUsage) { t.Helper() + writeProfileTokenCheckpointModel(ctx, t, store, checkpointID, sessionID, "", usage) +} + +// writeProfileTokenCheckpointModel is writeProfileTokenCheckpoint with an explicit +// model, so the display-side local cost estimate has a model to price under. +func writeProfileTokenCheckpointModel(ctx context.Context, t *testing.T, store *checkpoint.GitStore, checkpointID string, sessionID string, model string, usage *agent.TokenUsage) { + t.Helper() if err := store.Write(ctx, checkpoint.Session{ CheckpointID: id.MustCheckpointID(checkpointID), @@ -314,6 +433,7 @@ func writeProfileTokenCheckpoint(ctx context.Context, t *testing.T, store *check Strategy: strategy.StrategyNameManualCommit, Branch: "tokens-profile", Agent: testAgentClaude, + Model: model, Transcript: redact.AlreadyRedacted([]byte(`{"type":"user","message":{"content":[{"type":"text","text":"profile"}]}}` + "\n")), AuthorName: "Test", AuthorEmail: "test@example.com", diff --git a/e2e/testutil/metadata.go b/e2e/testutil/metadata.go index efa33803eb..5aadfd8197 100644 --- a/e2e/testutil/metadata.go +++ b/e2e/testutil/metadata.go @@ -3,11 +3,20 @@ package testutil import "time" type TokenUsage struct { - InputTokens int `json:"input_tokens"` - CacheCreationTokens int `json:"cache_creation_tokens"` - CacheReadTokens int `json:"cache_read_tokens"` - OutputTokens int `json:"output_tokens"` - APICallCount int `json:"api_call_count"` + InputTokens int `json:"input_tokens"` + CacheCreationTokens int `json:"cache_creation_tokens"` + CacheReadTokens int `json:"cache_read_tokens"` + OutputTokens int `json:"output_tokens"` + APICallCount int `json:"api_call_count"` + CostUSD *float64 `json:"cost_usd,omitempty"` + CostSource string `json:"cost_source,omitempty"` +} + +// ModelUsage mirrors the real per-model breakdown: a model identifier paired with +// its token usage under a nested "token_usage" key (the canonical backend shape). +type ModelUsage struct { + Model string `json:"model"` + TokenUsage TokenUsage `json:"token_usage"` } type Attribution struct { @@ -29,6 +38,7 @@ type CheckpointMetadata struct { FilesTouched []string `json:"files_touched"` Sessions []SessionRef `json:"sessions"` TokenUsage TokenUsage `json:"token_usage"` + ModelUsage []ModelUsage `json:"model_usage,omitempty"` } type SessionRef struct { @@ -40,19 +50,20 @@ type SessionRef struct { } type SessionMetadata struct { - CLIVersion string `json:"cli_version"` - CheckpointID string `json:"checkpoint_id"` - SessionID string `json:"session_id"` - Strategy string `json:"strategy"` - CreatedAt time.Time `json:"created_at"` - Branch string `json:"branch"` - Agent string `json:"agent"` - Model string `json:"model"` - CheckpointsCount int `json:"checkpoints_count"` - FilesTouched []string `json:"files_touched"` - TokenUsage TokenUsage `json:"token_usage"` - Attribution Attribution `json:"initial_attribution"` - TranscriptPath string `json:"transcript_path"` + CLIVersion string `json:"cli_version"` + CheckpointID string `json:"checkpoint_id"` + SessionID string `json:"session_id"` + Strategy string `json:"strategy"` + CreatedAt time.Time `json:"created_at"` + Branch string `json:"branch"` + Agent string `json:"agent"` + Model string `json:"model"` + CheckpointsCount int `json:"checkpoints_count"` + FilesTouched []string `json:"files_touched"` + TokenUsage TokenUsage `json:"token_usage"` + ModelUsage []ModelUsage `json:"model_usage,omitempty"` + Attribution Attribution `json:"initial_attribution"` + TranscriptPath string `json:"transcript_path"` // CheckpointTranscriptStart is the transcript.jsonl line offset where this // checkpoint's contributions begin. For the first checkpoint in a session