diff --git a/pkg/agent/core.go b/pkg/agent/core.go index 8e0fdd4..72e2f8e 100644 --- a/pkg/agent/core.go +++ b/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": diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index b748beb..7b5389a 100644 --- a/pkg/agent/loop.go +++ b/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 diff --git a/pkg/agent/session.go b/pkg/agent/session.go index 6781cc0..358b5a5 100644 --- a/pkg/agent/session.go +++ b/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 {