| 1 | package event |
| 2 | |
| 3 | import ( |
| 4 | "strings" |
| 5 | "sync" |
| 6 | "time" |
| 7 | |
| 8 | "reasonix/internal/billing" |
| 9 | "reasonix/internal/provider" |
| 10 | ) |
| 11 | |
| 12 | // QuoteContext supplies an explicit display request for the CostQuote |
| 13 | // middleware. Empty means auto and keeps the price-book currency. |
| 14 | type QuoteContext struct { |
| 15 | mu sync.RWMutex |
| 16 | DisplayCurrency string |
| 17 | DisplayRequest billing.DisplayRequest |
| 18 | // Now overrides the clock in tests. |
| 19 | Now func() time.Time |
| 20 | // BillingModeForModel resolves provider-owned billing semantics from the |
| 21 | // immutable boot config. It keeps billing_mode authoritative even when a |
| 22 | // custom provider name does not contain a token-plan heuristic. |
| 23 | BillingModeForModel func(modelRef string) string |
| 24 | // PricingContextForModel resolves trusted catalog/schedule identity from the |
| 25 | // immutable boot config. It supersedes BillingModeForModel when present. |
| 26 | PricingContextForModel func(modelRef string) billing.PricingContext |
| 27 | } |
| 28 | |
| 29 | // SetDisplay updates the resolved display currency used for Selected. |
| 30 | func (c *QuoteContext) SetDisplay(currency string) { |
| 31 | if c == nil { |
| 32 | return |
| 33 | } |
| 34 | c.mu.Lock() |
| 35 | c.DisplayCurrency = billing.NormalizeCurrency(currency) |
| 36 | c.DisplayRequest = billing.DisplayRequest{Currency: c.DisplayCurrency, Source: billing.DisplaySourceExplicit} |
| 37 | c.mu.Unlock() |
| 38 | } |
| 39 | |
| 40 | func (c *QuoteContext) SetDisplayRequest(request billing.DisplayRequest) { |
| 41 | if c == nil { |
| 42 | return |
| 43 | } |
| 44 | request.Currency = billing.NormalizeCurrency(request.Currency) |
| 45 | if request.Source == "" { |
| 46 | request.Source = billing.DisplaySourceAuto |
| 47 | } |
| 48 | c.mu.Lock() |
| 49 | c.DisplayRequest = request |
| 50 | c.DisplayCurrency = request.Currency |
| 51 | c.mu.Unlock() |
| 52 | } |
| 53 | |
| 54 | func (c *QuoteContext) snapshot() (request billing.DisplayRequest, now time.Time) { |
| 55 | if c == nil { |
| 56 | return billing.DisplayRequest{Source: billing.DisplaySourceAuto}, time.Now().UTC() |
| 57 | } |
| 58 | c.mu.RLock() |
| 59 | request = c.DisplayRequest |
| 60 | if request.Currency == "" && c.DisplayCurrency != "" { |
| 61 | request.Currency = c.DisplayCurrency |
| 62 | } |
| 63 | if request.Source == "" { |
| 64 | if request.Currency != "" { |
| 65 | request.Source = billing.DisplaySourceExplicit |
| 66 | } else { |
| 67 | request.Source = billing.DisplaySourceAuto |
| 68 | } |
| 69 | } |
| 70 | nowFn := c.Now |
| 71 | c.mu.RUnlock() |
| 72 | if nowFn != nil { |
| 73 | now = nowFn() |
| 74 | } else { |
| 75 | now = time.Now().UTC() |
| 76 | } |
| 77 | return request, now |
| 78 | } |
| 79 | |
| 80 | func (c *QuoteContext) billingMode(modelRef string) string { |
| 81 | if c == nil { |
| 82 | return "" |
| 83 | } |
| 84 | c.mu.RLock() |
| 85 | resolve := c.BillingModeForModel |
| 86 | c.mu.RUnlock() |
| 87 | if resolve == nil { |
| 88 | return "" |
| 89 | } |
| 90 | return strings.TrimSpace(resolve(modelRef)) |
| 91 | } |
| 92 | |
| 93 | func (c *QuoteContext) pricingContext(modelRef string) billing.PricingContext { |
| 94 | if c == nil { |
| 95 | return billing.PricingContext{} |
| 96 | } |
| 97 | c.mu.RLock() |
| 98 | resolve := c.PricingContextForModel |
| 99 | c.mu.RUnlock() |
| 100 | if resolve != nil { |
| 101 | return resolve(modelRef) |
| 102 | } |
| 103 | return billing.PricingContext{BillingMode: c.billingMode(modelRef)} |
| 104 | } |
| 105 | |
| 106 | // CostQuoteSink fills CostQuote on Usage events before forwarding. Frontends |
| 107 | // must consume e.CostQuote and must not call Pricing.Cost for aggregation. |
| 108 | // The embedded forwarder carries the optional audit capabilities past this |
| 109 | // link: without it every recorder below a quoting sink — trajectory and stats |
| 110 | // among them — silently receives nothing. |
| 111 | type CostQuoteSink struct { |
| 112 | AuditForwarder |
| 113 | Ctx *QuoteContext |
| 114 | } |
| 115 | |
| 116 | var _ OptionalSinkCapabilities = (*CostQuoteSink)(nil) |
| 117 | |
| 118 | // NewCostQuoteSink wraps inner with quoting. A nil ctx still quotes the |
| 119 | // original price-book currency. |
| 120 | func NewCostQuoteSink(inner Sink, ctx *QuoteContext) *CostQuoteSink { |
| 121 | if ctx == nil { |
| 122 | ctx = &QuoteContext{} |
| 123 | } |
| 124 | return &CostQuoteSink{AuditForwarder: AuditForwarder{Inner: inner}, Ctx: ctx} |
| 125 | } |
| 126 | |
| 127 | func (s *CostQuoteSink) Emit(e Event) { |
| 128 | if s == nil { |
| 129 | return |
| 130 | } |
| 131 | if e.Kind == Usage && e.Usage != nil && e.CostQuote == nil { |
| 132 | e.CostQuote = EnsureCostQuote(e, s.Ctx) |
| 133 | } |
| 134 | if s.Inner != nil { |
| 135 | s.Inner.Emit(e) |
| 136 | } |
| 137 | } |
| 138 | |
| 139 | // EnsureCostQuote builds a CostQuote for an event when missing. |
| 140 | func EnsureCostQuote(e Event, ctx *QuoteContext) *billing.CostQuote { |
| 141 | if e.Usage == nil { |
| 142 | return nil |
| 143 | } |
| 144 | display, now := ctx.snapshot() |
| 145 | if e.Usage.Unknown { |
| 146 | known := *e.Usage |
| 147 | known.Unknown = false |
| 148 | e.Usage = &known |
| 149 | q := EnsureCostQuote(e, ctx) |
| 150 | q.CostComplete, q.DisplayComplete, q.Complete = false, false, false |
| 151 | q.Estimated, q.DisplayStatus, q.IncompleteReason = true, billing.DisplayStatusUnavailable, "usage_unknown" |
| 152 | return q |
| 153 | } |
| 154 | if e.Pricing == nil { |
| 155 | q := billing.CostQuote{ |
| 156 | Estimated: true, CostComplete: false, DisplayComplete: false, Complete: false, |
| 157 | DisplayStatus: billing.DisplayStatusUnavailable, IncompleteReason: "no_price", |
| 158 | ModelRef: e.ModelRef, UsageSource: e.UsageSource, |
| 159 | } |
| 160 | return &q |
| 161 | } |
| 162 | pricingCtx := ctx.pricingContext(e.ModelRef) |
| 163 | mode := billing.BillingModePAYG |
| 164 | if configured := strings.TrimSpace(pricingCtx.BillingMode); configured != "" { |
| 165 | mode = configured |
| 166 | } else if strings.Contains(strings.ToLower(e.ModelRef), "token-plan") || |
| 167 | strings.Contains(strings.ToLower(e.Source), "token-plan") { |
| 168 | mode = billing.BillingModeSubscriptionEquivalent |
| 169 | } |
| 170 | card := rateCardFromPricing(e.Pricing) |
| 171 | q := billing.BuildQuote(billing.QuoteInput{ |
| 172 | Usage: usageTokens(e.Usage), |
| 173 | Rates: card, |
| 174 | OccurredAt: now, |
| 175 | Display: display, |
| 176 | BillingMode: mode, |
| 177 | ModelRef: e.ModelRef, |
| 178 | UsageSource: firstUsageSource(e), |
| 179 | ProviderKind: pricingCtx.ProviderKind, |
| 180 | ModelID: pricingCtx.ModelID, |
| 181 | ScheduleID: pricingCtx.ScheduleID, |
| 182 | CatalogSource: pricingCtx.CatalogSource, |
| 183 | }) |
| 184 | return &q |
| 185 | } |
| 186 | |
| 187 | func firstUsageSource(e Event) string { |
| 188 | if s := strings.TrimSpace(e.UsageSource); s != "" { |
| 189 | return s |
| 190 | } |
| 191 | if s := strings.TrimSpace(e.Source); s != "" { |
| 192 | return s |
| 193 | } |
| 194 | return UsageSourceExecutor |
| 195 | } |
| 196 | |
| 197 | func rateCardFromPricing(p *provider.Pricing) billing.RateCard { |
| 198 | if p == nil { |
| 199 | return billing.RateCard{} |
| 200 | } |
| 201 | return billing.RateCard{ |
| 202 | CacheHit: p.CacheHit, |
| 203 | Input: p.Input, |
| 204 | Output: p.Output, |
| 205 | Currency: billing.NormalizeCurrency(p.Currency), |
| 206 | } |
| 207 | } |
| 208 | |
| 209 | func usageTokens(u *provider.Usage) billing.UsageTokens { |
| 210 | if u == nil { |
| 211 | return billing.UsageTokens{} |
| 212 | } |
| 213 | return billing.UsageTokens{ |
| 214 | PromptTokens: u.PromptTokens, |
| 215 | CompletionTokens: u.CompletionTokens, |
| 216 | CacheHitTokens: u.CacheHitTokens, |
| 217 | CacheMissTokens: u.CacheMissTokens, |
| 218 | CacheWriteTokens: u.CacheWriteTokens, |
| 219 | CacheWriteBilledTokens: u.CacheWriteBilledTokens, |
| 220 | Estimated: u.Estimated, |
| 221 | } |
| 222 | } |
| 223 |