feat(cost): real usage across all providers, /cost subcommands, budget gates and session persistence

README.md

@@ -563,7 +563,7 @@ Credentials are stored with **AES-256-GCM** at `~/.chatcli/auth-profiles.json`.
 | **Custom personas** | Markdown with YAML frontmatter (model, tools, skills). |
 | **Hooks** | PreToolUse, PostToolUse, SessionStart/End, UserPromptSubmit, Pre/PostCompact — shell or webhook. |
 | **WebFetch / WebSearch** | DuckDuckGo + fetch with text extraction. |
-| **Cost tracking** | Per-session cost with per-provider pricing tables. |
+| **Cost tracking** | Real API usage across all providers, `/cost` (+ `reset`, `last`, `sessions`, `export`), session budgets with optional hard stop, persisted snapshots. |
 | **Git Worktrees** | Isolated work on parallel branches. |
 | **K8s Watcher** | Multi-target: metrics, logs, events, Prometheus scraping. |
 | **i18n** | Portuguese and English with automatic detection. |

README_PT.md

@@ -554,7 +554,7 @@ Credenciais armazenadas com **AES-256-GCM** em `~/.chatcli/auth-profiles.json`.
 | **Personas customizáveis** | Markdown com frontmatter YAML (model, tools, skills). |
 | **Hooks** | PreToolUse, PostToolUse, SessionStart/End, UserPromptSubmit, Compact pre/post — shell ou webhook. |
 | **WebFetch / WebSearch** | DuckDuckGo + fetch com extração de texto. |
-| **Cost tracking** | Custo por sessão com pricing tables por provider. |
+| **Cost tracking** | Uso real de API em todos os providers, `/cost` (+ `reset`, `last`, `sessions`, `export`), orçamento de sessão com hard stop opcional, snapshots persistidos. |
 | **Git Worktrees** | Trabalho isolado em branches paralelas. |
 | **K8s Watcher** | Multi-target: metrics, logs, events, Prometheus scraping. |
 | **i18n** | Português e Inglês com detecção automática. |

cli/agent/workers/dispatcher.go

@@ -10,6 +10,7 @@ import (
 	"github.com/diillson/chatcli/cli/agent/runs"
 	"github.com/diillson/chatcli/llm/client"
 	"github.com/diillson/chatcli/llm/manager"
+	"github.com/diillson/chatcli/models"
 	"go.uber.org/zap"
 )
 
@@ -36,6 +37,10 @@ type Dispatcher struct {
 	config        DispatcherConfig
 	policyChecker PolicyChecker
 	pipeline      ExecutionPipeline
+	// usageRecorder/budgetGate are boxed behind pointers so Dispatcher
+	// stays a comparable type (bare func fields would break that contract).
+	usageRecorder *usageRecorderBox
+	budgetGate    *budgetGateBox
 	logger        *zap.Logger
 }
 
@@ -362,6 +367,25 @@ func (d *Dispatcher) executeAgent(ctx context.Context, call AgentCall) AgentResu
 	workerCtx, cancel := context.WithTimeout(runCtx, d.config.WorkerTimeout)
 	defer cancel()
 
