account for stripCold in cost

This commit is contained in:
Ollie Agent 2026-08-15 12:37:13 +02:00
parent e4c24e8a5f
commit 329edc3218
3 changed files with 25 additions and 7 deletions

View File

@ -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.

View File

@ -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.

View File

@ -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()