| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | |
| 6 | "reasonix/internal/billing" |
| 7 | "reasonix/internal/event" |
| 8 | "reasonix/internal/provider" |
| 9 | ) |
| 10 | |
| 11 | // estimateFailedAttemptUsage fills Estimated usage when a body attempt ends |
| 12 | // without a terminal provider usage record, so billing and observational Goal |
| 13 | // usage still include the issued request plus any observed speculative output. |
| 14 | // Non-interrupt failures that already carry usage (e.g. client reasoning limit) |
| 15 | // are left intact. |
| 16 | // |
| 17 | // httpRequests is the SendWithRetry attempt-counter delta for this body attempt. |
| 18 | // When it is 0 and there was no speculative output, the failure was local or |
| 19 | // came from a provider without observable transport accounting; return nil or |
| 20 | // its existing usage rather than inventing billable tokens. |
| 21 | func estimateFailedAttemptUsage(usage *provider.Usage, frozen samplingRequest, result streamedTurn, httpRequests int) *provider.Usage { |
| 22 | if result.err == nil { |
| 23 | return usage |
| 24 | } |
| 25 | // Preserve exact client-side finish reasons that already computed usage. |
| 26 | if usage != nil && usage.FinishReason != "" && usage.FinishReason != "interrupted" { |
| 27 | return usage |
| 28 | } |
| 29 | // A zero-output, non-interrupted failure with no observed HTTP request is a |
| 30 | // local/provider validation failure. It is not a billable sampling attempt. |
| 31 | usage = unmeteredUsage(usage, result, httpRequests) |
| 32 | |
| 33 | preBodyLocal := httpRequests <= 0 && !result.interrupted && |
| 34 | !provider.IsStreamInterrupted(result.err) && !sawSpeculativeSamplingOutput(result) |
| 35 | if preBodyLocal { |
| 36 | if usage != nil && usageTotalTokens(usage) > 0 { |
| 37 | return usage |
| 38 | } |
| 39 | return nil |
| 40 | } |
| 41 | if !provider.IsStreamInterrupted(result.err) && !result.interrupted { |
| 42 | // Auth/cancel/decode/limit paths keep their own accounting. |
| 43 | if usage != nil { |
| 44 | return usage |
| 45 | } |
| 46 | if httpRequests <= 0 { |
| 47 | return nil |
| 48 | } |
| 49 | } |
| 50 | textBytes := len(result.text) |
| 51 | reasoningBytes := len(result.reasoning) |
| 52 | maxArg := result.maxArgChars |
| 53 | for _, call := range result.partialCalls { |
| 54 | if n := len(call.Arguments); n > maxArg { |
| 55 | maxArg = n |
| 56 | } |
| 57 | } |
| 58 | for _, call := range result.calls { |
| 59 | if n := len(call.Arguments); n > maxArg { |
| 60 | maxArg = n |
| 61 | } |
| 62 | } |
| 63 | if usage != nil && !usage.Estimated && usage.TotalTokens > 0 { |
| 64 | return usage |
| 65 | } |
| 66 | finish := "interrupted" |
| 67 | if usage != nil && usage.FinishReason != "" { |
| 68 | finish = usage.FinishReason |
| 69 | } |
| 70 | est := bestEffortStreamUsage(usage, textBytes, reasoningBytes, finish) |
| 71 | if est == nil { |
| 72 | est = &provider.Usage{Estimated: true, FinishReason: finish} |
| 73 | } |
| 74 | if est.PromptTokens <= 0 { |
| 75 | est.PromptTokens = estimateSamplingRequestInputTokens(frozen.req) |
| 76 | est.Estimated = true |
| 77 | } |
| 78 | // Estimated failed attempts without cache split still need Cost() to see |
| 79 | // billable input — Price falls back to PromptTokens only when hit+miss=0. |
| 80 | if est.CacheHitTokens+est.CacheMissTokens == 0 && est.PromptTokens > 0 { |
| 81 | est.CacheMissTokens = est.PromptTokens |
| 82 | } |
| 83 | if maxArg > 0 { |
| 84 | argTokens := (maxArg + 3) / 4 |
| 85 | if est.CompletionTokens < argTokens+estimateTokensFromBytes(textBytes)+estimateTokensFromBytes(reasoningBytes) { |
| 86 | est.CompletionTokens = argTokens + estimateTokensFromBytes(textBytes) + estimateTokensFromBytes(reasoningBytes) |
| 87 | est.Estimated = true |
| 88 | } |
| 89 | } |
| 90 | if minTotal := est.PromptTokens + est.CompletionTokens; est.TotalTokens < minTotal { |
| 91 | est.TotalTokens = minTotal |
| 92 | est.Estimated = true |
| 93 | } |
| 94 | return est |
| 95 | } |
| 96 | |
| 97 | func sawSpeculativeSamplingOutput(result streamedTurn) bool { |
| 98 | return result.text != "" || result.reasoning != "" || result.maxArgChars > 0 || |
| 99 | result.partialToolStarted || len(result.calls) > 0 || len(result.partialCalls) > 0 |
| 100 | } |
| 101 | |
| 102 | // estimateSamplingRequestInputTokens reconstructs a conservative input count |
| 103 | // only when an interrupted attempt closed before terminal provider usage. It is |
| 104 | // accounting telemetry, not request admission: the estimate never changes the |
| 105 | // frozen provider request or imposes a token ceiling. |
| 106 | func estimateSamplingRequestInputTokens(req provider.Request) int { |
| 107 | total := 3 |
| 108 | for _, msg := range provider.ModelMessages(req.Messages) { |
| 109 | total += 4 |
| 110 | total += estimateTextTokens(msg.Content) |
| 111 | total += estimateTextTokens(msg.ReasoningContent) |
| 112 | total += estimateTextTokens(msg.ReasoningSignature) |
| 113 | total += estimateTextTokens(msg.Name) |
| 114 | total += estimateTextTokens(msg.ToolCallID) |
| 115 | for _, image := range msg.Images { |
| 116 | total += estimateTextTokens(image) |
| 117 | } |
| 118 | for _, call := range msg.ToolCalls { |
| 119 | total += 8 + estimateTextTokens(call.ID) + estimateTextTokens(call.Name) + estimateTextTokens(call.Arguments) |
| 120 | } |
| 121 | for _, item := range msg.ResponsesItems { |
| 122 | total += estimateTextTokens(string(item)) |
| 123 | } |
| 124 | for _, search := range msg.ServerSearch { |
| 125 | provider.WalkServerSearchEstimate(search, func(s string) { |
| 126 | total += estimateTextTokens(s) |
| 127 | }) |
| 128 | } |
| 129 | } |
| 130 | for _, schema := range req.Tools { |
| 131 | encoded, _ := json.Marshal(schema) |
| 132 | total += 8 + estimateTextTokens(string(encoded)) |
| 133 | } |
| 134 | return max(total, 1) |
| 135 | } |
| 136 | |
| 137 | // mergeSamplingUsage accumulates billable counters across body attempts. |
| 138 | // PromptTokens is the billable input total (aligned with cache hit+miss). |
| 139 | // ContextPromptTokens is set later by finalizeSamplingUsage from the latest attempt. |
| 140 | func mergeSamplingUsage(acc, attempt *provider.Usage) *provider.Usage { |
| 141 | if attempt == nil { |
| 142 | return acc |
| 143 | } |
| 144 | billableHitMiss := func(u *provider.Usage) (hit, miss int) { |
| 145 | if u == nil { |
| 146 | return 0, 0 |
| 147 | } |
| 148 | if u.CacheHitTokens+u.CacheMissTokens > 0 { |
| 149 | return u.CacheHitTokens, u.CacheMissTokens |
| 150 | } |
| 151 | // No cache split: treat PromptTokens as uncached billable input. |
| 152 | return 0, u.PromptTokens |
| 153 | } |
| 154 | billablePrompt := func(hit, miss, prompt int) int { |
| 155 | if hit+miss > 0 { |
| 156 | return hit + miss |
| 157 | } |
| 158 | return prompt |
| 159 | } |
| 160 | if acc == nil { |
| 161 | merged := *attempt |
| 162 | if merged.RequestCount <= 0 { |
| 163 | merged.RequestCount = 1 |
| 164 | } |
| 165 | hit, miss := billableHitMiss(attempt) |
| 166 | merged.CacheHitTokens = hit |
| 167 | merged.CacheMissTokens = miss |
| 168 | merged.PromptTokens = billablePrompt(hit, miss, attempt.PromptTokens) |
| 169 | return &merged |
| 170 | } |
| 171 | merged := *acc |
| 172 | merged.Unknown = merged.Unknown || attempt.Unknown |
| 173 | // Billable input for Cost: sum hit/miss (prompt when no cache split). |
| 174 | ah, am := billableHitMiss(acc) |
| 175 | bh, bm := billableHitMiss(attempt) |
| 176 | // If acc was previously merged, CacheHit+Miss already holds the sum and |
| 177 | // PromptTokens may still be the first attempt's value — prefer stored sums. |
| 178 | if acc.CacheHitTokens+acc.CacheMissTokens > 0 { |
| 179 | ah, am = acc.CacheHitTokens, acc.CacheMissTokens |
| 180 | } |
| 181 | merged.CacheHitTokens = ah + bh |
| 182 | merged.CacheMissTokens = am + bm |
| 183 | merged.CacheWriteTokens += attempt.CacheWriteTokens |
| 184 | merged.CacheWriteBilledTokens += attempt.CacheWriteBilledTokens |
| 185 | merged.PromptTokens = billablePrompt(merged.CacheHitTokens, merged.CacheMissTokens, 0) |
| 186 | if merged.PromptTokens == 0 { |
| 187 | merged.PromptTokens = acc.PromptTokens + attempt.PromptTokens |
| 188 | } |
| 189 | merged.CompletionTokens += attempt.CompletionTokens |
| 190 | merged.ReasoningTokens += attempt.ReasoningTokens |
| 191 | merged.TotalTokens += usageTotalTokens(attempt) |
| 192 | merged.RequestCount = usageRequestCount(acc) + usageRequestCount(attempt) |
| 193 | if attempt.Estimated { |
| 194 | merged.Estimated = true |
| 195 | } |
| 196 | if attempt.FinishReason != "" { |
| 197 | merged.FinishReason = attempt.FinishReason |
| 198 | } |
| 199 | return &merged |
| 200 | } |
| 201 | |
| 202 | // storeLatestRequestUsage records single-request usage, never a billable aggregate. |
| 203 | func (a *Agent) storeLatestRequestUsage(attempt *provider.Usage) { |
| 204 | if a == nil || attempt == nil { |
| 205 | return |
| 206 | } |
| 207 | // Skip request-only shells with no token shape. |
| 208 | if attempt.PromptTokens <= 0 && attempt.CompletionTokens <= 0 && attempt.TotalTokens <= 0 { |
| 209 | return |
| 210 | } |
| 211 | clone := *attempt |
| 212 | // Keep the per-attempt RequestCount; context calculations do not use it. |
| 213 | a.sess.output.lastUsage.Store(&clone) |
| 214 | a.setPromptTokenCalibrationFromUsage(&clone) |
| 215 | } |
| 216 | |
| 217 | // finalizeSamplingUsage builds the Usage event payload for consumers that |
| 218 | // expect one coherent billable record: |
| 219 | // - PromptTokens / cache hit+miss / Completion / Total / RequestCount: billable aggregate |
| 220 | // - Context* fields: latest attempt only (context gauges + rebind telemetry) |
| 221 | func finalizeSamplingUsage(billable, latest *provider.Usage) *provider.Usage { |
| 222 | if billable == nil && latest == nil { |
| 223 | return nil |
| 224 | } |
| 225 | if billable == nil { |
| 226 | out := *latest |
| 227 | applyLatestContextShape(&out, latest) |
| 228 | return &out |
| 229 | } |
| 230 | out := *billable |
| 231 | if latest != nil { |
| 232 | applyLatestContextShape(&out, latest) |
| 233 | out.FinishReason = latest.FinishReason |
| 234 | } |
| 235 | // Ensure PromptTokens matches billable input (hit+miss) for CLI/ACP/Desktop |
| 236 | // telemetry that requires cache totals to align with PromptTokens. |
| 237 | if hitMiss := out.CacheHitTokens + out.CacheMissTokens; hitMiss > 0 { |
| 238 | out.PromptTokens = hitMiss |
| 239 | } |
| 240 | if out.TotalTokens < out.PromptTokens+out.CompletionTokens { |
| 241 | out.TotalTokens = out.PromptTokens + out.CompletionTokens |
| 242 | } |
| 243 | return &out |
| 244 | } |
| 245 | |
| 246 | // mergeStreamUsage remains for missing-reasoning style single-repair merges that |
| 247 | // need a simple sum. Sampling recovery uses mergeSamplingUsage instead. |
| 248 | func mergeStreamUsage(first, retry *provider.Usage) *provider.Usage { |
| 249 | return mergeSamplingUsage(first, retry) |
| 250 | } |
| 251 | |
| 252 | func usageTotalTokens(u *provider.Usage) int { |
| 253 | if u == nil { |
| 254 | return 0 |
| 255 | } |
| 256 | if u.TotalTokens > 0 { |
| 257 | return u.TotalTokens |
| 258 | } |
| 259 | return u.PromptTokens + u.CompletionTokens |
| 260 | } |
| 261 | |
| 262 | func usageRequestCount(usage *provider.Usage) int { |
| 263 | if usage == nil { |
| 264 | return 0 |
| 265 | } |
| 266 | if usage.RequestCount > 0 { |
| 267 | return usage.RequestCount |
| 268 | } |
| 269 | return 1 |
| 270 | } |
| 271 | |
| 272 | func (a *Agent) emitTurnUsage(usage *provider.Usage, cacheDiagnostics *CacheDiagnostics) *billing.CostQuote { |
| 273 | if usage == nil || (usage.TotalTokens <= 0 && usage.RequestCount <= 0) { |
| 274 | return nil |
| 275 | } |
| 276 | // lastUsage must stay as the latest single-request shape (set during |
| 277 | // sampling recovery). Never overwrite it with a multi-attempt billable |
| 278 | // aggregate — that would inflate ContextSnapshot and compaction decisions. |
| 279 | if a.sess.output.lastUsage.Load() == nil && usage.PromptTokens > 0 { |
| 280 | a.storeLatestRequestUsage(usage) |
| 281 | } |
| 282 | e := event.Event{Kind: event.Usage, ModelRef: a.modelRef, Usage: usage, Pricing: a.svc.pricing, |
| 283 | UsageSource: a.usageSource, |
| 284 | CacheDiagnostics: cacheDiagnostics, |
| 285 | SessionHit: int(a.sess.cacheHit.Load()), SessionMiss: int(a.sess.cacheMiss.Load())} |
| 286 | e.CostQuote = event.EnsureCostQuote(e, a.svc.quoteContext) |
| 287 | a.svc.sink.Emit(e) |
| 288 | return e.CostQuote |
| 289 | } |
| 290 |