+	// Decorate the worker's client so every LLM round-trip it makes lands in
+	// the session cost tracker AS IT HAPPENS, attributed to the
+	// provider+model that served this worker — per-call recording keeps the
+	// provider-billed accounting correct and lets the budget gate see live
+	// spend. The gate itself refuses further calls once the session budget
+	// hard stop trips, so an in-flight dispatch wave stops mid-run instead
+	// of finishing its ReAct loops on borrowed money.
+	if rec, gate := d.usageRecorder, d.budgetGate; rec != nil || gate != nil {
+		record := func(*models.UsageInfo) {}
+		if rec != nil {
+			record = func(u *models.UsageInfo) { rec.fn(effProvider, effModel, u) }
+		}
+		var gateFn func() error
+		if gate != nil {
+			gateFn = gate.fn
+		}
+		llmClient = wrapWithUsageRecording(llmClient, record, gateFn)
+	}
+
 	deps := &WorkerDeps{
 		LLMClient:     llmClient,
 		LockMgr:       d.lockMgr,

cli/agent/workers/usage_tally.go

@@ -0,0 +1,153 @@
+/*
+ * ChatCLI - Command Line Interface for LLM interaction
+ * Copyright (c) 2024 Edilson Freitas
+ * License: Apache-2.0
+ */
+package workers
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/diillson/chatcli/llm/client"
+	"github.com/diillson/chatcli/models"
+)
+
+// UsageRecorder receives the token usage of ONE worker LLM call so the
+// session cost tracker can account for subagent spend — historically the
+// largest untracked slice of an agent-mode session. Called per call, not
+// per worker run: merging N calls into one UsageInfo would collapse the
+// tracker's per-call provider-billed accounting (one call carrying
+// usage.cost would mark ALL merged tokens as billed), and per-call
+// recording also keeps the tracker live mid-run so the budget gate sees
+// spend as it happens instead of only when the worker finishes.
+type UsageRecorder func(provider, model string, usage *models.UsageInfo)
+
+// BudgetGate is consulted before every worker LLM call; a non-nil error
+// refuses the call (the session budget hard stop). Kept as a callback so
+// the workers package stays decoupled from the CLI's cost tracker.
+type BudgetGate func() error
+
+// usageRecorderBox wraps the recorder func so Dispatcher can hold it via a
+// comparable pointer field.
+type usageRecorderBox struct{ fn UsageRecorder }
+
+// budgetGateBox wraps the gate func so Dispatcher stays comparable.
+type budgetGateBox struct{ fn BudgetGate }
+
+// SetUsageRecorder wires the callback that receives each worker LLM call's
+// usage, attributed to the provider+model that actually served the worker.
+// Nil disables recording (the default).
+func (d *Dispatcher) SetUsageRecorder(fn UsageRecorder) {
+	if fn == nil {
+		d.usageRecorder = nil
+		return
+	}
+	d.usageRecorder = &usageRecorderBox{fn: fn}
+}
+
+// SetBudgetGate wires the session budget hard stop into every worker LLM
+// call. Without it a dispatch wave in flight kept spending after the
+// budget was exhausted — up to a full ReAct loop per worker, times the
+// parallel wave. Nil disables the gate (the default).
+func (d *Dispatcher) SetBudgetGate(fn BudgetGate) {
+	if fn == nil {
+		d.budgetGate = nil
+		return
+	}
+	d.budgetGate = &budgetGateBox{fn: fn}
+}
+
+// recordingClient decorates a worker's LLM client so every SendPrompt /
+// SendPromptWithTools round-trip is (a) refused when the budget gate says
+// so and (b) recorded immediately — real API usage when the inner client
+// reports it, a character estimate otherwise. It preserves the tool-use
+// capability of the inner client (SupportsNativeTools answers for it), so
+// RunWorkerReAct routes exactly as it would undecorated.
+type recordingClient struct {
+	inner  client.LLMClient
+	record func(*models.UsageInfo) // never nil
+	gate   func() error            // may be nil (no budget gate)
+}
+
+func wrapWithUsageRecording(inner client.LLMClient, record func(*models.UsageInfo), gate func() error) client.LLMClient {
+	return &recordingClient{inner: inner, record: record, gate: gate}
+}
+
+func (rc *recordingClient) GetModelName() string { return rc.inner.GetModelName() }
+
+func (rc *recordingClient) SendPrompt(ctx context.Context, prompt string, history []models.Message, maxTokens int) (string, error) {
+	if rc.gate != nil {
+		if err := rc.gate(); err != nil {
+			return "", err
+		}
+	}
+	resp, err := rc.inner.SendPrompt(ctx, prompt, history, maxTokens)
+	if err == nil {
+		rc.record(client.GetUsageOrEstimate(rc.inner, promptChars(prompt, history), len(resp)))
+	}
+	return resp, err
+}
+
+// SendPromptWithTools delegates to the inner client's native tool path.
+// Only reachable when SupportsNativeTools() returned true.
+func (rc *recordingClient) SendPromptWithTools(ctx context.Context, prompt string, history []models.Message, tools []models.ToolDefinition, maxTokens int) (*models.LLMResponse, error) {
+	tac, ok := client.AsToolAware(rc.inner)
+	if !ok {
+		return nil, fmt.Errorf("worker LLM client does not support native tools")
+	}
+	if rc.gate != nil {
+		if err := rc.gate(); err != nil {
+			return nil, err
+		}
+	}
+	resp, err := tac.SendPromptWithTools(ctx, prompt, history, tools, maxTokens)
+	if err == nil {
+		if resp != nil && resp.Usage != nil {
+			rc.record(resp.Usage)
+		} else {
+			outChars := 0
+			if resp != nil {
+				outChars = len(resp.Content)
+			}
+			rc.record(client.GetUsageOrEstimate(rc.inner, promptChars(prompt, history), outChars))
+		}
+	}
+	return resp, err
+}
+
+// SupportsNativeTools answers for the inner client so the decorated client
+// keeps its exact routing behavior.
+func (rc *recordingClient) SupportsNativeTools() bool {
+	if tac, ok := client.AsToolAware(rc.inner); ok {
+		return tac.SupportsNativeTools()
+	}
+	return false
+}
+
+// LastUsage forwards the inner client's real usage so nested consumers of
+// the decorated client still see it.
+func (rc *recordingClient) LastUsage() *models.UsageInfo {
+	if uac, ok := client.AsUsageAware(rc.inner); ok {
+		return uac.LastUsage()
+	}
+	return nil
+}
+
+// LastStopReason forwards the inner client's stop reason (the ReAct loop
+// uses it to detect max_tokens truncation).
+func (rc *recordingClient) LastStopReason() string {
+	if src, ok := client.AsStopReasonAware(rc.inner); ok {
+		return src.LastStopReason()
+	}
+	return ""
+}
+
+// promptChars sizes the outgoing payload for the estimate fallback.
+func promptChars(prompt string, history []models.Message) int {
+	n := len(prompt)
+	for _, m := range history {
+		n += len(m.Content)
+	}
+	return n
+}

