From f9ba6c8bce135960f83b1421d1e6791dd5b78bfe Mon Sep 17 00:00:00 2001
From: Levi Neely <lkn@darkstar.example.net>
Date: Wed, 1 Apr 2026 00:31:03 +0200
Subject: [PATCH] Track and display token usage per request
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add Usage{InputTokens, OutputTokens} to backend.Response. Both
OllamaBackend (prompt_eval_count/eval_count) and OpenAIBackend
(prompt_tokens/completion_tokens) populate it. The loop emits a
'usage' OutputMsg after each Chat call; the UI displays it as
[↑N ↓N tokens].

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 agent/loop.go      | 11 ++++++++---
 backend/backend.go |  7 +++++++
 backend/ollama.go  | 12 +++++++++---
 backend/openai.go  | 12 +++++++++++-
 main.go            |  2 ++
 5 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/agent/loop.go b/agent/loop.go
index dc25e87..a9c7dd6 100644
--- a/agent/loop.go
+++ b/agent/loop.go
@@ -13,10 +13,10 @@ import (
 type ToolExecutor func(name string, args json.RawMessage) (string, error)
 
 // OutputMsg is emitted by the loop for each visible event: assistant replies,
-// tool calls, tool results, and errors.
+// tool calls, tool results, errors, and token usage.
 type OutputMsg struct {
-	Role    string // "assistant" | "tool" | "error"
-	Name    string // tool name, for Role=="tool"
+	Role    string // "assistant" | "call" | "tool" | "error" | "usage"
+	Name    string // tool name, for Role=="call"/"tool"
 	Content string
 }
 
@@ -97,6 +97,11 @@ func (l *Loop) Run(ctx context.Context, state State) error {
 			l.emit(OutputMsg{Role: "assistant", Content: resp.Message.Content})
 		}
 
+		l.emit(OutputMsg{
+			Role:    "usage",
+			Content: fmt.Sprintf("↑%d ↓%d tokens", resp.Usage.InputTokens, resp.Usage.OutputTokens),
+		})
+
 		// 4. Update: append assistant message and tool results to state.
 		if err := state.Update(resp.Message, results); err != nil {
 			return fmt.Errorf("step %d update: %w", step, err)
diff --git a/backend/backend.go b/backend/backend.go
index 4e118f8..09af2d6 100644
--- a/backend/backend.go
+++ b/backend/backend.go
@@ -30,10 +30,17 @@ type ToolCall struct {
 	Arguments json.RawMessage `json:"arguments"` // always a JSON object
 }
 
+// Usage holds token counts for a single Chat call.
+type Usage struct {
+	InputTokens  int
+	OutputTokens int
+}
+
 // Response is the model's reply for one Chat call.
 type Response struct {
 	Message    Message
 	StopReason string // "stop" | "tool_calls" | "length" | ...
+	Usage      Usage
 }
 
 // Backend is the interface all LLM providers must implement.
diff --git a/backend/ollama.go b/backend/ollama.go
index 9123fb0..95daf3a 100644
--- a/backend/ollama.go
+++ b/backend/ollama.go
@@ -58,8 +58,10 @@ type ollamaChatRequest struct {
 }
 
 type ollamaChatResponse struct {
-	Message ollamaMessage `json:"message"`
-	Done    bool          `json:"done"`
+	Message         ollamaMessage `json:"message"`
+	Done            bool          `json:"done"`
+	PromptEvalCount int           `json:"prompt_eval_count"`
+	EvalCount       int           `json:"eval_count"`
 }
 
 // -- implementation --
@@ -134,5 +136,9 @@ func (b *OllamaBackend) Chat(ctx context.Context, model string, messages []Messa
 		stopReason = "tool_calls"
 	}
 
-	return &Response{Message: msg, StopReason: stopReason}, nil
+	return &Response{
+		Message:    msg,
+		StopReason: stopReason,
+		Usage:      Usage{InputTokens: wire.PromptEvalCount, OutputTokens: wire.EvalCount},
+	}, nil
 }
diff --git a/backend/openai.go b/backend/openai.go
index 9ac0ff5..b10e658 100644
--- a/backend/openai.go
+++ b/backend/openai.go
@@ -62,8 +62,14 @@ type openAIChatRequest struct {
 	Stream   bool            `json:"stream"`
 }
 
+type openAIUsage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+}
+
 type openAIChatResponse struct {
 	Choices []openAIChoice `json:"choices"`
+	Usage   openAIUsage    `json:"usage"`
 }
 
 type openAIChoice struct {
@@ -158,5 +164,9 @@ func (b *OpenAIBackend) Chat(ctx context.Context, model string, messages []Messa
 		})
 	}
 
-	return &Response{Message: msg, StopReason: choice.FinishReason}, nil
+	return &Response{
+		Message:    msg,
+		StopReason: choice.FinishReason,
+		Usage:      Usage{InputTokens: wire.Usage.PromptTokens, OutputTokens: wire.Usage.CompletionTokens},
+	}, nil
 }
diff --git a/main.go b/main.go
index 5c54f21..1032aea 100644
--- a/main.go
+++ b/main.go
@@ -278,6 +278,8 @@ func (m model) runLoop(input string) tea.Cmd {
 				lines = append(lines, fmt.Sprintf("→ %s(%s)", msg.Name, msg.Content))
 			case "tool":
 				lines = append(lines, fmt.Sprintf("  = %s", msg.Content))
+			case "usage":
+				lines = append(lines, "["+msg.Content+"]")
 			case "error":
 				lines = append(lines, "Error: "+msg.Content)
 			}