diff --git a/cmd/olliesrv/internal/agent/cancellation_test.go b/cmd/olliesrv/internal/agent/cancellation_test.go index 8c77b5f..e93e562 100644 --- a/cmd/olliesrv/internal/agent/cancellation_test.go +++ b/cmd/olliesrv/internal/agent/cancellation_test.go @@ -61,7 +61,7 @@ func TestStripColdCancellation(t *testing.T) { done := make(chan struct{}) go func() { - h.stripCold(ctx, b) + _, _ = h.stripCold(ctx, b) close(done) }() @@ -99,7 +99,7 @@ func TestPreRunCompactionCancellation(t *testing.T) { Backend: b, }, history: newHistory("test"), - log: noOpLogger, + log: noOpLogger, } // Add enough messages to exceed the compact threshold. diff --git a/cmd/olliesrv/internal/agent/history.go b/cmd/olliesrv/internal/agent/history.go index 6f1653c..b7afdc4 100644 --- a/cmd/olliesrv/internal/agent/history.go +++ b/cmd/olliesrv/internal/agent/history.go @@ -381,10 +381,10 @@ func (s *History) estimateTokens() int { // stripCold summarizes large tool-result messages outside the hot tail // using a single batched LLM call. Messages in the last hotTailSize slots // are left verbatim. -func (s *History) stripCold(ctx context.Context, b backend.Backend) { +func (s *History) stripCold(ctx context.Context, b backend.Backend) (backend.Usage, bool) { hot := len(s.messages) - hotTailSize if hot < 0 { - return + return backend.Usage{}, false } // Collect cold-zone tool results exceeding 200 chars. @@ -404,7 +404,7 @@ func (s *History) stripCold(ctx context.Context, b backend.Backend) { } } if len(items) == 0 { - return + return backend.Usage{}, false } // Build a single prompt listing each item. @@ -416,17 +416,27 @@ func (s *History) stripCold(ctx context.Context, b backend.Backend) { ch, err := b.ChatStream(ctx, []backend.Message{{Role: "user", Content: sb.String()}}, nil, backend.GenerationParams{}) if err != nil { - return + return backend.Usage{}, false } var resp strings.Builder + var usage backend.Usage + var reported bool for ev := range ch { if ev.Content != "" { resp.WriteString(ev.Content) } if ev.Done { + usage = ev.Usage + reported = true break } } + if !reported { + usage.InputTokens = len(sb.String()) / 4 + usage.OutputTokens = resp.Len() / 4 + } + estimated := !reported + summaries := make(map[string]string) json.Unmarshal([]byte(resp.String()), &summaries) @@ -437,6 +447,7 @@ func (s *History) stripCold(ctx context.Context, b backend.Backend) { s.messages[it.idx].Content = summary } } + return usage, estimated } // resolveCompactionModel returns the model to use for compaction. diff --git a/cmd/olliesrv/internal/agent/turn.go b/cmd/olliesrv/internal/agent/turn.go index d2097c0..a002d85 100644 --- a/cmd/olliesrv/internal/agent/turn.go +++ b/cmd/olliesrv/internal/agent/turn.go @@ -204,7 +204,11 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string { // Strip cold-zone tool results (one batched LLM call per turn). if ag.history != nil { - ag.history.stripCold(actCtx, ag.runtime.Backend) + usage, estimated := ag.history.stripCold(actCtx, ag.runtime.Backend) + if usage.InputTokens > 0 || usage.OutputTokens > 0 { + ag.history.addUsage(usage, estimated) + ag.notifyChange() + } } // Inject the preamble system message into history before running. @@ -257,6 +261,9 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string { ag.flush() if err != nil { + if ag.history != nil { + ag.history.recordTurnCost(ag.runtime.Backend.Model()) + } // Keep completed work — only remove cancelled tool results. if ag.history != nil { ag.history.removeCancelledToolResults()