cli/agent/workers/usage_tally_test.go

@@ -0,0 +1,109 @@
+/*
+ * ChatCLI - worker usage recording tests
+ * Copyright (c) 2024 Edilson Freitas
+ * License: Apache-2.0
+ */
+package workers
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/diillson/chatcli/llm/client"
+	"github.com/diillson/chatcli/models"
+)
+
+// fakeWorkerClient is UsageAware + ToolAware with canned responses.
+type fakeWorkerClient struct {
+	usage *models.UsageInfo
+	tools bool
+}
+
+func (f *fakeWorkerClient) GetModelName() string { return "fake-model" }
+func (f *fakeWorkerClient) SendPrompt(_ context.Context, _ string, _ []models.Message, _ int) (string, error) {
+	return "answer", nil
+}
+func (f *fakeWorkerClient) SendPromptWithTools(_ context.Context, _ string, _ []models.Message, _ []models.ToolDefinition, _ int) (*models.LLMResponse, error) {
+	return &models.LLMResponse{Content: "tool answer", Usage: f.usage}, nil
+}
+func (f *fakeWorkerClient) SupportsNativeTools() bool    { return f.tools }
+func (f *fakeWorkerClient) LastUsage() *models.UsageInfo { return f.usage }
+func (f *fakeWorkerClient) LastStopReason() string       { return "end_turn" }
+
+// TestRecordingClientRecordsEveryCallSeparately: two rounds through the
+// wrapped client must land as TWO records — per-call recording is what
+// keeps the tracker's provider-billed accounting correct (a merged total
+// would let one billed call mark all tokens as billed).
+func TestRecordingClientRecordsEveryCallSeparately(t *testing.T) {
+	inner := &fakeWorkerClient{
+		usage: &models.UsageInfo{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120, IsReal: true},
+		tools: true,
+	}
+	var recorded []*models.UsageInfo
+	wrapped := wrapWithUsageRecording(inner, func(u *models.UsageInfo) { recorded = append(recorded, u) }, nil)
+
+	if _, err := wrapped.SendPrompt(context.Background(), "p", nil, 0); err != nil {
+		t.Fatal(err)
+	}
+	tac, ok := client.AsToolAware(wrapped)
+	if !ok || !tac.SupportsNativeTools() {
+		t.Fatal("wrapper lost the native-tools capability")
+	}
+	if _, err := tac.SendPromptWithTools(context.Background(), "p", nil, nil, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	if len(recorded) != 2 {
+		t.Fatalf("recorded %d calls, want 2 (per-call recording)", len(recorded))
+	}
+	for i, u := range recorded {
+		if u.PromptTokens != 100 || u.CompletionTokens != 20 || !u.IsReal {
+			t.Fatalf("call %d usage = %+v, want 100/20 real", i, u)
+		}
+	}
+}
+
+// TestRecordingClientEstimatesWhenInnerSilent: a usage-less inner client
+// still produces a character estimate instead of zero.
+func TestRecordingClientEstimatesWhenInnerSilent(t *testing.T) {
+	inner := &fakeWorkerClient{usage: nil}
+	var recorded []*models.UsageInfo
+	wrapped := wrapWithUsageRecording(inner, func(u *models.UsageInfo) { recorded = append(recorded, u) }, nil)
+
+	if _, err := wrapped.SendPrompt(context.Background(), "a 40-character prompt string goes here!!", nil, 0); err != nil {
+		t.Fatal(err)
+	}
+	if len(recorded) != 1 || recorded[0].PromptTokens == 0 || recorded[0].IsReal {
+		t.Fatalf("estimate fallback broken: %+v", recorded)
+	}
+}
+
+// TestRecordingClientBudgetGateRefusesCalls: once the gate errors, no
+// provider call happens — the in-flight dispatch wave stops mid-run.
+func TestRecordingClientBudgetGateRefusesCalls(t *testing.T) {
+	inner := &fakeWorkerClient{tools: true}
+	blocked := errors.New("budget exhausted")
+	calls := 0
+	wrapped := wrapWithUsageRecording(inner, func(*models.UsageInfo) { calls++ }, func() error { return blocked })
+
+	if _, err := wrapped.SendPrompt(context.Background(), "p", nil, 0); !errors.Is(err, blocked) {
+		t.Fatalf("SendPrompt not gated: %v", err)
+	}
+	tac, _ := client.AsToolAware(wrapped)
+	if _, err := tac.SendPromptWithTools(context.Background(), "p", nil, nil, 0); !errors.Is(err, blocked) {
+		t.Fatalf("SendPromptWithTools not gated: %v", err)
+	}
+	if calls != 0 {
+		t.Fatalf("gated calls still recorded usage: %d", calls)
+	}
+}
+
+// TestRecordingClientHidesToolsWhenInnerLacksThem: capability parity.
+func TestRecordingClientHidesToolsWhenInnerLacksThem(t *testing.T) {
+	inner := &fakeWorkerClient{tools: false}
+	wrapped := wrapWithUsageRecording(inner, func(*models.UsageInfo) {}, nil)
+	if tac, ok := client.AsToolAware(wrapped); ok && tac.SupportsNativeTools() {
+		t.Fatal("wrapper claims native tools the inner client lacks")
+	}
+}

cli/agent_mode.go

@@ -1947,6 +1947,13 @@ func (a *AgentMode) processAIResponseAndAct(ctx context.Context, maxTurns int) e
 		default:
 		}
 
