mirror of
https://github.com/esengine/DeepSeek-Reasonix.git
synced 2026-11-07 10:45:59 -10:00
feat(agent): observe what a turn actually spends (#8309)
Rounds are the axis we bound and the one nobody reports a failure in. The turn behind #8228 was four hours, 524 requests and 194M tokens; not one of those numbers is a round count, and the same hundred rounds cost minutes or hours depending on what each one read and how long the model thought. This adds the axis that is missing: tokens, requests, wall clock and money, accumulated per turn and reported after every round to sinks that opt into event.RunBudgetSink. Cost goes through provider.Pricing, so a cache hit prices at a fiftieth of a miss instead of reading as a full-price prompt — on a cache-first agent, anything else would misreport the bill by an order of magnitude. Shadow only. Nothing reads a verdict from it and no threshold exists yet: the distribution these samples record is what a threshold should be chosen from, and picking one first would repeat how 100 was picked. Rounds still count when a round's usage never arrived, so the axis never reads cheaper than the turn was. run_loop.go was 3 lines over the file-size ceiling with this, so the sampling-usage accounting it already carried — estimate, merge, finalize, emit — moved to run_usage.go as its own concern. That drops run_loop.go under 800 lines and retires its file-size debt; the baseline diff carries the two findings that moved with the extraction and records one fewer finding overall.
This commit is contained in:
7 files changed
+525
-271
No files matched your search
@@ -0,0 +1,79 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"reasonix/internal/event"
|
||||
"reasonix/internal/provider"
|
||||
)
|
||||
|
||||
// runBudget accumulates what a turn has actually spent. Rounds are a poor proxy
|
||||
// for it: the same hundred of them cost minutes or hours depending on what each
|
||||
// one read and how long the model thought, and the failures worth stopping are
|
||||
// reported in hours and tokens, never in rounds.
|
||||
type runBudget struct {
|
||||
started time.Time
|
||||
rounds int
|
||||
requests int
|
||||
promptTokens int
|
||||
outputTokens int
|
||||
cost float64
|
||||
pricedRounds int
|
||||
unpricedTurns bool
|
||||
}
|
||||
|
||||
// observe folds one round's provider usage into the turn's running total.
|
||||
// A round whose usage never arrived still counts as a round, so the axis never
|
||||
// reads cheaper than the turn actually was.
|
||||
func (b *runBudget) observe(usage *provider.Usage, pricing *provider.Pricing) {
|
||||
b.rounds++
|
||||
if usage == nil {
|
||||
return
|
||||
}
|
||||
b.requests += usageRequestCount(usage)
|
||||
b.promptTokens += usage.PromptTokens
|
||||
b.outputTokens += usage.CompletionTokens
|
||||
if pricing == nil {
|
||||
b.unpricedTurns = true
|
||||
return
|
||||
}
|
||||
b.cost += pricing.Cost(usage)
|
||||
b.pricedRounds++
|
||||
}
|
||||
|
||||
func (b *runBudget) elapsed() time.Duration {
|
||||
if b.started.IsZero() {
|
||||
return 0
|
||||
}
|
||||
return time.Since(b.started)
|
||||
}
|
||||
|
||||
// sample is the per-round shadow reading: counts and money, never content.
|
||||
func (b *runBudget) sample(currency string) event.RunBudgetSample {
|
||||
return event.RunBudgetSample{
|
||||
Rounds: b.rounds,
|
||||
Requests: b.requests,
|
||||
PromptTokens: b.promptTokens,
|
||||
OutputTokens: b.outputTokens,
|
||||
Cost: b.cost,
|
||||
Currency: currency,
|
||||
Priced: !b.unpricedTurns && b.pricedRounds > 0,
|
||||
ElapsedMs: b.elapsed().Milliseconds(),
|
||||
}
|
||||
}
|
||||
|
||||
// observeRunBudget records the turn's spend after a round and reports it to
|
||||
// sinks that opt in. Shadow only: nothing reads the verdict yet, and no
|
||||
// threshold exists to read it against until the recorded distribution says
|
||||
// where one belongs.
|
||||
func (a *Agent) observeRunBudget(state *runLoopState, usage *provider.Usage) {
|
||||
if state == nil {
|
||||
return
|
||||
}
|
||||
state.budget.observe(usage, a.pricing)
|
||||
currency := ""
|
||||
if a.pricing != nil {
|
||||
currency = a.pricing.Symbol()
|
||||
}
|
||||
event.RecordRunBudget(a.sink, state.budget.sample(currency))
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"reasonix/internal/event"
|
||||
"reasonix/internal/provider"
|
||||
"reasonix/internal/tool"
|
||||
)
|
||||
|
||||
// budgetSink opts into the shadow axis; an ordinary sink would receive nothing.
|
||||
type budgetSink struct {
|
||||
event.FuncSink
|
||||
samples []event.RunBudgetSample
|
||||
}
|
||||
|
||||
func newBudgetSink() *budgetSink {
|
||||
s := &budgetSink{}
|
||||
s.FuncSink = event.FuncSink(func(event.Event) {})
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *budgetSink) RecordRunBudget(sample event.RunBudgetSample) {
|
||||
s.samples = append(s.samples, sample)
|
||||
}
|
||||
|
||||
// spendingProvider bills a fixed usage per round and reads one file, so a turn
|
||||
// costs a predictable amount without depending on a real backend.
|
||||
type spendingProvider struct {
|
||||
rounds atomic.Int32
|
||||
max int32
|
||||
}
|
||||
|
||||
func (p *spendingProvider) Name() string { return "spending" }
|
||||
|
||||
func (p *spendingProvider) Stream(context.Context, provider.Request) (<-chan provider.Chunk, error) {
|
||||
round := p.rounds.Add(1)
|
||||
ch := make(chan provider.Chunk, 4)
|
||||
usage := &provider.Usage{
|
||||
PromptTokens: 1000, CompletionTokens: 100, TotalTokens: 1100,
|
||||
CacheHitTokens: 900, CacheMissTokens: 100, RequestCount: 1,
|
||||
}
|
||||
if round > p.max {
|
||||
ch <- provider.Chunk{Type: provider.ChunkText, Text: "Done."}
|
||||
ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: usage}
|
||||
ch <- provider.Chunk{Type: provider.ChunkDone}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
ch <- provider.Chunk{Type: provider.ChunkToolCall, ToolCall: &provider.ToolCall{
|
||||
ID: fmt.Sprintf("call-%d", round),
|
||||
Name: "read_file",
|
||||
Arguments: fmt.Sprintf(`{"path":"pkg%d/file.go"}`, round),
|
||||
}}
|
||||
ch <- provider.Chunk{Type: provider.ChunkUsage, Usage: usage}
|
||||
ch <- provider.Chunk{Type: provider.ChunkDone}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// The axis must read what the turn actually spent, through the real Run loop:
|
||||
// a component-level accumulator that never reaches a sink proves nothing.
|
||||
func TestRunBudgetTracksRealTurnSpend(t *testing.T) {
|
||||
sink := newBudgetSink()
|
||||
reg := tool.NewRegistry()
|
||||
reg.Add(readProbe{})
|
||||
pricing := &provider.Pricing{CacheHit: 0.02, Input: 1, Output: 2, Currency: "CNY"}
|
||||
a := New(&spendingProvider{max: 3}, reg, NewSession("sys"), Options{Pricing: pricing}, sink)
|
||||
|
||||
if err := a.Run(context.Background(), "read a few files"); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
if len(sink.samples) != 4 {
|
||||
t.Fatalf("samples = %d, want one per model round (3 tool rounds + 1 final)", len(sink.samples))
|
||||
}
|
||||
|
||||
last := sink.samples[len(sink.samples)-1]
|
||||
if last.Rounds != 4 || last.Requests != 4 {
|
||||
t.Fatalf("last sample = %+v, want 4 rounds and 4 requests", last)
|
||||
}
|
||||
if last.PromptTokens != 4000 || last.OutputTokens != 400 {
|
||||
t.Fatalf("tokens = prompt %d output %d, want 4000/400", last.PromptTokens, last.OutputTokens)
|
||||
}
|
||||
if !last.Priced || last.Currency != "¥" {
|
||||
t.Fatalf("sample = %+v, want a priced reading in ¥", last)
|
||||
}
|
||||
// Cache hits are 50x cheaper than misses; a turn that bills 900 hits per
|
||||
// round must not read as if all 1000 prompt tokens were misses.
|
||||
wantCost := 4 * (900*0.02 + 100*1 + 100*2) / 1e6
|
||||
if diff := last.Cost - wantCost; diff > 1e-12 || diff < -1e-12 {
|
||||
t.Fatalf("cost = %v, want %v (cache-hit priced)", last.Cost, wantCost)
|
||||
}
|
||||
if last.ElapsedMs < 0 {
|
||||
t.Fatalf("elapsed = %d, want a wall-clock reading", last.ElapsedMs)
|
||||
}
|
||||
}
|
||||
|
||||
// Every round counts even when its usage never arrived, so the axis never
|
||||
// reads cheaper than the turn was.
|
||||
func TestRunBudgetCountsRoundsWithoutUsage(t *testing.T) {
|
||||
var b runBudget
|
||||
b.observe(nil, nil)
|
||||
b.observe(&provider.Usage{PromptTokens: 10, CompletionTokens: 1, RequestCount: 1}, nil)
|
||||
got := b.sample("")
|
||||
if got.Rounds != 2 || got.Requests != 1 || got.PromptTokens != 10 {
|
||||
t.Fatalf("sample = %+v, want 2 rounds / 1 request / 10 prompt tokens", got)
|
||||
}
|
||||
if got.Priced {
|
||||
t.Fatal("an unpriced turn must not report a priced reading")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBudgetIgnoresSinksThatDoNotOptIn(t *testing.T) {
|
||||
plain := event.FuncSink(func(event.Event) {})
|
||||
a := &Agent{sink: plain}
|
||||
state := &runLoopState{}
|
||||
a.observeRunBudget(state, &provider.Usage{PromptTokens: 5, RequestCount: 1})
|
||||
if state.budget.rounds != 1 || state.budget.promptTokens != 5 {
|
||||
t.Fatalf("budget = %+v, want the round still accumulated locally", state.budget)
|
||||
}
|
||||
}
|
||||
+3
-269
@@ -217,6 +217,7 @@ func (a *Agent) beginRunTurn(ctx context.Context, input string) (rawInput string
|
||||
seenTodoProgress: make(map[string]struct{}),
|
||||
executorHandoff: a.executorHandoffGuard && strings.Contains(input, executorHandoffMarker),
|
||||
input: input,
|
||||
budget: runBudget{started: time.Now()},
|
||||
}
|
||||
state.todoProgress, state.trackingTodoProgress = a.canonicalTodoProgress()
|
||||
if a.evidence != nil {
|
||||
@@ -267,6 +268,7 @@ func (a *Agent) runToolLoop(ctx context.Context, state *runLoopState) error {
|
||||
cacheDiagnostics := CompareShape(prevPrefixShape, prefixShape, usage, contentReasons)
|
||||
if err != nil {
|
||||
a.emitTurnUsage(usage, &cacheDiagnostics)
|
||||
a.observeRunBudget(state, usage)
|
||||
if msg, ok := finishReasonMessage(usage); ok {
|
||||
a.sink.Emit(event.Event{Kind: event.Notice, Level: event.LevelWarn, Text: msg})
|
||||
}
|
||||
@@ -279,6 +281,7 @@ func (a *Agent) runToolLoop(ctx context.Context, state *runLoopState) error {
|
||||
a.lastPrefixShape = prefixShape
|
||||
a.haveLastPrefixShape = true
|
||||
a.emitTurnUsage(usage, &cacheDiagnostics)
|
||||
a.observeRunBudget(state, usage)
|
||||
if msg, ok := finishReasonMessage(usage); ok {
|
||||
a.sink.Emit(event.Event{Kind: event.Notice, Level: event.LevelWarn, Text: msg})
|
||||
}
|
||||
@@ -485,275 +488,6 @@ func sleepStreamRetryBackoff(ctx context.Context, attempt int) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// estimateFailedAttemptUsage fills Estimated usage when a body attempt ends
|
||||
// without a terminal provider usage record, so billing and observational Goal
|
||||
// usage still include the issued request plus any observed speculative output.
|
||||
// Non-interrupt failures that already carry usage (e.g. client reasoning limit)
|
||||
// are left intact.
|
||||
//
|
||||
// httpRequests is the SendWithRetry attempt-counter delta for this body attempt.
|
||||
// When it is 0 and there was no speculative output, the failure was local or
|
||||
// came from a provider without observable transport accounting; return nil or
|
||||
// its existing usage rather than inventing billable tokens.
|
||||
func estimateFailedAttemptUsage(usage *provider.Usage, frozen samplingRequest, result streamedTurn, httpRequests int) *provider.Usage {
|
||||
if result.err == nil {
|
||||
return usage
|
||||
}
|
||||
// Preserve exact client-side finish reasons that already computed usage.
|
||||
if usage != nil && usage.FinishReason != "" && usage.FinishReason != "interrupted" {
|
||||
return usage
|
||||
}
|
||||
// A zero-output, non-interrupted failure with no observed HTTP request is a
|
||||
// local/provider validation failure. It is not a billable sampling attempt.
|
||||
preBodyLocal := httpRequests <= 0 && !result.interrupted &&
|
||||
!provider.IsStreamInterrupted(result.err) && !sawSpeculativeSamplingOutput(result)
|
||||
if preBodyLocal {
|
||||
if usage != nil && usageTotalTokens(usage) > 0 {
|
||||
return usage
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if !provider.IsStreamInterrupted(result.err) && !result.interrupted {
|
||||
// Auth/cancel/decode/limit paths keep their own accounting.
|
||||
if usage != nil {
|
||||
return usage
|
||||
}
|
||||
if httpRequests <= 0 {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
textBytes := len(result.text)
|
||||
reasoningBytes := len(result.reasoning)
|
||||
maxArg := result.maxArgChars
|
||||
for _, call := range result.partialCalls {
|
||||
if n := len(call.Arguments); n > maxArg {
|
||||
maxArg = n
|
||||
}
|
||||
}
|
||||
for _, call := range result.calls {
|
||||
if n := len(call.Arguments); n > maxArg {
|
||||
maxArg = n
|
||||
}
|
||||
}
|
||||
if usage != nil && !usage.Estimated && usage.TotalTokens > 0 {
|
||||
return usage
|
||||
}
|
||||
finish := "interrupted"
|
||||
if usage != nil && usage.FinishReason != "" {
|
||||
finish = usage.FinishReason
|
||||
}
|
||||
est := bestEffortStreamUsage(usage, textBytes, reasoningBytes, finish)
|
||||
if est == nil {
|
||||
est = &provider.Usage{Estimated: true, FinishReason: finish}
|
||||
}
|
||||
if est.PromptTokens <= 0 {
|
||||
est.PromptTokens = estimateSamplingRequestInputTokens(frozen.req)
|
||||
est.Estimated = true
|
||||
}
|
||||
// Estimated failed attempts without cache split still need Cost() to see
|
||||
// billable input — Price falls back to PromptTokens only when hit+miss=0.
|
||||
if est.CacheHitTokens+est.CacheMissTokens == 0 && est.PromptTokens > 0 {
|
||||
est.CacheMissTokens = est.PromptTokens
|
||||
}
|
||||
if maxArg > 0 {
|
||||
argTokens := (maxArg + 3) / 4
|
||||
if est.CompletionTokens < argTokens+estimateTokensFromBytes(textBytes)+estimateTokensFromBytes(reasoningBytes) {
|
||||
est.CompletionTokens = argTokens + estimateTokensFromBytes(textBytes) + estimateTokensFromBytes(reasoningBytes)
|
||||
est.Estimated = true
|
||||
}
|
||||
}
|
||||
if minTotal := est.PromptTokens + est.CompletionTokens; est.TotalTokens < minTotal {
|
||||
est.TotalTokens = minTotal
|
||||
est.Estimated = true
|
||||
}
|
||||
return est
|
||||
}
|
||||
|
||||
func sawSpeculativeSamplingOutput(result streamedTurn) bool {
|
||||
return result.text != "" || result.reasoning != "" || result.maxArgChars > 0 ||
|
||||
result.partialToolStarted || len(result.calls) > 0 || len(result.partialCalls) > 0
|
||||
}
|
||||
|
||||
// estimateSamplingRequestInputTokens reconstructs a conservative input count
|
||||
// only when an interrupted attempt closed before terminal provider usage. It is
|
||||
// accounting telemetry, not request admission: the estimate never changes the
|
||||
// frozen provider request or imposes a token ceiling.
|
||||
func estimateSamplingRequestInputTokens(req provider.Request) int {
|
||||
total := 3
|
||||
for _, msg := range provider.ModelMessages(req.Messages) {
|
||||
total += 4
|
||||
total += estimateTextTokens(msg.Content)
|
||||
total += estimateTextTokens(msg.ReasoningContent)
|
||||
total += estimateTextTokens(msg.ReasoningSignature)
|
||||
total += estimateTextTokens(msg.Name)
|
||||
total += estimateTextTokens(msg.ToolCallID)
|
||||
for _, image := range msg.Images {
|
||||
total += estimateTextTokens(image)
|
||||
}
|
||||
for _, call := range msg.ToolCalls {
|
||||
total += 8 + estimateTextTokens(call.ID) + estimateTextTokens(call.Name) + estimateTextTokens(call.Arguments)
|
||||
}
|
||||
for _, item := range msg.ResponsesItems {
|
||||
total += estimateTextTokens(string(item))
|
||||
}
|
||||
}
|
||||
for _, schema := range req.Tools {
|
||||
encoded, _ := json.Marshal(schema)
|
||||
total += 8 + estimateTextTokens(string(encoded))
|
||||
}
|
||||
return max(total, 1)
|
||||
}
|
||||
|
||||
// mergeSamplingUsage accumulates billable counters across body attempts.
|
||||
// PromptTokens is the billable input total (aligned with cache hit+miss).
|
||||
// ContextPromptTokens is set later by finalizeSamplingUsage from the latest attempt.
|
||||
func mergeSamplingUsage(acc, attempt *provider.Usage) *provider.Usage {
|
||||
if attempt == nil {
|
||||
return acc
|
||||
}
|
||||
billableHitMiss := func(u *provider.Usage) (hit, miss int) {
|
||||
if u == nil {
|
||||
return 0, 0
|
||||
}
|
||||
if u.CacheHitTokens+u.CacheMissTokens > 0 {
|
||||
return u.CacheHitTokens, u.CacheMissTokens
|
||||
}
|
||||
// No cache split: treat PromptTokens as uncached billable input.
|
||||
return 0, u.PromptTokens
|
||||
}
|
||||
billablePrompt := func(hit, miss, prompt int) int {
|
||||
if hit+miss > 0 {
|
||||
return hit + miss
|
||||
}
|
||||
return prompt
|
||||
}
|
||||
if acc == nil {
|
||||
merged := *attempt
|
||||
if merged.RequestCount <= 0 {
|
||||
merged.RequestCount = 1
|
||||
}
|
||||
hit, miss := billableHitMiss(attempt)
|
||||
merged.CacheHitTokens = hit
|
||||
merged.CacheMissTokens = miss
|
||||
merged.PromptTokens = billablePrompt(hit, miss, attempt.PromptTokens)
|
||||
return &merged
|
||||
}
|
||||
merged := *acc
|
||||
// Billable input for Cost: sum hit/miss (prompt when no cache split).
|
||||
ah, am := billableHitMiss(acc)
|
||||
bh, bm := billableHitMiss(attempt)
|
||||
// If acc was previously merged, CacheHit+Miss already holds the sum and
|
||||
// PromptTokens may still be the first attempt's value — prefer stored sums.
|
||||
if acc.CacheHitTokens+acc.CacheMissTokens > 0 {
|
||||
ah, am = acc.CacheHitTokens, acc.CacheMissTokens
|
||||
}
|
||||
merged.CacheHitTokens = ah + bh
|
||||
merged.CacheMissTokens = am + bm
|
||||
merged.CacheWriteTokens += attempt.CacheWriteTokens
|
||||
merged.CacheWriteBilledTokens += attempt.CacheWriteBilledTokens
|
||||
merged.PromptTokens = billablePrompt(merged.CacheHitTokens, merged.CacheMissTokens, 0)
|
||||
if merged.PromptTokens == 0 {
|
||||
merged.PromptTokens = acc.PromptTokens + attempt.PromptTokens
|
||||
}
|
||||
merged.CompletionTokens += attempt.CompletionTokens
|
||||
merged.ReasoningTokens += attempt.ReasoningTokens
|
||||
merged.TotalTokens += usageTotalTokens(attempt)
|
||||
merged.RequestCount = usageRequestCount(acc) + usageRequestCount(attempt)
|
||||
if attempt.Estimated {
|
||||
merged.Estimated = true
|
||||
}
|
||||
if attempt.FinishReason != "" {
|
||||
merged.FinishReason = attempt.FinishReason
|
||||
}
|
||||
return &merged
|
||||
}
|
||||
|
||||
// storeLatestRequestUsage records single-request usage, never a billable aggregate.
|
||||
func (a *Agent) storeLatestRequestUsage(attempt *provider.Usage) {
|
||||
if a == nil || attempt == nil {
|
||||
return
|
||||
}
|
||||
// Skip request-only shells with no token shape.
|
||||
if attempt.PromptTokens <= 0 && attempt.CompletionTokens <= 0 && attempt.TotalTokens <= 0 {
|
||||
return
|
||||
}
|
||||
clone := *attempt
|
||||
// Keep the per-attempt RequestCount; context calculations do not use it.
|
||||
a.lastUsage.Store(&clone)
|
||||
a.setPromptTokenCalibrationFromUsage(&clone)
|
||||
}
|
||||
|
||||
// finalizeSamplingUsage builds the Usage event payload for consumers that
|
||||
// expect one coherent billable record:
|
||||
// - PromptTokens / cache hit+miss / Completion / Total / RequestCount: billable aggregate
|
||||
// - Context* fields: latest attempt only (context gauges + rebind telemetry)
|
||||
func finalizeSamplingUsage(billable, latest *provider.Usage) *provider.Usage {
|
||||
if billable == nil && latest == nil {
|
||||
return nil
|
||||
}
|
||||
if billable == nil {
|
||||
out := *latest
|
||||
applyLatestContextShape(&out, latest)
|
||||
return &out
|
||||
}
|
||||
out := *billable
|
||||
if latest != nil {
|
||||
applyLatestContextShape(&out, latest)
|
||||
out.FinishReason = latest.FinishReason
|
||||
}
|
||||
// Ensure PromptTokens matches billable input (hit+miss) for CLI/ACP/Desktop
|
||||
// telemetry that requires cache totals to align with PromptTokens.
|
||||
if hitMiss := out.CacheHitTokens + out.CacheMissTokens; hitMiss > 0 {
|
||||
out.PromptTokens = hitMiss
|
||||
}
|
||||
if out.TotalTokens < out.PromptTokens+out.CompletionTokens {
|
||||
out.TotalTokens = out.PromptTokens + out.CompletionTokens
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// mergeStreamUsage remains for missing-reasoning style single-repair merges that
|
||||
// need a simple sum. Sampling recovery uses mergeSamplingUsage instead.
|
||||
func mergeStreamUsage(first, retry *provider.Usage) *provider.Usage {
|
||||
return mergeSamplingUsage(first, retry)
|
||||
}
|
||||
|
||||
func usageTotalTokens(u *provider.Usage) int {
|
||||
if u == nil {
|
||||
return 0
|
||||
}
|
||||
if u.TotalTokens > 0 {
|
||||
return u.TotalTokens
|
||||
}
|
||||
return u.PromptTokens + u.CompletionTokens
|
||||
}
|
||||
|
||||
func usageRequestCount(usage *provider.Usage) int {
|
||||
if usage == nil {
|
||||
return 0
|
||||
}
|
||||
if usage.RequestCount > 0 {
|
||||
return usage.RequestCount
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func (a *Agent) emitTurnUsage(usage *provider.Usage, cacheDiagnostics *CacheDiagnostics) {
|
||||
if usage == nil || (usage.TotalTokens <= 0 && usage.RequestCount <= 0) {
|
||||
return
|
||||
}
|
||||
// lastUsage must stay as the latest single-request shape (set during
|
||||
// sampling recovery). Never overwrite it with a multi-attempt billable
|
||||
// aggregate — that would inflate ContextSnapshot and compaction decisions.
|
||||
if a.lastUsage.Load() == nil && usage.PromptTokens > 0 {
|
||||
a.storeLatestRequestUsage(usage)
|
||||
}
|
||||
a.sink.Emit(event.Event{Kind: event.Usage, ModelRef: a.modelRef, Usage: usage, Pricing: a.pricing,
|
||||
UsageSource: a.usageSource,
|
||||
CacheDiagnostics: cacheDiagnostics,
|
||||
SessionHit: int(a.sessCacheHit.Load()), SessionMiss: int(a.sessCacheMiss.Load())})
|
||||
}
|
||||
|
||||
// handleFinalResponse processes a no-tool assistant turn: recovery pause,
|
||||
// readiness retry, empty final retry, executor handoff nudge, steer drain, and
|
||||
// final compaction. cont=true continues the tool loop; cont=false returns err
|
||||
|
||||
@@ -26,4 +26,8 @@ type runLoopState struct {
|
||||
executorHandoff bool
|
||||
input string
|
||||
workDurationMs func() int64
|
||||
|
||||
// budget is the turn's spend axis: tokens, money, wall clock. Observed
|
||||
// every round, read by nothing yet.
|
||||
budget runBudget
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"reasonix/internal/event"
|
||||
"reasonix/internal/provider"
|
||||
)
|
||||
|
||||
// estimateFailedAttemptUsage fills Estimated usage when a body attempt ends
|
||||
// without a terminal provider usage record, so billing and observational Goal
|
||||
// usage still include the issued request plus any observed speculative output.
|
||||
// Non-interrupt failures that already carry usage (e.g. client reasoning limit)
|
||||
// are left intact.
|
||||
//
|
||||
// httpRequests is the SendWithRetry attempt-counter delta for this body attempt.
|
||||
// When it is 0 and there was no speculative output, the failure was local or
|
||||
// came from a provider without observable transport accounting; return nil or
|
||||
// its existing usage rather than inventing billable tokens.
|
||||
func estimateFailedAttemptUsage(usage *provider.Usage, frozen samplingRequest, result streamedTurn, httpRequests int) *provider.Usage {
|
||||
if result.err == nil {
|
||||
return usage
|
||||
}
|
||||
// Preserve exact client-side finish reasons that already computed usage.
|
||||
if usage != nil && usage.FinishReason != "" && usage.FinishReason != "interrupted" {
|
||||
return usage
|
||||
}
|
||||
// A zero-output, non-interrupted failure with no observed HTTP request is a
|
||||
// local/provider validation failure. It is not a billable sampling attempt.
|
||||
preBodyLocal := httpRequests <= 0 && !result.interrupted &&
|
||||
!provider.IsStreamInterrupted(result.err) && !sawSpeculativeSamplingOutput(result)
|
||||
if preBodyLocal {
|
||||
if usage != nil && usageTotalTokens(usage) > 0 {
|
||||
return usage
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if !provider.IsStreamInterrupted(result.err) && !result.interrupted {
|
||||
// Auth/cancel/decode/limit paths keep their own accounting.
|
||||
if usage != nil {
|
||||
return usage
|
||||
}
|
||||
if httpRequests <= 0 {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
textBytes := len(result.text)
|
||||
reasoningBytes := len(result.reasoning)
|
||||
maxArg := result.maxArgChars
|
||||
for _, call := range result.partialCalls {
|
||||
if n := len(call.Arguments); n > maxArg {
|
||||
maxArg = n
|
||||
}
|
||||
}
|
||||
for _, call := range result.calls {
|
||||
if n := len(call.Arguments); n > maxArg {
|
||||
maxArg = n
|
||||
}
|
||||
}
|
||||
if usage != nil && !usage.Estimated && usage.TotalTokens > 0 {
|
||||
return usage
|
||||
}
|
||||
finish := "interrupted"
|
||||
if usage != nil && usage.FinishReason != "" {
|
||||
finish = usage.FinishReason
|
||||
}
|
||||
est := bestEffortStreamUsage(usage, textBytes, reasoningBytes, finish)
|
||||
if est == nil {
|
||||
est = &provider.Usage{Estimated: true, FinishReason: finish}
|
||||
}
|
||||
if est.PromptTokens <= 0 {
|
||||
est.PromptTokens = estimateSamplingRequestInputTokens(frozen.req)
|
||||
est.Estimated = true
|
||||
}
|
||||
// Estimated failed attempts without cache split still need Cost() to see
|
||||
// billable input — Price falls back to PromptTokens only when hit+miss=0.
|
||||
if est.CacheHitTokens+est.CacheMissTokens == 0 && est.PromptTokens > 0 {
|
||||
est.CacheMissTokens = est.PromptTokens
|
||||
}
|
||||
if maxArg > 0 {
|
||||
argTokens := (maxArg + 3) / 4
|
||||
if est.CompletionTokens < argTokens+estimateTokensFromBytes(textBytes)+estimateTokensFromBytes(reasoningBytes) {
|
||||
est.CompletionTokens = argTokens + estimateTokensFromBytes(textBytes) + estimateTokensFromBytes(reasoningBytes)
|
||||
est.Estimated = true
|
||||
}
|
||||
}
|
||||
if minTotal := est.PromptTokens + est.CompletionTokens; est.TotalTokens < minTotal {
|
||||
est.TotalTokens = minTotal
|
||||
est.Estimated = true
|
||||
}
|
||||
return est
|
||||
}
|
||||
|
||||
func sawSpeculativeSamplingOutput(result streamedTurn) bool {
|
||||
return result.text != "" || result.reasoning != "" || result.maxArgChars > 0 ||
|
||||
result.partialToolStarted || len(result.calls) > 0 || len(result.partialCalls) > 0
|
||||
}
|
||||
|
||||
// estimateSamplingRequestInputTokens reconstructs a conservative input count
|
||||
// only when an interrupted attempt closed before terminal provider usage. It is
|
||||
// accounting telemetry, not request admission: the estimate never changes the
|
||||
// frozen provider request or imposes a token ceiling.
|
||||
func estimateSamplingRequestInputTokens(req provider.Request) int {
|
||||
total := 3
|
||||
for _, msg := range provider.ModelMessages(req.Messages) {
|
||||
total += 4
|
||||
total += estimateTextTokens(msg.Content)
|
||||
total += estimateTextTokens(msg.ReasoningContent)
|
||||
total += estimateTextTokens(msg.ReasoningSignature)
|
||||
total += estimateTextTokens(msg.Name)
|
||||
total += estimateTextTokens(msg.ToolCallID)
|
||||
for _, image := range msg.Images {
|
||||
total += estimateTextTokens(image)
|
||||
}
|
||||
for _, call := range msg.ToolCalls {
|
||||
total += 8 + estimateTextTokens(call.ID) + estimateTextTokens(call.Name) + estimateTextTokens(call.Arguments)
|
||||
}
|
||||
for _, item := range msg.ResponsesItems {
|
||||
total += estimateTextTokens(string(item))
|
||||
}
|
||||
}
|
||||
for _, schema := range req.Tools {
|
||||
encoded, _ := json.Marshal(schema)
|
||||
total += 8 + estimateTextTokens(string(encoded))
|
||||
}
|
||||
return max(total, 1)
|
||||
}
|
||||
|
||||
// mergeSamplingUsage accumulates billable counters across body attempts.
|
||||
// PromptTokens is the billable input total (aligned with cache hit+miss).
|
||||
// ContextPromptTokens is set later by finalizeSamplingUsage from the latest attempt.
|
||||
func mergeSamplingUsage(acc, attempt *provider.Usage) *provider.Usage {
|
||||
if attempt == nil {
|
||||
return acc
|
||||
}
|
||||
billableHitMiss := func(u *provider.Usage) (hit, miss int) {
|
||||
if u == nil {
|
||||
return 0, 0
|
||||
}
|
||||
if u.CacheHitTokens+u.CacheMissTokens > 0 {
|
||||
return u.CacheHitTokens, u.CacheMissTokens
|
||||
}
|
||||
// No cache split: treat PromptTokens as uncached billable input.
|
||||
return 0, u.PromptTokens
|
||||
}
|
||||
billablePrompt := func(hit, miss, prompt int) int {
|
||||
if hit+miss > 0 {
|
||||
return hit + miss
|
||||
}
|
||||
return prompt
|
||||
}
|
||||
if acc == nil {
|
||||
merged := *attempt
|
||||
if merged.RequestCount <= 0 {
|
||||
merged.RequestCount = 1
|
||||
}
|
||||
hit, miss := billableHitMiss(attempt)
|
||||
merged.CacheHitTokens = hit
|
||||
merged.CacheMissTokens = miss
|
||||
merged.PromptTokens = billablePrompt(hit, miss, attempt.PromptTokens)
|
||||
return &merged
|
||||
}
|
||||
merged := *acc
|
||||
// Billable input for Cost: sum hit/miss (prompt when no cache split).
|
||||
ah, am := billableHitMiss(acc)
|
||||
bh, bm := billableHitMiss(attempt)
|
||||
// If acc was previously merged, CacheHit+Miss already holds the sum and
|
||||
// PromptTokens may still be the first attempt's value — prefer stored sums.
|
||||
if acc.CacheHitTokens+acc.CacheMissTokens > 0 {
|
||||
ah, am = acc.CacheHitTokens, acc.CacheMissTokens
|
||||
}
|
||||
merged.CacheHitTokens = ah + bh
|
||||
merged.CacheMissTokens = am + bm
|
||||
merged.CacheWriteTokens += attempt.CacheWriteTokens
|
||||
merged.CacheWriteBilledTokens += attempt.CacheWriteBilledTokens
|
||||
merged.PromptTokens = billablePrompt(merged.CacheHitTokens, merged.CacheMissTokens, 0)
|
||||
if merged.PromptTokens == 0 {
|
||||
merged.PromptTokens = acc.PromptTokens + attempt.PromptTokens
|
||||
}
|
||||
merged.CompletionTokens += attempt.CompletionTokens
|
||||
merged.ReasoningTokens += attempt.ReasoningTokens
|
||||
merged.TotalTokens += usageTotalTokens(attempt)
|
||||
merged.RequestCount = usageRequestCount(acc) + usageRequestCount(attempt)
|
||||
if attempt.Estimated {
|
||||
merged.Estimated = true
|
||||
}
|
||||
if attempt.FinishReason != "" {
|
||||
merged.FinishReason = attempt.FinishReason
|
||||
}
|
||||
return &merged
|
||||
}
|
||||
|
||||
// storeLatestRequestUsage records single-request usage, never a billable aggregate.
|
||||
func (a *Agent) storeLatestRequestUsage(attempt *provider.Usage) {
|
||||
if a == nil || attempt == nil {
|
||||
return
|
||||
}
|
||||
// Skip request-only shells with no token shape.
|
||||
if attempt.PromptTokens <= 0 && attempt.CompletionTokens <= 0 && attempt.TotalTokens <= 0 {
|
||||
return
|
||||
}
|
||||
clone := *attempt
|
||||
// Keep the per-attempt RequestCount; context calculations do not use it.
|
||||
a.lastUsage.Store(&clone)
|
||||
a.setPromptTokenCalibrationFromUsage(&clone)
|
||||
}
|
||||
|
||||
// finalizeSamplingUsage builds the Usage event payload for consumers that
|
||||
// expect one coherent billable record:
|
||||
// - PromptTokens / cache hit+miss / Completion / Total / RequestCount: billable aggregate
|
||||
// - Context* fields: latest attempt only (context gauges + rebind telemetry)
|
||||
func finalizeSamplingUsage(billable, latest *provider.Usage) *provider.Usage {
|
||||
if billable == nil && latest == nil {
|
||||
return nil
|
||||
}
|
||||
if billable == nil {
|
||||
out := *latest
|
||||
applyLatestContextShape(&out, latest)
|
||||
return &out
|
||||
}
|
||||
out := *billable
|
||||
if latest != nil {
|
||||
applyLatestContextShape(&out, latest)
|
||||
out.FinishReason = latest.FinishReason
|
||||
}
|
||||
// Ensure PromptTokens matches billable input (hit+miss) for CLI/ACP/Desktop
|
||||
// telemetry that requires cache totals to align with PromptTokens.
|
||||
if hitMiss := out.CacheHitTokens + out.CacheMissTokens; hitMiss > 0 {
|
||||
out.PromptTokens = hitMiss
|
||||
}
|
||||
if out.TotalTokens < out.PromptTokens+out.CompletionTokens {
|
||||
out.TotalTokens = out.PromptTokens + out.CompletionTokens
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// mergeStreamUsage remains for missing-reasoning style single-repair merges that
|
||||
// need a simple sum. Sampling recovery uses mergeSamplingUsage instead.
|
||||
func mergeStreamUsage(first, retry *provider.Usage) *provider.Usage {
|
||||
return mergeSamplingUsage(first, retry)
|
||||
}
|
||||
|
||||
func usageTotalTokens(u *provider.Usage) int {
|
||||
if u == nil {
|
||||
return 0
|
||||
}
|
||||
if u.TotalTokens > 0 {
|
||||
return u.TotalTokens
|
||||
}
|
||||
return u.PromptTokens + u.CompletionTokens
|
||||
}
|
||||
|
||||
func usageRequestCount(usage *provider.Usage) int {
|
||||
if usage == nil {
|
||||
return 0
|
||||
}
|
||||
if usage.RequestCount > 0 {
|
||||
return usage.RequestCount
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func (a *Agent) emitTurnUsage(usage *provider.Usage, cacheDiagnostics *CacheDiagnostics) {
|
||||
if usage == nil || (usage.TotalTokens <= 0 && usage.RequestCount <= 0) {
|
||||
return
|
||||
}
|
||||
// lastUsage must stay as the latest single-request shape (set during
|
||||
// sampling recovery). Never overwrite it with a multi-attempt billable
|
||||
// aggregate — that would inflate ContextSnapshot and compaction decisions.
|
||||
if a.lastUsage.Load() == nil && usage.PromptTokens > 0 {
|
||||
a.storeLatestRequestUsage(usage)
|
||||
}
|
||||
a.sink.Emit(event.Event{Kind: event.Usage, ModelRef: a.modelRef, Usage: usage, Pricing: a.pricing,
|
||||
UsageSource: a.usageSource,
|
||||
CacheDiagnostics: cacheDiagnostics,
|
||||
SessionHit: int(a.sessCacheHit.Load()), SessionMiss: int(a.sessCacheMiss.Load())})
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package event
|
||||
|
||||
import "reasonix/internal/nilutil"
|
||||
|
||||
// RunBudgetSample is one turn's accumulated spend after a round: counts and
|
||||
// money, never content. Priced is false when no price table covered the turn,
|
||||
// so a reader can tell a genuinely free turn from an unpriced one.
|
||||
type RunBudgetSample struct {
|
||||
Rounds int `json:"rounds"`
|
||||
Requests int `json:"requests"`
|
||||
PromptTokens int `json:"promptTokens"`
|
||||
OutputTokens int `json:"outputTokens"`
|
||||
Cost float64 `json:"cost"`
|
||||
Currency string `json:"currency,omitempty"`
|
||||
Priced bool `json:"priced"`
|
||||
ElapsedMs int64 `json:"elapsedMs"`
|
||||
}
|
||||
|
||||
// RunBudgetSink is an optional sink capability for the per-round spend axis.
|
||||
// Shadow means observed, not enforced: no threshold reads these yet.
|
||||
type RunBudgetSink interface {
|
||||
RecordRunBudget(RunBudgetSample)
|
||||
}
|
||||
|
||||
// RecordRunBudget forwards a turn's spend reading only to sinks that opt in.
|
||||
// Ordinary UI sinks receive nothing.
|
||||
func RecordRunBudget(s Sink, sample RunBudgetSample) {
|
||||
if nilutil.IsNil(s) {
|
||||
return
|
||||
}
|
||||
if rb, ok := s.(RunBudgetSink); ok {
|
||||
rb.RecordRunBudget(sample)
|
||||
}
|
||||
}
|
||||
@@ -571,9 +571,11 @@
|
||||
"essay": 1
|
||||
},
|
||||
"internal/agent/run_loop.go": {
|
||||
"essay": 27
|
||||
},
|
||||
"internal/agent/run_usage.go": {
|
||||
"complexity": 1,
|
||||
"essay": 32,
|
||||
"file-size": 179
|
||||
"essay": 5
|
||||
},
|
||||
"internal/agent/save.go": {
|
||||
"complexity": 49,
|
||||
|
||||
Reference in new issue
Block a user