usage: track responses, estimate tokens when backend doesn't report

pkg/agent/core.go

@@ -446,8 +446,13 @@ func (s *agentCore) Usage() string {
 	if s.session == nil {
 		return "no active session"
 	}
-	return fmt.Sprintf("%d in, %d out, %d requests",
-		s.session.TotalInputTokens, s.session.TotalOutputTokens, s.session.TotalRequests)
+	str := fmt.Sprintf("%d in, %d out, %d requests, %d responses",
+		s.session.TotalInputTokens, s.session.TotalOutputTokens,
+		s.session.TotalRequests, s.session.TotalResponses)
+	if s.session.Estimated {
+		str += " [estimated]"
+	}
+	return str
 }
 
 func (s *agentCore) ListModels() string {
@@ -570,9 +575,9 @@ func (s *agentCore) Submit(ctx context.Context, input string, handler EventHandl
 			s.mu.Unlock()
 		}
 		if ev.Role == "usage" && s.session != nil {
-			var in, out int
-			fmt.Sscanf(ev.Content, "%d %d", &in, &out)
-			s.session.addUsage(backend.Usage{InputTokens: in, OutputTokens: out})
+			var in, out, est int
+			fmt.Sscanf(ev.Content, "%d %d %d", &in, &out, &est)
+			s.session.addUsage(backend.Usage{InputTokens: in, OutputTokens: out}, est != 0)
 		}
 		handler(ev)
 	}
@@ -599,6 +604,10 @@ func (s *agentCore) Submit(ctx context.Context, input string, handler EventHandl
 	actCancel(nil)
 	s.currentAction.CompareAndSwap(handle, nil)
 
+	if s.session != nil {
+		s.session.addResponse()
+	}
+
 	s.mu.Lock()
 	s.reply = replyBuf.String()
 	s.state = "idle"
@@ -827,19 +836,20 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve
 			handler(infoEvent("no active session"))
 			return true
 		}
-		total := s.session.TotalInputTokens + s.session.TotalOutputTokens
 		ctxLen := s.loopcfg.Backend.ContextLength(ctx)
 		if ctxLen <= 0 {
 			ctxLen = defaultContextLength
 		}
-		estimated := s.session.estimateTokens()
-		pct := estimated * 100 / ctxLen
-		handler(infoEvent(fmt.Sprintf("~%d / %d tokens (%d%%) | %d in, %d out, %d requests",
-			estimated, ctxLen, pct,
-			s.session.TotalInputTokens, s.session.TotalOutputTokens, s.session.TotalRequests)))
-		if total == 0 {
-			handler(infoEvent("(backend does not report token usage)"))
+		ctxEstimated := s.session.estimateTokens()
+		pct := ctxEstimated * 100 / ctxLen
+		usageStr := fmt.Sprintf("~%d / %d tokens (%d%%) | %d in, %d out, %d requests, %d responses",
+			ctxEstimated, ctxLen, pct,
+			s.session.TotalInputTokens, s.session.TotalOutputTokens,
+			s.session.TotalRequests, s.session.TotalResponses)
+		if s.session.Estimated {
+			usageStr += " [estimated]"
 		}
+		handler(infoEvent(usageStr))
 		return true
 
 	case "/history":

pkg/agent/loop.go

@@ -76,7 +76,20 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
 				if ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0 {
 					emit(cfg, Event{
 						Role:    "usage",
-						Content: fmt.Sprintf("%d %d", ev.Usage.InputTokens, ev.Usage.OutputTokens),
+						Content: fmt.Sprintf("%d %d 0", ev.Usage.InputTokens, ev.Usage.OutputTokens),
+					})
+				} else {
+					// Backend didn't report usage; estimate from content.
+					inChars := 0
+					for _, m := range history {
+						inChars += len(m.Content)
+						for _, tc := range m.ToolCalls {
+							inChars += len(tc.Name) + len(tc.Arguments)
+						}
+					}
+					emit(cfg, Event{
+						Role:    "usage",
+						Content: fmt.Sprintf("%d %d 1", inChars/4, content.Len()/4),
 					})
 				}
 				break

pkg/agent/session.go

@@ -73,6 +73,8 @@ type Session struct {
 	TotalInputTokens  int
 	TotalOutputTokens int
 	TotalRequests     int
+	TotalResponses    int
+	Estimated         bool // true if any usage was estimated rather than reported by the backend
 }
 
 // newSession creates a new Session with an initial user message.
@@ -90,10 +92,17 @@ func (s *Session) history() []backend.Message {
 
 func (s *Session) isComplete() bool { return s.complete }
 
-func (s *Session) addUsage(u backend.Usage) {
+func (s *Session) addUsage(u backend.Usage, estimated bool) {
 	s.TotalInputTokens += u.InputTokens
 	s.TotalOutputTokens += u.OutputTokens
 	s.TotalRequests++
+	if estimated {
+		s.Estimated = true
+	}
+}
+
+func (s *Session) addResponse() {
+	s.TotalResponses++
 }
 
 func (s *Session) update(assistant backend.Message, results []toolResult) error {