+		// Budget hard stop: end the run before the next provider call when
+		// the session budget is exhausted and CHATCLI_BUDGET_HARD_STOP is on.
+		if err := a.cli.budgetBlockedErr(); err != nil {
+			fmt.Println(colorize("  "+err.Error(), ColorRed))
+			return err
+		}
+
 		// Honor the session /max-tokens override exactly like chat mode does.
 		// Re-read every turn (raising OR lowering reflects here), while
 		// preserving any truncation-driven escalation applied meanwhile.
@@ -2269,6 +2276,7 @@ func (a *AgentMode) processAIResponseAndAct(ctx context.Context, maxTurns int) e
 				effModel = resolution.Model
 			}
 			a.cli.costTracker.RecordRealUsage(effProvider, effModel, turnUsage)
+			a.cli.maybeAnnounceBudget()
 		}
 
 		// Para o timer e obtém a duração
@@ -4016,6 +4024,18 @@ func (a *AgentMode) initMultiAgent(ctx context.Context) bool {
 
 	a.agentDispatcher = workers.NewDispatcher(a.agentRegistry, a.cli.manager, cfg, a.logger)
 
+	// Every worker's LLM spend lands in the session cost tracker per call,
+	// attributed to the provider+model that served the worker — /cost covers
+	// subagents live, and the budget hard stop reaches into in-flight
+	// dispatch waves instead of only gating the orchestrator's next turn.
+	if a.cli.costTracker != nil {
+		tracker := a.cli.costTracker
+		a.agentDispatcher.SetUsageRecorder(func(provider, model string, usage *models.UsageInfo) {
+			tracker.RecordRealUsage(provider, model, usage)
+		})
+		a.agentDispatcher.SetBudgetGate(a.cli.budgetBlockedErr)
+	}
+
 	// Attach policy enforcement so parallel workers respect security rules
 	if pa, err := newWorkerPolicyAdapter(a.logger); err == nil {
 		pa.unattended = a.cli.unattended // gateway: auto-approve "ask" instead of blocking on stdin

cli/chat_ask.go

@@ -129,6 +129,11 @@ func (cli *ChatCLI) executeChatAskNative(
 		}
 		if resp != nil && resp.Usage != nil && cli.costTracker != nil {
 			cli.costTracker.RecordRealUsage(resolution.Provider, resolution.Model, resp.Usage)
+			// Flag the turn as recorded: handleChatTurnResult would otherwise
+			// record the SAME usage again via LastUsage() on the buffered
+			// path, doubling every chat-ask turn's tokens and dollars.
+			cli.turnUsageRecorded = true
+			cli.maybeAnnounceBudget()
 		}
 
 		var calls chatExceptionCalls

cli/chat_envelope_footer_test.go

@@ -19,16 +19,16 @@ import (
 
 func TestChatEnvelopeFooter_EmptyWithoutUsage(t *testing.T) {
 	cli := &ChatCLI{}
-	assert.Equal(t, "", cli.chatEnvelopeFooter(nil), "no usage → no footer")
-	assert.Equal(t, "", cli.chatEnvelopeFooter(&models.UsageInfo{}), "zero usage → no footer")
+	assert.Equal(t, "", cli.chatEnvelopeFooter("", "", nil), "no usage → no footer")
+	assert.Equal(t, "", cli.chatEnvelopeFooter("", "", &models.UsageInfo{}), "zero usage → no footer")
 }
 
 func TestChatEnvelopeFooter_ShowsCostAndContext(t *testing.T) {
 	t.Cleanup(func() { theme.SetProfile(theme.DetectProfile()) })
 	theme.SetProfile(theme.ProfileANSI) // keep ANSI so colorize doesn't strip in test
 
 	cli := &ChatCLI{Provider: "OPENAI", Model: "gpt-4o"}
-	footer := cli.chatEnvelopeFooter(&models.UsageInfo{PromptTokens: 1000, CompletionTokens: 500})
+	footer := cli.chatEnvelopeFooter("OPENAI", "gpt-4o", &models.UsageInfo{PromptTokens: 1000, CompletionTokens: 500})
 
 	// A known-priced model yields a cost token and a context percentage.
 	assert.Contains(t, footer, "$", "footer shows a cost")