diff --git a/docs/TELEMETRY.md b/docs/TELEMETRY.md index d89fdc973..6312116e0 100644 --- a/docs/TELEMETRY.md +++ b/docs/TELEMETRY.md @@ -50,11 +50,17 @@ any of the three drift apart. | usage_delta | mode | chat or task | | usage_delta | input_tokens | provider-reported input tokens since the preceding usage event | | usage_delta | output_tokens | provider-reported output tokens since the preceding usage event | -| usage_delta | total_tokens | the sum of this event's input and output tokens | +| usage_delta | total_tokens | the sum of this event's measured components, or the provider total with no split for total_only | | usage_delta | routing_provider | openrouter, openai, anthropic, google, bedrock, azure, ollama, other, or unknown | | usage_delta | model_family | deepseek, claude, gpt, gemini, llama, qwen, kimi, glm, other, or unknown | -| usage_delta | usage_status | reported or missing | +| usage_delta | usage_status | reported, total_only, mixed, estimated, partial, or missing | | usage_delta | accounting_source | provider or engine | +| usage_delta | estimated_input_tokens | estimated input tokens for a missing provider count | +| usage_delta | estimated_output_tokens | estimated observed output tokens for a missing provider count | +| usage_delta | estimated_total_tokens | the sum of this event's estimated input and output tokens | +| usage_delta | estimate_method | utf8_bytes_4_v1 | +| usage_delta | estimate_scope | text_only, observed_partial, multimodal_input, or unknown | +| usage_delta | usage_receipt | sha256 of the session and local receipt identity; groups estimates with later provider receipts | | session_ended | mode | chat or task | | session_ended | duration | a band: under 1m, 1-5m, 5-30m, 30m-2h, 2h or more | | session_ended | turns | a count band | @@ -71,8 +77,14 @@ any of the three drift apart. Count bands are 0, 1, 2-5, 6-20, 21-100 and 100+. Dollar bands are 0, under 0.01, 0.01-0.1, 0.1-1, 1-10 and 10+. +Streaming requests ask compatible APIs for token usage with `stream_options.include_usage`. +Connected Ollama services are classified as `ollama`, including remote runners; +other local or custom APIs remain `other`. Counts depend on the API returning +a usage block. Missing counts may carry separately labeled byte-based estimates; +these never replace provider-reported fields or enter billing. + Each provider usage receipt is appended locally before provider accounting returns, including -tokens read from cache once. Missing provider receipts produce a `usage_status = missing` event without token fields; a missing receipt never invents a zero total. Routing categories follow the actual attempt, including an explicitly authorized billing-door switch and any later generation receipt. Bounded routing and model-family categories describe reported usage without sending exact model names or endpoint addresses. Hosted engines currently report aggregate turn deltas with unknown routing and family, labelled `accounting_source = engine`; those aggregates are not per-call receipts and can be lost before a turn ends. Deltas, rather +tokens read from cache once. Missing token components carry `estimated_*` fields when text observations permit an estimate, otherwise remain absent. `usage_status` distinguishes reported, total-only, mixed, estimated, partial, and missing. A missing receipt never invents a measured zero total. Routing categories follow the actual attempt, including an explicitly authorized billing-door switch and any later generation receipt. Bounded routing and model-family categories describe reported usage without sending exact model names or endpoint addresses. Hosted engines currently report aggregate turn deltas with unknown routing and family, labelled `accounting_source = engine`; those aggregates are not per-call receipts and can be lost before a turn ends. Deltas, rather than repeated session totals, make the stream safe to sum while a run is open. first_run carries only the every-event properties and is sent once per @@ -154,3 +166,33 @@ the count GitHub already shows on the release page. Sidecar files such as batch it would send. The release-download workflow fails visibly when its production PostHog key or GitHub token is missing; dry-run still works without a PostHog key. + +## Estimates and recovered receipts + +See [Model usage accounting](model-usage-accounting.md) for the estimator, +limitations, privacy fields, receipt-aware PostHog query and deployment order. +A later provider receipt supersedes each measured component's earlier estimate by `usage_receipt`; +raw estimates and recovered measurements must not be added together. Estimates +never change billing, response usage or token limits. +The receipt key includes the physical attempt's identity so separate answers +stay separate even when a compatible service reuses a generation id. +Generation ids are trimmed before both observations and recovery. A one-sided +recovery leaves the other component unknown; a cost-only receipt adds no token +event. Invalid or unsafe counts and sums are omitted, and estimates require a +receipt identity. The final status describes the components actually emitted. + +Native `input_tokens`/`output_tokens` aliases now supply measured counts where +they previously stayed zero. Native cache reads and writes are included in +prompt tokens, matching the canonical cache-inclusive prompt convention. +Provider cost is banked unchanged; the provider billing boundary does not +calculate a token price when cost is absent. Estimates never supply a price. +Unmetered answers from a direct service retain the prior no-lookup behavior, +including overflow answers; their telemetry still follows the actual attempt. + +Connected services request streaming usage with `stream_options.include_usage`. +An explicit unsupported-option response on HTTP 400 or 422 retries once without +that option; later calls to that endpoint omit it. Unrelated errors do not turn +the option off. The default router road keeps its existing `usage.include` +request and sends no `stream_options`. +Independent reasoning or parameter repairs can compose with this fallback; +each distinct repair is used at most once in the repair chain. diff --git a/docs/changes/unreleased/1887-model-usage-accounting.md b/docs/changes/unreleased/1887-model-usage-accounting.md new file mode 100644 index 000000000..3c2811db9 --- /dev/null +++ b/docs/changes/unreleased/1887-model-usage-accounting.md @@ -0,0 +1,26 @@ +--- +kind: fixed +title: model usage accounting covers compatible providers and separates estimates from measurements +pr: 1887 +surface: [engine, docs] +invalidates: + - "Streaming requests previously sent only the router usage opt-in. Connected compatible APIs now receive stream_options.include_usage, with a narrow fallback for endpoints that explicitly reject it on 400 or 422. The default router road keeps its original wire shape; unrelated refusals gain no new repair." + - "Missing provider counts previously stayed unknown even when text was observed. Missing text components now have separate versioned estimates; reported zeroes, partial counts and total-only receipts remain measurements, and media or unseen output remains unknown." + - "A raw stream previously closed its call log without accounting for the completed response. It now records usage through the same provider boundary as normal calls." + - "Estimates and later provider receipts are not additive. They share a run-scoped physical-attempt receipt hash so analytics can prefer a measurement for each component without double counting. A provider that reuses generation ids no longer collapses distinct answers into one receipt." + - "Native input_tokens/output_tokens aliases previously recorded zero counts in responses, call logs and ledgers. They now provide measured counts, including native cache reads and writes in the cache-inclusive prompt count; equivalent canonical counts yield the same billing and token price." + - "Usage-option refusal repair previously stopped after another repair, or prevented a later reasoning repair. Independent repairs now compose once each; unrelated refusals keep their existing handling." + - "Repeated streamed tool names and opaque replay could inflate text estimates, and metadata-only cuts could invent zero output. Names now count once per call, opaque input remains unknown, and only observed text or normal completion establishes output." + - "Late cost-only and one-sided receipts previously fabricated measured zeroes. Recovery now preserves component presence and trimmed receipt identity without changing its price or ledger row." +--- + +Deploy the companion relay change in Agent-Field/website-open-infra#82 before +releasing the client. The prior relay rejects the new status values. Estimates +never enter billing, response usage or token-limit enforcement; existing opt-out +settings apply to the new count fields. The contract, limits and receipt-aware +query are documented in docs/model-usage-accounting.md. +The provider billing boundary copies a supplied cost and does not derive it from +tokens when cost is absent. Unmetered direct-service answers retain the prior +no-lookup behavior even on overflow; telemetry follows the actual attempt. +Telemetry status is derived from emitted components, unsafe counts and sums are +omitted, and estimates require a receipt identity. diff --git a/docs/model-usage-accounting.md b/docs/model-usage-accounting.md new file mode 100644 index 000000000..a921d71b0 --- /dev/null +++ b/docs/model-usage-accounting.md @@ -0,0 +1,67 @@ +# Model usage accounting + +Token usage is collected at the provider attempt boundary for every compatible endpoint, including OpenRouter, Ollama, vLLM, LM Studio, llama.cpp and custom or hosted APIs. The fallback does not depend on a provider name. + +## Measurements and estimates + +- Connected services request `stream_options.include_usage` on streamed chat completions. An explicit unsupported-field 400 or 422 retries without that option and remembers the refusal for that endpoint. Other validation errors do not disable usage or earn a new repair. The default router road keeps its existing `usage.include` opt-in and sends no `stream_options`. +- Keep provider-reported counts, including explicit zero. Normalize input/output aliases and native cache input without adding cache tokens twice. +- Keep a provider's total-only measurement without inventing its input/output split. +- Estimate only absent components, using `utf8_bytes_4_v1`: round observed UTF-8 bytes up at four bytes per token. Input includes the encoded text messages and tool definitions after shaping; opaque sidecars decline the input estimate. Output includes observed content, each tool name once per call, every argument fragment and visible textual reasoning. Metadata alone observes no output; a normally finished empty answer observes zero. +- Do not estimate media from URL or base64 length. Non-text input stays unknown; output text can still be estimated. Unseen output, encrypted reasoning, hidden thinking and work inside opaque hosted engines cannot be reconstructed. +- Mark interrupted-stream estimates `observed_partial`, never extrapolate to a completed answer. Failed or refused requests without an accepted response do not imply consumption. +- Estimates are approximate, particularly for non-English text, tokenizers and provider chat templates. They are suitable for usage visibility, not billing, token-limit enforcement, benchmark comparisons or exact cost attribution. They never modify response usage or the spend ledger. + +## Wire contract + +`input_tokens`, `output_tokens` and `total_tokens` remain measurements. `estimated_input_tokens`, `estimated_output_tokens` and `estimated_total_tokens` are separate, optional numeric fields, emitted only for unknown measured components and with a receipt identity. Each total sums only its own present components, except a provider's total-only measurement has no split. Negative counts, counts above 2^53-1 and groups whose sum exceeds that bound remain unknown. Missing is distinct from an explicitly reported or estimated zero. The final status follows the emitted components. + +`usage_status` is `reported`, `total_only`, `mixed`, `estimated`, `partial` or `missing`. `estimate_method` and `estimate_scope` label estimates; scopes are `text_only`, `observed_partial`, `multimodal_input` or `unknown`. + +`usage_receipt` is a SHA-256 hash scoped to the run and physical attempt. When a generation id exists, its actual endpoint and id participate only locally in that hash alongside the attempt identity. Providers that reuse generation ids cannot collapse separate answers into one receipt. Recovery keeps the original attempt identity. No prompt, response, endpoint, model slug, generation id, API key or other new raw identity is sent. Existing opt-out gates apply to estimates too. + +## Recovered receipts and dashboards + +An initial estimate and a later recovered provider receipt carry the same `usage_receipt`, but different event UUIDs. Generation ids are trimmed at both boundaries, and blank ids use the attempt identity alone. For each component, prefer a reported count, even if it is zero; retain an estimate only while no measurement exists. One-sided recovery is partial; cost-only recovery adds no token event and keeps its recovered price. A measured total-only receipt supersedes both component estimates. Do not add raw estimates and later measured receipts together. + +Native `input_tokens`/`output_tokens` now populate response, call-log and ledger counts where those aliases previously recorded zero. Native cache reads and writes join the prompt count to match canonical cache-inclusive `prompt_tokens`. The provider billing boundary copies any supplied cost and does not price tokens when cost is absent; downstream token pricing sees equivalent native and canonical counts. No estimate is banked, priced or shown as spend. Unmetered direct-service answers, including authorized overflow, retain their prior no-lookup behavior while telemetry identifies the actual attempt. + +The existing measured-only series still sums `total_tokens`. The new PostHog view groups by installation, event and receipt hash (falling back to the event UUID for older builds), then selects the best available components. Legacy session-end totals remain separate. Receipt time uses the earliest observation, with a one-day lookback margin around the displayed 14-day range. The query also exposes records with unknown input/output components. This is coverage of recorded receipts, not proof that every model call or anonymous installation reported. + +The estimate-aware line chart is configured separately in PostHog. Its estimated series starts only after rollout; historical missing receipts cannot be reconstructed. Verify the actual project receiving `codeaf:usage_delta` before installing a query. + +## Deployment order + +1. Merge and deploy the companion change in `Agent-Field/website-open-infra`, `site/src/app/api/oss/codeaf/telemetry/`. The previous relay drops unknown status values and strips new properties. +2. Merge the client PR into `dev`, then promote releases normally. +3. Confirm new numeric fields arrive in PostHog. No synthetic events are sent to the production relay by the tests. + +The client and relay carry identical `usage-accounting.json` fixtures. Client constructors are checked against that payload and the actual relay handler is checked through a mocked PostHog capture call, including estimates, measured zeroes, partial receipts, privacy stripping, safe integer totals and late-receipt identity. A full 50-event enriched batch is checked against the relay's 64 KiB body cap. + +## Receipt-aware PostHog query + +```sql +WITH receipts AS ( + SELECT event, distinct_id, + coalesce(nullIf(toString(properties.usage_receipt), ''), toString(uuid)) AS receipt_key, + min(timestamp) AS receipt_time, + max(toFloat(properties.total_tokens)) AS measured_total, + max(toFloat(properties.input_tokens)) AS measured_input, + max(toFloat(properties.output_tokens)) AS measured_output, + max(toFloat(properties.estimated_input_tokens)) AS estimated_input, + max(toFloat(properties.estimated_output_tokens)) AS estimated_output + FROM events + WHERE event IN ('codeaf:usage_delta', 'codeaf:session_ended') + AND properties.channel='stable' AND coalesce(properties.usage_context,'')!='ci' + AND timestamp >= today()-INTERVAL 15 DAY + GROUP BY event, distinct_id, receipt_key +) +SELECT toDate(receipt_time) AS day, + sumIf(coalesce(measured_total,0), event='codeaf:session_ended') AS legacy_reported_tokens, + sumIf(if(isNull(measured_input) AND isNull(measured_output),coalesce(measured_total,0),coalesce(measured_input,0)+coalesce(measured_output,0)), event='codeaf:usage_delta') AS reported_tokens, + sumIf(if(isNotNull(measured_total) AND isNull(measured_input) AND isNull(measured_output),0,if(isNull(measured_input),coalesce(estimated_input,0),0)+if(isNull(measured_output),coalesce(estimated_output,0),0)),event='codeaf:usage_delta') AS estimated_tokens, + countIf(event='codeaf:usage_delta' AND ((isNull(measured_input) AND isNull(estimated_input)) OR (isNull(measured_output) AND isNull(estimated_output)))) AS records_with_unknown_components +FROM receipts WHERE receipt_time >= today()-INTERVAL 14 DAY +GROUP BY day ORDER BY day +``` + diff --git a/internal/config/config.go b/internal/config/config.go index 6fb49d5d1..fc87aad5f 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -988,6 +988,7 @@ func ClientConfigFor(sources modelsource.Set, model string) provider.Config { effort, _ := provider.ParseEffort(level) service, bare := sources.For(model) configured := provider.Config{ + SourceID: service.Source.ID, APIKey: service.Key, BaseURL: service.Address, Model: bare, diff --git a/internal/config/sources_test.go b/internal/config/sources_test.go index 54c50daba..77460daaf 100644 --- a/internal/config/sources_test.go +++ b/internal/config/sources_test.go @@ -517,7 +517,7 @@ func TestAnUnqualifiedIdResolvesExactlyAsItDidBefore(t *testing.T) { for _, model := range models { configured := Config{APIKey: "default-key", BaseURL: base, Model: model}.ClientConfig(model) assertComparableProviderConfig(t, configured, provider.Config{ - APIKey: "default-key", BaseURL: base, Model: model, Effort: provider.EffortNone, + SourceID: modelsource.DefaultID, APIKey: "default-key", BaseURL: base, Model: model, Effort: provider.EffortNone, }) } } @@ -546,7 +546,7 @@ func TestAnUnqualifiedIdResolvesExactlyAsItDidBefore(t *testing.T) { for _, model := range models { configured := loaded.ClientConfig(model) assertComparableProviderConfig(t, configured, provider.Config{ - APIKey: "sk-profile-captured-1234567890", BaseURL: wantBase, + SourceID: modelsource.DefaultID, APIKey: "sk-profile-captured-1234567890", BaseURL: wantBase, Model: model, Effort: provider.EffortNone, Timeout: DefaultTimeout, }) } @@ -665,3 +665,14 @@ func TestOnlyAnExplicitlyKeyOptionalServiceAcceptsBlank(t *testing.T) { t.Fatalf("explicitly optional key = %+v, %v", outcome, err) } } + +func TestOllamaClientConfigPreservesServiceIdentity(t *testing.T) { + sources := modelsource.NewSet(modelsource.Connected{ + Source: modelsource.Source{ID: "ollama", Written: "local-models", KeyOptional: true}, + Address: "http://private-runner.example:9000/v1", + }) + configured := ClientConfigFor(sources, "local-models/qwen3") + if configured.SourceID != "ollama" || configured.BaseURL != "http://private-runner.example:9000/v1" || !configured.KeyOptional || !configured.Direct { + t.Fatalf("Ollama identity was lost: %+v", configured) + } +} diff --git a/internal/manual/chat/running-from-the-terminal.md b/internal/manual/chat/running-from-the-terminal.md index 8dfe73d69..b6c91fa10 100644 --- a/internal/manual/chat/running-from-the-terminal.md +++ b/internal/manual/chat/running-from-the-terminal.md @@ -1115,7 +1115,11 @@ codeaf sends anonymous usage counts to AgentField, **on by default**: version, O mode (chat or task), session counts and errors in bands, and the tokens each provider call used. Never prompts, code, file names, paths, repo names, keys, email, IP or machine name, and no model names. Each completed provider call queues a `usage_delta` -with exact input, output and total tokens. The queue is sent every 30 seconds while +with provider-reported token counts when available and separate estimates for missing +text counts. A hashed receipt key pairs an estimate with a later provider receipt, +so analytics select reported counts per component instead of adding both. A recovered +receipt with just one token count leaves the other component's estimate available; +a cost-only receipt adds no token measurement. No text is sent. The queue is sent every 30 seconds while the session stays open and once more when it ends; `session_ended` carries the outcome and bucketed counts without repeating those tokens. @@ -1125,6 +1129,25 @@ The relay is `https://agentfield.ai/api/oss/codeaf/telemetry`; The product prints no notice. With `model_pool` on, one scored row per crew seat also leaves for the Model Pool after a task lands (*Model Pool* above). +## Missing streaming token counts — connected services and refused usage options + +Connected services ask for streaming token counts with `stream_options.include_usage`. +If a service explicitly refuses that option with HTTP 400 or 422, codeaf retries +once without it and omits it on later calls to that endpoint. An unrelated error +does not turn usage reporting off. The default router uses its existing usage +request without the extra option. +The usage-option repair can follow or precede an independent reasoning repair; +each distinct repair is used once in that chain. + +Missing text counts may have separately labelled estimates in anonymous usage +events. Estimates never change spending, prices or token limits. An interrupted +answer estimates only what arrived. Provider-reported counts, including explicit +zeroes, remain measurements. Separate answers keep separate hashed receipts even +if a connected service repeats its own answer id. +Opaque reasoning replay in the input leaves its input estimate unknown. Role or +metadata frames alone do not establish output on a cut stream; a normally +finished empty answer can honestly estimate zero output. + ## How to turn telemetry off now — /settings stops sending immediately On the *display* tab of `/settings`, turn the `telemetry` switch off. It writes diff --git a/internal/manual/chat/settings.md b/internal/manual/chat/settings.md index 64bff9234..1aa467760 100644 --- a/internal/manual/chat/settings.md +++ b/internal/manual/chat/settings.md @@ -144,5 +144,12 @@ and **shared model recommendations** (formerly Model Pool). Recommendations can control is not a promise that all the others are disabled. Search accepts old labels as well as the new organization. +Anonymous token usage includes local models when their API returns token counts. +Streaming requests ask for those counts. A connected Ollama service is labeled +`ollama`; other local APIs are `other`. Missing text counts use a separately labeled byte-based estimate; media and unseen +output remain unknown. Opaque reasoning in the input also leaves its input estimate +unknown. Estimates do not enter billing or provider-reported totals. Provider +receipts replace estimates only for the token components they actually report. + There is no practice limit any more: codeaf no longer practises on its own, and `/budget practice` answers with exactly that. diff --git a/internal/provider/billing.go b/internal/provider/billing.go index feee7921e..aa7fa8f40 100644 --- a/internal/provider/billing.go +++ b/internal/provider/billing.go @@ -235,9 +235,15 @@ func (c *Client) bill(ctx context.Context, model string, response *ai.Response) // billAnswered bills an answer that arrived whole. One with no usage block is // settled like a cut one when the work asked for that ([WithUnmeteredReceipts]), // and billed the ordinary way — which banks nothing for it — otherwise. -func (c *Client) billAnswered(ctx context.Context, model string, response *ai.Response, answerBytes int) { +func (c *Client) billAnswered(ctx context.Context, model string, response *ai.Response, answerBytes int, traces ...*callTrace) { if response != nil && response.Usage == nil && unmeteredReceiptsFrom(ctx) { - c.settle(ctx, model, response, receiptUnmeteredReason, answerBytes) + // A direct client's receipt bearer belongs to its original service. + // Preserve the whole-answer boundary until recovery can carry the + // overflow attempt's own credential to its different host. + if c.config.Direct { + return + } + c.settle(ctx, model, response, receiptUnmeteredReason, answerBytes, traces...) return } c.bill(ctx, model, response) diff --git a/internal/provider/calllog.go b/internal/provider/calllog.go index d85cf9f7d..02809f6ec 100644 --- a/internal/provider/calllog.go +++ b/internal/provider/calllog.go @@ -121,6 +121,11 @@ type callTrace struct { // row written after the transport returned always names the attempt that // actually produced the answer. attemptID string + // Usage observations contain counts only, never retained prompt or reply text. + usageObservation usageObservation + // repairs survives physical attempts so relaxation cannot repeat a repair + // already used by this call's earlier request shape. + repairs map[string]bool // attemptBase is the address this attempt actually used, including an // explicitly authorized plan overflow. It never leaves the process. attemptBase string @@ -280,6 +285,7 @@ func (t *callTrace) begin() { } t.attempts++ t.attemptID = calllog.NewID() + t.usageObservation = usageObservation{inputBytes: t.usageObservation.inputBytes, inputTextOnly: t.usageObservation.inputTextOnly} } func (t *callTrace) note(learned ...string) { @@ -555,15 +561,8 @@ func (c *Client) record(facts recordFacts) { // sentence is the failure verdict and its cost the spend, 0 when absent. if record.Phase != calllog.PhaseStart { telemetry.CountModelCall(record.Error == "", record.Cost) - status := "missing" - if facts.response != nil && facts.response.Usage != nil { - status = "reported" - } - telemetry.CountUsage(record.PromptTokens, record.CompletionTokens, telemetry.UsageDimensions{ - RoutingProvider: telemetry.RoutingProvider(facts.baseFor(c.config.BaseURL)), - ModelFamily: telemetry.ModelFamily(model), UsageStatus: status, - AccountingSource: "provider", ReceiptID: record.ID, - }) + dimensions := c.usageDimensions(facts, record.ID, model) + telemetry.CountUsage(record.PromptTokens, record.CompletionTokens, dimensions) } c.recordBodies(facts, record, model) } diff --git a/internal/provider/client.go b/internal/provider/client.go index f9ad8ae1c..c5747a1d8 100644 --- a/internal/provider/client.go +++ b/internal/provider/client.go @@ -32,9 +32,13 @@ import ( // as is a constant (attribution.go), and a config field for it is exactly how a // caller ends up sending a different app — or none. type Config struct { - APIKey string - BaseURL string - Model string + // SourceID identifies the configured service for usage classification when + // its address alone cannot identify it, such as a remote Ollama runner. + // It stays local; only a bounded routing category reaches telemetry. + SourceID string + APIKey string + BaseURL string + Model string // Direct says this account is a connected service with one road rather than // a router with a set of serving lanes. It suppresses every lane preference, // sheet and probe at the transport boundary; a direct service must never be @@ -125,8 +129,10 @@ type Config struct { // TextStreamer interfaces structurally, so nothing above it knows a wire // format, and it owns the only outbound provider path in the process. type Client struct { - config Config - http *http.Client + // Unsupported stream-usage options are remembered per endpoint, never per model. + streamUsageUnsupported sync.Map + config Config + http *http.Client // stream is the same client with the total deadline removed. A streamed // answer is bounded by silence, not by duration — see send. stream *http.Client @@ -808,76 +814,115 @@ func (c *Client) sendRecovered(ctx context.Context, request *ai.Request, knobs c // The retry costs nothing: a 400 generated no tokens, and the answer is // remembered so only the first call on a model pays for the discovery. func (c *Client) sendRepaired(ctx context.Context, request *ai.Request, knobs callKnobs, stream bool) (*http.Response, error) { + repaired := make(map[string]bool) + if knobs.trace != nil { + if knobs.trace.repairs == nil { + knobs.trace.repairs = repaired + } else { + repaired = knobs.trace.repairs + } + } body, err := c.encodeRequest(request, knobs) if err != nil { return nil, encodeFailure(err) } began := logNow() response, err := c.send(ctx, request, knobs, body, stream) - if err == nil && response.StatusCode == http.StatusPaymentRequired && knobs.affordable == 0 { - peek, readErr := io.ReadAll(io.LimitReader(response.Body, maxErrorPeek)) - if readErr == nil { - if affordable := affordableTokens(peek); affordable > 0 { - c.record(recordFacts{ctx: ctx, request: request, knobs: knobs, stream: stream, - attempt: c.attemptsSoFar(knobs), began: began, status: response.StatusCode, - err: apiError(response.StatusCode, peek), responseBody: peek}) - knobs.affordable = affordable - return c.resend(ctx, request, knobs, stream, response) + // Each repair returns here so the next refusal can name a different field. + for { + if err == nil && response.StatusCode == http.StatusPaymentRequired && knobs.affordable == 0 && !repaired["affordable"] { + peek, readErr := io.ReadAll(io.LimitReader(response.Body, maxErrorPeek)) + if readErr == nil { + if affordable := affordableTokens(peek); affordable > 0 { + c.record(recordFacts{ctx: ctx, request: request, knobs: knobs, stream: stream, + attempt: c.attemptsSoFar(knobs), began: began, status: response.StatusCode, + err: apiError(response.StatusCode, peek), responseBody: peek}) + knobs.affordable = affordable + repaired["affordable"] = true + began = logNow() + response, err = c.resend(ctx, request, knobs, stream, response) + continue + } } + response.Body = rewound(peek, response.Body) } - response.Body = rewound(peek, response.Body) - } - if err != nil || !endpointRefusalStatus(response.StatusCode) { - return response, err - } - model := c.modelFor(request) - // Two repairs read the body, and they are told apart by what they cost the - // memo. A learned quirk is a 400 about a knob this adapter chose. Foreign - // reasoning is a 404 (OpenRouter's spelling) about content the transcript - // carried from a model the person has since switched away from — it is - // true of this conversation, not of the model, so it is fixed for this - // request and remembered nowhere. - foreign := len(knobs.reasoning) > 0 - repairable := response.StatusCode == http.StatusBadRequest && c.repairable(model, knobs) - if !foreign && !repairable { - return response, nil - } - peek, readErr := io.ReadAll(io.LimitReader(response.Body, maxErrorPeek)) - switch { - case readErr != nil: - case foreign && refusesForeignReasoning(peek): - // THE WORDS STILL GO; ONLY THE OTHER MODEL'S THINKING STAYS HOME. A - // sidecar tagged with its model never gets here (reasoning.go's - // producedElsewhere); this is the journal written before the tag. - // - // The refused shape gets its own row before the repair, because "the - // call that went out first was refused and the one that came back was a - // different request" is precisely the fact a log with only the answer - // on it cannot tell anybody (calllog.go). - c.record(recordFacts{ - ctx: ctx, request: request, knobs: knobs, stream: stream, - attempt: c.attemptsSoFar(knobs), began: began, - status: response.StatusCode, err: apiError(response.StatusCode, peek), - responseBody: peek, - }) - knobs.reasoning = nil - return c.resend(ctx, request, knobs, stream, response) - case repairable: - if learned := c.learn(model, knobs, peek); len(learned) > 0 { - knobs.trace.note(learned...) + streamUsage := stream && c.streamUsageSupported(knobs) + usageRefusalStatus := response != nil && (response.StatusCode == http.StatusBadRequest || response.StatusCode == http.StatusUnprocessableEntity) + if err != nil || (!endpointRefusalStatus(response.StatusCode) && !(streamUsage && usageRefusalStatus)) { + return response, err + } + model := c.modelFor(request) + // Two repairs read the body, and they are told apart by what they cost the + // memo. A learned quirk is a 400 about a knob this adapter chose. Foreign + // reasoning is a 404 (OpenRouter's spelling) about content the transcript + // carried from a model the person has since switched away from — it is + // true of this conversation, not of the model, so it is fixed for this + // request and remembered nowhere. + foreign := endpointRefusalStatus(response.StatusCode) && len(knobs.reasoning) > 0 + // The compatible usage option earns only its own repair. Peeking at an + // unrelated refusal must not admit it to the older learning road. + repairable := response.StatusCode == http.StatusBadRequest && c.repairable(model, knobs) + if !foreign && !repairable && !(streamUsage && usageRefusalStatus) { + return response, nil + } + peek, readErr := io.ReadAll(io.LimitReader(response.Body, maxErrorPeek)) + switch { + case readErr != nil: + case foreign && !repaired["foreign_reasoning"] && refusesForeignReasoning(peek): + // THE WORDS STILL GO; ONLY THE OTHER MODEL'S THINKING STAYS HOME. A + // sidecar tagged with its model never gets here (reasoning.go's + // producedElsewhere); this is the journal written before the tag. + // + // The refused shape gets its own row before the repair, because "the + // call that went out first was refused and the one that came back was a + // different request" is precisely the fact a log with only the answer + // on it cannot tell anybody (calllog.go). c.record(recordFacts{ ctx: ctx, request: request, knobs: knobs, stream: stream, attempt: c.attemptsSoFar(knobs), began: began, status: response.StatusCode, err: apiError(response.StatusCode, peek), - learned: learned, responseBody: peek, + responseBody: peek, }) - return c.resend(ctx, request, knobs, stream, response) + knobs.reasoning = nil + repaired["foreign_reasoning"] = true + began = logNow() + response, err = c.resend(ctx, request, knobs, stream, response) + continue + case streamUsage && usageRefusalStatus && !repaired["stream_usage"] && refusesStreamUsage(peek): + c.record(recordFacts{ctx: ctx, request: request, knobs: knobs, stream: stream, + attempt: c.attemptsSoFar(knobs), began: began, status: response.StatusCode, + err: apiError(response.StatusCode, peek), responseBody: peek}) + c.streamUsageUnsupported.Store(c.usageEndpoint(knobs), true) + repaired["stream_usage"] = true + began = logNow() + response, err = c.resend(ctx, request, knobs, stream, response) + continue + case repairable: + var learned []string + for _, repair := range c.learn(model, knobs, peek) { + if !repaired[repair] { + repaired[repair] = true + learned = append(learned, repair) + } + } + if len(learned) > 0 { + knobs.trace.note(learned...) + c.record(recordFacts{ + ctx: ctx, request: request, knobs: knobs, stream: stream, + attempt: c.attemptsSoFar(knobs), began: began, + status: response.StatusCode, err: apiError(response.StatusCode, peek), + learned: learned, responseBody: peek, + }) + began = logNow() + response, err = c.resend(ctx, request, knobs, stream, response) + continue + } } + // Not ours to fix. The body is handed back whole — the caller still has + // to read the provider's own words to build the error it reports. + response.Body = rewound(peek, response.Body) + return response, nil } - // Not ours to fix. The body is handed back whole — the caller still has - // to read the provider's own words to build the error it reports. - response.Body = rewound(peek, response.Body) - return response, nil } // resend closes a refused answer and sends the request again as the knobs now @@ -1109,6 +1154,10 @@ func (c *Client) completionInOnePiece( // has never heard of. var reasoningTokens int response.Usage, reasoningTokens = decoded.Usage.mergeInto(response.Usage, 0) + if knobs.trace != nil { + knobs.trace.usageObservation.observeUsage(decoded.Usage) + knobs.trace.usageObservation.observeWhole(payload) + } served := servedProvider(decoded.Provider) if err := c.terminalResponseError(ctx, request, knobs, &response, served); err != nil { c.record(recordFacts{ @@ -1170,7 +1219,7 @@ func (c *Client) completionInOnePiece( // for it an answer with no usage block is priced by its receipt // ([Client.billAnswered]). The bare door banks nothing without usage, so // that charge reached no book at all. - c.billAnswered(ctx, c.modelFor(request), &response, len(responseText(&response))) + c.billAnswered(ctx, c.modelFor(request), &response, len(responseText(&response)), knobs.trace) return nil, false, cut } reasonWord, servedWell := answerOutcome(&response) @@ -1193,7 +1242,7 @@ func (c *Client) completionInOnePiece( }) // The money, banked at the same instant the log row is written and for the // same reason: this is where the fact is known. See billing.go. - c.billAnswered(ctx, c.modelFor(request), &response, len(responseText(&response))) + c.billAnswered(ctx, c.modelFor(request), &response, len(responseText(&response)), knobs.trace) return &response, len(relearned) > 0, nil } @@ -1241,6 +1290,9 @@ type servedResponse struct { // into "the thinking pass spent the whole ceiling", which is the bug that made // the model-call log worth building. type usageWire struct { + inputReported bool + outputReported bool + totalReported bool ai.Usage CompletionTokensDetails struct { ReasoningTokens int `json:"reasoning_tokens"` @@ -1942,11 +1994,17 @@ func (c *Client) completeWithMessagesStreaming( // [usageWire.mergeInto], which both transports read a usage frame // through). response.Usage, reasoningTokens = chunk.Usage.mergeInto(response.Usage, reasoningTokens) + if knobs.trace != nil { + knobs.trace.usageObservation.observeUsage(chunk.Usage) + } } for _, choice := range chunk.Choices { if choice.Index != 0 { continue } + if knobs.trace != nil { + knobs.trace.usageObservation.observeDelta(choice.Delta) + } // THE SPLIT IS TAKEN FIRST AND ONCE (answer.go), because every line // below wants its answer rather than the channel the bytes came in // on: what the person is shown, what the response accumulates, what @@ -2167,6 +2225,7 @@ func (c *Client) completeWithMessagesStreaming( c.settle(ctx, c.modelFor(request), response, receiptRefusalReason, content.Len(), knobs.trace) return nil, false, err } + knobs.trace.usageObservation.outputObserved = true // The last call has no successor to close it, so the clean end of the stream // does. This runs only past the decode loop's error returns: a stream that // died mid-call announces nothing, because the fragment it stopped on may be @@ -2250,7 +2309,7 @@ func (c *Client) completeWithMessagesStreaming( ttft: firstTokenAfter(began, firstToken), }) // Through the answered door, for the reason the whole-body twin gives. - c.billAnswered(ctx, c.modelFor(request), response, content.Len()) + c.billAnswered(ctx, c.modelFor(request), response, content.Len(), knobs.trace) return nil, false, cut } // A RESCUE IS NOT THE TURN UNTIL IT READS AS LANGUAGE. The hedge used @@ -2266,7 +2325,7 @@ func (c *Client) completeWithMessagesStreaming( }) // A rescue that is not language was still paid for, so it goes through // the answered door too. - c.billAnswered(ctx, c.modelFor(request), response, content.Len()) + c.billAnswered(ctx, c.modelFor(request), response, content.Len(), knobs.trace) return nil, false, cut } // PAST EVERY GUARD, SO THIS LANE SERVED — the recovery half of the quality @@ -2296,7 +2355,7 @@ func (c *Client) completeWithMessagesStreaming( // Both paths or neither, exactly as the learning above: a streamed answer // is billed by the provider the same way a whole-body one is, and a ledger // blind to one of the two transports is a ledger nobody can reconcile. - c.billAnswered(ctx, c.modelFor(request), response, content.Len()) + c.billAnswered(ctx, c.modelFor(request), response, content.Len(), knobs.trace) finished = true observer(StreamEvent{Kind: StreamFinished, Session: session}) return response, relearned, nil @@ -2415,17 +2474,58 @@ func (c *Client) StreamComplete(ctx context.Context, prompt string, options ...a } decoder := newSSEDecoder(httpResponse.Body) + response := &ai.Response{} + began := logNow() + var terminalErr error + finished := false + failedFinish := false + defer func() { + if finished && terminalErr == nil && !failedFinish { + rawKnobs.trace.usageObservation.outputObserved = true + } + c.record(recordFacts{ctx: ctx, request: request, knobs: rawKnobs, stream: true, + began: began, status: httpResponse.StatusCode, response: response, err: terminalErr}) + if finished { + c.billAnswered(ctx, c.modelFor(request), response, rawKnobs.trace.usageObservation.outputBytes, rawKnobs.trace) + } else { + c.settle(ctx, c.modelFor(request), response, receiptTornReason, rawKnobs.trace.usageObservation.outputBytes, rawKnobs.trace) + } + }() for { - chunk, err := decoder.Decode() + wire, err := decoder.DecodeChunk() if err != nil { - if !errors.Is(err, io.EOF) { - errs <- fmt.Errorf("decode stream: %w", err) + if errors.Is(err, io.EOF) && (decoder.done || finished) { + finished = true + } else { + terminalErr = err + if errors.Is(err, io.EOF) { + terminalErr = io.ErrUnexpectedEOF + } + errs <- fmt.Errorf("decode stream: %w", terminalErr) } return } + if response.ID == "" { + response.ID = wire.ID + } + rawKnobs.trace.usageObservation.observeUsage(wire.Usage) + if wire.Usage != nil { + response.Usage, _ = wire.Usage.mergeInto(response.Usage, 0) + } + chunk := ai.StreamChunk{ID: wire.ID, Object: wire.Object, Created: wire.Created, Model: wire.Model, Usage: wire.Usage.usage()} + for _, choice := range wire.Choices { + rawKnobs.trace.usageObservation.observeDelta(choice.Delta) + if choice.FinishReason != nil { + finished = true + failedFinish = failedFinish || strings.EqualFold(strings.TrimSpace(*choice.FinishReason), "error") + } + chunk.Choices = append(chunk.Choices, ai.StreamDelta{Index: choice.Index, + Delta: ai.MessageDelta{Role: choice.Delta.Role, Content: choice.Delta.Content}, FinishReason: choice.FinishReason}) + } select { case <-ctx.Done(): - errs <- ctx.Err() + terminalErr = ctx.Err() + errs <- terminalErr return case chunks <- chunk: } diff --git a/internal/provider/receipt.go b/internal/provider/receipt.go index 5afd976c7..82da23e54 100644 --- a/internal/provider/receipt.go +++ b/internal/provider/receipt.go @@ -101,7 +101,7 @@ type receiptWork struct { result Reconciled sink ReconcileSink queued time.Time - usage func(int, int) + usage func(int, int, bool, bool) base string } @@ -190,7 +190,7 @@ func (c *Client) settle(ctx context.Context, model string, response *ai.Response Hedged: hedgeLaneFrom(ctx) != "", } if response != nil { - result.Ref = strings.TrimSpace(response.ID) + result.Ref = canonicalGenerationID(response.ID) } // NO ID MEANS NO REQUEST. Text still proves that the call got somewhere and // may be charged, so its missing price is said; without text there is no @@ -201,10 +201,14 @@ func (c *Client) settle(ctx context.Context, model string, response *ai.Response } return } - work := receiptWork{result: result, sink: sink, base: base, usage: telemetry.CaptureUsageRecorder(telemetry.UsageDimensions{ - RoutingProvider: telemetry.RoutingProvider(base), + attemptID := "" + if len(traces) > 0 && traces[0] != nil { + attemptID = traces[0].attemptID + } + work := receiptWork{result: result, sink: sink, base: base, usage: telemetry.CaptureUsageRecorderWithPresence(telemetry.UsageDimensions{ + RoutingProvider: c.usageRoutingProvider(base), ModelFamily: telemetry.ModelFamily(result.Model), UsageStatus: "reported", - AccountingSource: "provider", ReceiptID: "generation:" + result.Ref, + AccountingSource: "provider", ReceiptID: usageReceiptID(base, result.Ref, attemptID), })} // THE WORK IS TOLD A RECEIPT IS OWED BEFORE IT IS QUEUED, and told it was // answered only after the sink has banked it, so a caller waiting for its @@ -287,7 +291,14 @@ func (c *Client) reconcile(work receiptWork) { ctx, cancel := context.WithDeadline(context.Background(), work.deadline()) defer cancel() for attempt := 0; attempt < receiptAttempts; attempt++ { - billed, found, noRoute := c.fetchReceiptAt(ctx, result.Ref, base) + var presence [2]bool + billed, found, noRoute := c.fetchReceiptAt(ctx, result.Ref, base, &presence) + // Explicit zero counts are measurements even when the old money door + // still considers this an empty, unpriced receipt. Keep its price and + // ledger behavior while allowing those counts to replace estimates. + if work.usage != nil && (presence[0] || presence[1]) { + work.usage(billed.PromptTokens, billed.CompletionTokens, presence[0], presence[1]) + } if noRoute { receiptRoutes.heard(base, time.Now()) work.sink(result) @@ -298,12 +309,6 @@ func (c *Client) reconcile(work receiptWork) { billed.Model = result.Model result.Billed = billed result.Found = true - // Recovery is reached only when the stream had no usage block, so - // these tokens replace an unknown receipt rather than add to a - // previously reported total. - if work.usage != nil { - work.usage(billed.PromptTokens, billed.CompletionTokens) - } work.sink(result) return } @@ -337,7 +342,7 @@ func (c *Client) fetchReceipt(ctx context.Context, ref string) (Billed, bool, bo } // fetchReceiptAt follows the route the generation actually used. -func (c *Client) fetchReceiptAt(ctx context.Context, ref, base string) (Billed, bool, bool) { +func (c *Client) fetchReceiptAt(ctx context.Context, ref, base string, observed ...*[2]bool) (Billed, bool, bool) { endpoint, err := url.Parse(strings.TrimSpace(base)) if err != nil { return Billed{}, false, false @@ -379,13 +384,19 @@ func (c *Client) fetchReceiptAt(ctx context.Context, ref, base string) (Billed, if wire.Data.TotalCost != nil { billed.Cost = *wire.Data.TotalCost } + presence := [2]bool{} if wire.Data.TokensPrompt != nil || wire.Data.TokensCompletion != nil { + presence = [2]bool{wire.Data.TokensPrompt != nil, wire.Data.TokensCompletion != nil} billed.PromptTokens = valueOrZero(wire.Data.TokensPrompt) billed.CompletionTokens = valueOrZero(wire.Data.TokensCompletion) } else { + presence = [2]bool{wire.Data.NativeTokensPrompt != nil, wire.Data.NativeTokensCompletion != nil} billed.PromptTokens = valueOrZero(wire.Data.NativeTokensPrompt) billed.CompletionTokens = valueOrZero(wire.Data.NativeTokensCompletion) } + if len(observed) > 0 && observed[0] != nil { + *observed[0] = presence + } if billed.Empty() { return Billed{}, false, false } diff --git a/internal/provider/sse.go b/internal/provider/sse.go index 55b28eaee..3edea60fd 100644 --- a/internal/provider/sse.go +++ b/internal/provider/sse.go @@ -395,17 +395,23 @@ func (a *toolCallAccumulator) current() (formingCall, bool) { // // The last call of a message has no such successor and is reported by // [toolCallAccumulator.flush] instead, at the clean end of the stream. -func (a *toolCallAccumulator) add(fragment toolCallDelta) (ai.ToolCall, bool) { +// toolCallIndex keeps observation and assembly on the same call boundaries. +func toolCallIndex(fragment toolCallDelta, order []int) int { index := 0 switch { case fragment.Index != nil: index = *fragment.Index - case len(a.order) > 0 && fragment.ID == "": + case len(order) > 0 && fragment.ID == "": // A continuation of the call already open. - index = a.order[len(a.order)-1] + index = order[len(order)-1] default: - index = len(a.order) + index = len(order) } + return index +} + +func (a *toolCallAccumulator) add(fragment toolCallDelta) (ai.ToolCall, bool) { + index := toolCallIndex(fragment, a.order) if a.calls == nil { a.calls, a.args = make(map[int]*ai.ToolCall, 2), make(map[int]*strings.Builder, 2) a.reported = make(map[int]bool, 2) diff --git a/internal/provider/telemetry_usage.go b/internal/provider/telemetry_usage.go new file mode 100644 index 000000000..7e574e0ee --- /dev/null +++ b/internal/provider/telemetry_usage.go @@ -0,0 +1,17 @@ +package provider + +import ( + "strings" + + "github.com/Agent-Field/codeaf/internal/telemetry" +) + +// usageRoutingProvider preserves a named Ollama connection without guessing +// from a loopback address or port, which may serve any compatible API. An +// attempt on another billing road must be classified from its actual address. +func (c *Client) usageRoutingProvider(base string) string { + if base == c.config.BaseURL && strings.EqualFold(strings.TrimSpace(c.config.SourceID), "ollama") { + return "ollama" + } + return telemetry.RoutingProvider(base) +} diff --git a/internal/provider/telemetry_usage_test.go b/internal/provider/telemetry_usage_test.go index fb30853d7..2ba8402d4 100644 --- a/internal/provider/telemetry_usage_test.go +++ b/internal/provider/telemetry_usage_test.go @@ -3,7 +3,9 @@ package provider import ( "context" "encoding/json" + "fmt" "net/http" + "strings" "testing" "github.com/Agent-Field/agentfield/sdk/go/ai" @@ -179,3 +181,118 @@ func TestLateReceiptFollowsTheOverflowRoute(t *testing.T) { t.Fatal(event.Props) } } + +// The fixture only sends a usage frame when the streaming request opts in, +// matching compatible local servers rather than gifting every caller counts. +func TestLocalStreamingUsageIsRequestedAndCounted(t *testing.T) { + for _, source := range []string{"ollama", "custom"} { + for _, reported := range []bool{true, false} { + name := fmt.Sprintf("%s/reported=%t", source, reported) + t.Run(name, func(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "local") + defer finish() + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request struct { + Stream bool `json:"stream"` + Options struct { + IncludeUsage bool `json:"include_usage"` + } `json:"stream_options"` + } + if err := json.NewDecoder(r.Body).Decode(&request); err != nil { + t.Fatal(err) + } + if !request.Stream || !request.Options.IncludeUsage { + t.Error("stream did not request token usage") + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, `data: {"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"}}]}`+"\n\n") + fmt.Fprint(w, `data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`+"\n\n") + if reported && request.Options.IncludeUsage { + fmt.Fprint(w, `data: {"choices":[],"usage":{"prompt_tokens":40,"completion_tokens":2,"total_tokens":42}}`+"\n\n") + } + fmt.Fprint(w, "data: [DONE]\n\n") + }) + client, err := NewClient(Config{SourceID: source, BaseURL: "http://private-runner.example:11434/v1", Model: "qwen3", Direct: true, KeyOptional: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + ctx := WithStreamObserver(context.Background(), func(StreamEvent) {}) + response, err := client.CompleteWithMessages(ctx, userMessages("hello")) + if err != nil { + t.Fatal(err) + } + if response.Text() != "ok" { + t.Fatalf("response = %q", response.Text()) + } + rows := telemetry.SpoolContents() + if len(rows) != 1 { + t.Fatalf("receipts = %d, want 1", len(rows)) + } + var event struct { + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &event); err != nil { + t.Fatal(err) + } + want := "other" + if source == "ollama" { + want = "ollama" + } + if event.Props["routing_provider"] != want { + t.Fatal(event.Props) + } + if reported { + if event.Props["total_tokens"] != float64(42) || event.Props["usage_status"] != "reported" { + t.Fatal(event.Props) + } + } else { + if _, exists := event.Props["total_tokens"]; exists { + t.Fatal("missing usage invented tokens") + } + if event.Props["usage_status"] != "estimated" || event.Props["estimated_total_tokens"] == nil { + t.Fatal(event.Props) + } + } + if strings.Contains(string(rows[0]), "private-runner") { + t.Fatal("private endpoint escaped") + } + }) + } + } +} + +func TestUsageProviderKeepsOllamaIdentityOnItsOwnRouteOnly(t *testing.T) { + for _, tc := range []struct{ source, base, attempt, want string }{ + {"ollama", "http://localhost:11434/v1", "http://localhost:11434/v1", "ollama"}, + {"ollama", "http://[::1]:9000/v1", "http://[::1]:9000/v1", "ollama"}, + {"ollama", "http://runner.example/v1", "https://openrouter.ai/api/v1", "openrouter"}, + {"custom", "http://localhost:11434/v1", "http://localhost:11434/v1", "other"}, + {"private-service", "http://localhost:8080/v1", "http://localhost:8080/v1", "other"}, + } { + client := &Client{config: Config{SourceID: tc.source, BaseURL: tc.base}} + if got := client.usageRoutingProvider(tc.attempt); got != tc.want { + t.Errorf("%+v: got %s", tc, got) + } + } +} + +func TestNonStreamingRequestDoesNotSendStreamOptions(t *testing.T) { + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://localhost:9000/v1", Model: "qwen3", Direct: true}) + if err != nil { + t.Fatal(err) + } + body, err := client.encodeRequest(&ai.Request{Messages: userMessages("hello")}, callKnobs{}) + if err != nil { + t.Fatal(err) + } + var request map[string]any + if err := json.Unmarshal(body, &request); err != nil { + t.Fatal(err) + } + if _, exists := request["stream_options"]; exists { + t.Fatal("non-streamed call sent stream_options") + } +} diff --git a/internal/provider/usage_c3_contract_test.go b/internal/provider/usage_c3_contract_test.go new file mode 100644 index 000000000..4cf3758fc --- /dev/null +++ b/internal/provider/usage_c3_contract_test.go @@ -0,0 +1,440 @@ +package provider + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "reflect" + "strings" + "testing" + "time" + + "github.com/Agent-Field/agentfield/sdk/go/ai" + "github.com/Agent-Field/codeaf/internal/home" + "github.com/Agent-Field/codeaf/internal/telemetry" +) + +func usageC3Session(t *testing.T) { + t.Helper() + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + t.Cleanup(telemetry.BeginUsageSession(telemetry.ModeChat, "c3")) +} + +func usageC3Events(t *testing.T) []map[string]any { + t.Helper() + var events []map[string]any + for _, row := range telemetry.SpoolContents() { + var e struct { + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(row, &e); err != nil { + t.Fatal(err) + } + events = append(events, e.Props) + } + return events +} + +func TestK10TrimmedGenerationPairsRecoveryAndBlankUsesAttempt(t *testing.T) { + usageC3Session(t) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if strings.HasSuffix(r.URL.Path, "/generation") { + if r.URL.Query().Get("id") != "late" { + t.Error("generation id was not canonicalized") + } + fmt.Fprint(w, `{"data":{"tokens_prompt":10,"tokens_completion":3,"total_cost":0.01}}`) + } else { + fmt.Fprint(w, `{"id":" late ","choices":[{"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}]}`) + } + }) + c, err := NewClient(Config{APIKey: "test", BaseURL: "https://openrouter.ai/api/v1", Model: "qwen3", HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + result := make(chan Reconciled, 1) + ctx := WithUnmeteredReceipts(WithReconcile(context.Background(), func(r Reconciled) { result <- r })) + if _, err := c.CompleteWithMessages(ctx, userMessages("hello")); err != nil { + t.Fatal(err) + } + if !receiptResult(t, result).Found { + t.Fatal("receipt missing") + } + rows := usageC3Events(t) + if len(rows) != 2 || rows[0]["usage_receipt"] != rows[1]["usage_receipt"] { + t.Fatalf("receipts: %+v", rows) + } + if rows[0]["usage_status"] != "estimated" || rows[1]["usage_status"] != "reported" { + t.Fatal(rows) + } + for _, generation := range []string{"", " \t "} { + if got := usageReceiptID("base", generation, "attempt"); got != "attempt" { + t.Fatal(got) + } + } +} + +func TestK5RepeatedToolNamesCountOnceAndEveryArgumentFragmentCounts(t *testing.T) { + for _, explicit := range []bool{true, false} { + o := usageObservation{} + zero := 0 + for i := 0; i < 10; i++ { + fragment := toolCallDelta{Function: toolFunctionDelta{Name: "read", Arguments: "x"}} + if explicit { + fragment.Index = &zero + } + o.observeDelta(streamDelta{ToolCalls: []toolCallDelta{fragment}}) + } + if o.outputBytes != 14 { + t.Fatalf("explicit=%v: bytes=%d, want 14", explicit, o.outputBytes) + } + } + // Without indexes, an id opens a call and id-less fragments continue it. + o := usageObservation{} + for _, f := range []toolCallDelta{ + {ID: "first", Function: toolFunctionDelta{Name: "read", Arguments: "x"}}, + {Function: toolFunctionDelta{Name: "read", Arguments: "x"}}, + {ID: "second", Function: toolFunctionDelta{Name: "read", Arguments: "x"}}, + {Function: toolFunctionDelta{Name: "read", Arguments: "x"}}, + } { + o.observeDelta(streamDelta{ToolCalls: []toolCallDelta{f}}) + } + if o.outputBytes != 12 { + t.Fatal(o.outputBytes) + } +} + +func TestK5OpaqueInputCannotGrowATextEstimate(t *testing.T) { + c, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "qwen3", Direct: true}) + if err != nil { + t.Fatal(err) + } + for _, size := range []int{4, 400} { + ctx := WithMessageReasoning(context.Background(), []MessageReasoning{{}, {Details: json.RawMessage(fmt.Sprintf(`[{"type":"reasoning.encrypted","data":%q}]`, strings.Repeat("x", size)))}}) + knobs := knobsFrom(ctx) + req := &ai.Request{Model: "qwen3", Messages: reasoningTranscript()[:2]} + if _, err := c.encodeRequest(req, knobs); err != nil { + t.Fatal(err) + } + d := c.usageDimensions(recordFacts{status: 200, response: &ai.Response{}, knobs: knobs}, "attempt", "qwen3") + if d.EstimatedInput != nil { + t.Fatalf("ciphertext %d estimated: %+v", size, d) + } + } + for _, field := range []string{"signature", "reasoning_details", "opaque"} { + raw := json.RawMessage(fmt.Sprintf(`{"role":"assistant","content":"hello",%q:"%s"}`, field, strings.Repeat("x", 400))) + if _, known := usageInputBytes([]json.RawMessage{raw}, nil, userMessages("hello")); known { + t.Fatalf("%s counted as text", field) + } + } +} + +func TestK2IndependentRepairsComposeInEitherOrderAndStayBounded(t *testing.T) { + for _, order := range []string{"reasoning-first", "usage-first", "usage-repeated", "reasoning-repeated", "learned-first"} { + t.Run(order, func(t *testing.T) { + model := "qwen/c3-" + order + quirksAt(t, model) + requests := 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + messages := body["messages"].([]any) + thought := messages[1].(map[string]any)["reasoning_content"] + usage := body["stream_options"] != nil + refusal := "" + status := 400 + switch { + case order == "learned-first" && requests == 1: + refusal = "reasoning cannot be disabled" + case order == "usage-repeated" || (order == "usage-first" && usage) || (order == "reasoning-first" && thought == nil && usage) || (order == "learned-first" && usage && requests > 1): + refusal = "unknown field stream_options" + case order == "reasoning-repeated" || thought != nil: + refusal = "Your request contains encrypted reasoning or compaction content that was produced under a different model. Encrypted payloads can only be replayed to the endpoint that created them." + status = 404 + } + if refusal != "" { + w.WriteHeader(status) + json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message": refusal}}) + return + } + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + }) + c, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: model, Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + ctx := WithMessageReasoning(context.Background(), []MessageReasoning{{}, {Field: "reasoning_content", Text: "carry"}}) + knobs := knobsFrom(ctx) + if order == "usage-repeated" { + knobs.reasoning = nil + } + if order == "learned-first" { + knobs.reasoning = nil + knobs.effort = effortRequest{effort: EffortOff, explicit: true} + } + request := &ai.Request{Model: model, Messages: reasoningTranscript()[:2], Stream: true} + response, err := c.sendRepaired(ctx, request, knobs, true) + if err != nil { + t.Fatal(err) + } + response.Body.Close() + repeated := strings.HasSuffix(order, "repeated") + want := 3 + if repeated { + want = 2 + } + if requests != want || (response.StatusCode >= 400) != repeated { + t.Fatalf("requests=%d status=%d", requests, response.StatusCode) + } + if !repeated && !strings.HasPrefix(order, "learned") { + // The endpoint memo survives the call, while foreign reasoning is repaired anew. + r, err := c.sendRepaired(ctx, request, knobsFrom(ctx), true) + if err != nil { + t.Fatal(err) + } + r.Body.Close() + if requests != 5 { + t.Fatalf("memo lost: requests=%d", requests) + } + } + }) + } +} + +func TestK7RoleOnlyCutKeepsOutputUnknownAndFinishedEmptyObservesZero(t *testing.T) { + for _, finish := range []bool{false, true} { + t.Run(fmt.Sprint(finish), func(t *testing.T) { + usageC3Session(t) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"reasoning_details\":[{\"data\":\"opaque\"}]}}],\"usage\":{\"prompt_tokens\":7}}\n\n") + if finish { + fmt.Fprint(w, "data: [DONE]\n\n") + } + }) + c, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "qwen3", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + chunks, errs := c.StreamComplete(context.Background(), "hello") + for range chunks { + } + var failed bool + for err := range errs { + failed = failed || err != nil + } + rows := usageC3Events(t) + if len(rows) != 1 { + t.Fatal(rows) + } + p := rows[0] + if finish { + if failed || p["usage_status"] != "mixed" || p["estimated_output_tokens"] != float64(0) { + t.Fatal(p) + } + } else if !failed || p["usage_status"] != "partial" || p["estimated_output_tokens"] != nil { + t.Fatal(p) + } + }) + } + for _, finished := range []bool{false, true} { + o := usageObservation{} + tail := "" + if finished { + tail = `,"finish_reason":"stop"` + } + o.observeWhole([]byte(`{"choices":[{"message":{"role":"assistant"}` + tail + `}]}`)) + if o.outputObserved != finished { + t.Fatalf("whole observation: %+v", o) + } + } +} + +func TestK2RefusalMustQualifyTheUsageOption(t *testing.T) { + for _, tc := range []struct { + body string + want bool + }{ + {`{"error":{"message":"stream_options is supported; unknown tool type"}}`, false}, + {`{"error":{"message":"unknown tool type and stream_options is supported"}}`, false}, + {`{"error":{"message":"Unrecognized request argument supplied: stream_options"}}`, true}, + {`{"error":{"message":"unknown field stream_options"}}`, true}, + {`{"error":{"message":"stream_options is not supported"}}`, true}, + {`{"error":{"message":"include_usage not supported"}}`, true}, + {`{"detail":[{"loc":["body","stream_options"],"msg":"Extra inputs are not permitted"}]}`, true}, + {`{"detail":[{"loc":["body","stream_options"],"msg":"supported"},{"loc":["body","tool"],"msg":"Extra inputs are not permitted"}]}`, false}, + {`{"error":{"message":"Extra inputs are not permitted: stream_options"}}`, true}, + } { + if got := refusesStreamUsage([]byte(tc.body)); got != tc.want { + t.Errorf("%s: %v", tc.body, got) + } + } +} + +func TestK9NativeCacheAliasesKeepCanonicalBillingAndPrice(t *testing.T) { + var bills []Billed + for _, body := range []string{ + `{"input_tokens":7,"output_tokens":3,"cache_read_input_tokens":5,"cache_creation_input_tokens":2}`, + `{"prompt_tokens":14,"completion_tokens":3,"cache_read_input_tokens":5,"cache_creation_input_tokens":2}`, + } { + var wire usageWire + if err := json.Unmarshal([]byte(body), &wire); err != nil { + t.Fatal(err) + } + c := &Client{} + c.bill(WithBilling(context.Background(), func(b Billed) { bills = append(bills, b) }), "qwen3", &ai.Response{Usage: wire.usage()}) + } + if len(bills) != 2 || !reflect.DeepEqual(bills[0], bills[1]) || bills[0].PromptTokens != 14 || bills[0].CachedTokens != 5 || bills[0].Cost != 0 { + t.Fatal(bills) + } + // SpendGuard.after subtracts cached reads from the inclusive prompt count + // before applying the regular prompt rate, and prices cache reads separately. + price := func(b Billed) float64 { + return float64(b.PromptTokens-b.CachedTokens)*2 + float64(b.CachedTokens)*0.5 + float64(b.CompletionTokens)*4 + } + if price(bills[0]) != 32.5 || price(bills[0]) != price(bills[1]) { + t.Fatal(bills) + } +} + +func TestK9UnmeteredDirectOverflowDoesNotLookUpWithVendorBearer(t *testing.T) { + usageC3Session(t) + requests := 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + if r.Method == http.MethodGet { + t.Error("unmetered overflow sent a receipt lookup") + } + if r.Header.Get("Authorization") != "Bearer vendor-key" { + t.Error("unexpected attempt bearer") + } + w.Header().Set("Content-Type", "application/json") + if r.URL.Host == "api.anthropic.com" { + w.WriteHeader(429) + fmt.Fprint(w, `{"code":"1316","message":"plan window exhausted"}`) + return + } + fmt.Fprint(w, `{"id":"overflow","choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`) + }) + c, err := NewClient(Config{APIKey: "vendor-key", BaseURL: "https://api.anthropic.com/v1", Model: "qwen3", Direct: true, PlanOverflow: "https://openrouter.ai/api/v1", OverflowOnPlanPause: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + ctx := WithUnmeteredReceipts(WithReconcile(context.Background(), func(Reconciled) { t.Error("unmetered direct overflow reconciled") })) + response, err := c.CompleteWithMessages(ctx, userMessages("hello")) + if err != nil || response.Text() != "ok" || requests != 2 { + t.Fatalf("response=%v err=%v requests=%d", response, err, requests) + } + rows := usageC3Events(t) + if len(rows) != 2 || rows[1]["routing_provider"] != "openrouter" || rows[1]["usage_status"] != "estimated" { + t.Fatal(rows) + } + // The pre-existing torn-stream route still uses the client's original key. + req, err := c.newHTTPRequestAt(context.Background(), &ai.Request{APIKeyOverride: "overflow-key"}, []byte(`{}`), true, "https://openrouter.ai/api/v1") + if err != nil || req.Header.Get("Authorization") != "Bearer overflow-key" { + t.Fatal(req, err) + } +} + +func TestK10RecoveredTokenPresenceNeverFabricatesZeroes(t *testing.T) { + for _, tc := range []struct { + tokens, status string + events int + in, out any + }{ + {``, "estimated", 1, nil, nil}, + {`,"tokens_prompt":10`, "partial", 2, float64(10), nil}, + {`,"tokens_completion":0`, "partial", 2, nil, float64(0)}, + {`,"tokens_prompt":0,"tokens_completion":0`, "reported", 2, float64(0), float64(0)}, + } { + t.Run(tc.tokens, func(t *testing.T) { + usageC3Session(t) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if strings.HasSuffix(r.URL.Path, "/generation") { + fmt.Fprint(w, `{"data":{"total_cost":0.01`+tc.tokens+`}}`) + return + } + fmt.Fprint(w, `{"id":"late","choices":[{"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}]}`) + }) + c, err := NewClient(Config{APIKey: "test", BaseURL: "https://openrouter.ai/api/v1", Model: "qwen3", HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + results := make(chan Reconciled, 1) + ctx := WithUnmeteredReceipts(WithReconcile(context.Background(), func(r Reconciled) { results <- r })) + if _, err := c.CompleteWithMessages(ctx, userMessages("hello")); err != nil { + t.Fatal(err) + } + answer := receiptResult(t, results) + if !answer.Found || answer.Cost != 0.01 { + t.Fatal(answer) + } + if tc.in != nil && answer.PromptTokens != int(tc.in.(float64)) { + t.Fatal(answer) + } + if tc.out != nil && answer.CompletionTokens != int(tc.out.(float64)) { + t.Fatal(answer) + } + rows := usageC3Events(t) + if len(rows) != tc.events { + t.Fatal(rows) + } + last := rows[len(rows)-1] + if last["usage_status"] != tc.status || last["input_tokens"] != tc.in || last["output_tokens"] != tc.out { + t.Fatal(last) + } + if tc.events == 2 && rows[0]["usage_receipt"] != last["usage_receipt"] { + t.Fatal(rows) + } + }) + } +} + +func TestK3SupportedUsageMentionDoesNotResendAnUnknownTool(t *testing.T) { + requests := 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + w.WriteHeader(400) + fmt.Fprint(w, `{"error":{"message":"stream_options is supported; unknown tool type"}}`) + }) + c, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + if _, err := c.CompleteWithMessages(context.Background(), userMessages("hello")); err == nil { + t.Fatal("expected the tool refusal") + } + if requests != 1 || !c.streamUsageSupported(callKnobs{}) { + t.Fatalf("requests=%d option disabled=%v", requests, !c.streamUsageSupported(callKnobs{})) + } +} + +func TestK10RecoveredExplicitZeroesKeepTheExistingEmptyMoneyDoor(t *testing.T) { + usageC3Session(t) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"data":{"total_cost":0,"tokens_prompt":0,"tokens_completion":0}}`) + }) + c, err := NewClient(Config{APIKey: "test", BaseURL: "https://openrouter.ai/api/v1", Model: "qwen3", HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + c.wait = func(context.Context, time.Duration) error { return nil } + result := make(chan Reconciled, 1) + c.settle(WithReconcile(context.Background(), func(r Reconciled) { result <- r }), "qwen3", &ai.Response{ID: "zero-recovery"}, receiptTornReason, 1) + if r := receiptResult(t, result); r.Found || r.Cost != 0 { + t.Fatal("empty money behavior changed", r) + } + rows := usageC3Events(t) + if len(rows) != 1 || rows[0]["usage_status"] != "reported" || rows[0]["input_tokens"] != float64(0) || rows[0]["output_tokens"] != float64(0) { + t.Fatal(rows) + } +} diff --git a/internal/provider/usage_contract_test.go b/internal/provider/usage_contract_test.go new file mode 100644 index 000000000..2f2075544 --- /dev/null +++ b/internal/provider/usage_contract_test.go @@ -0,0 +1,391 @@ +package provider + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + + "github.com/Agent-Field/agentfield/sdk/go/ai" + "github.com/Agent-Field/codeaf/internal/home" + "github.com/Agent-Field/codeaf/internal/telemetry" +) + +func TestK2StreamUsageOptInFollowsTheActualRoad(t *testing.T) { + for _, tc := range []struct { + name, base string + direct, want bool + }{ + {"router", "", false, false}, + {"connected", "", true, true}, + {"connected overflow to router", "https://openrouter.ai/api/v1", true, false}, + {"router overflow to connected", "https://vendor.example/v1", false, true}, + } { + t.Run(tc.name, func(t *testing.T) { + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: tc.direct}) + if err != nil { + t.Fatal(err) + } + knobs := knobsFrom(context.Background()) + knobs.trace.attemptBase = tc.base + request, err := client.newRequest(userMessages("hello"), []ai.Option{ai.WithStream()}) + if err != nil { + t.Fatal(err) + } + encoded, err := client.encodeRequest(request, knobs) + if err != nil { + t.Fatal(err) + } + var body map[string]json.RawMessage + if err := json.Unmarshal(encoded, &body); err != nil { + t.Fatal(err) + } + if got := string(body["stream_options"]); (got != "") != tc.want || tc.want && got != `{"include_usage":true}` { + t.Fatalf("stream_options=%s, want present=%t", got, tc.want) + } + if !tc.direct && tc.base == "" && string(body["usage"]) != `{"include":true}` { + t.Fatalf("router usage=%s", body["usage"]) + } + }) + } +} + +func TestK2StreamUsageRefusalIsRememberedOn400And422(t *testing.T) { + for _, status := range []int{http.StatusBadRequest, http.StatusUnprocessableEntity} { + t.Run(fmt.Sprint(status), func(t *testing.T) { + requested, accepted := 0, 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + if body["stream_options"] != nil { + requested++ + w.WriteHeader(status) + fmt.Fprint(w, `{"detail":[{"loc":["body","stream_options"],"msg":"Extra inputs are not permitted"}]}`) + return + } + accepted++ + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://legacy/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + if _, err := client.CompleteWithMessages(context.Background(), userMessages("hello")); err != nil { + t.Fatal(err) + } + } + if requested != 1 || accepted != 2 { + t.Fatalf("requests with option=%d, accepted=%d; want 1, 2", requested, accepted) + } + }) + } +} + +func TestK3UnrelatedDirectRefusalKeepsTheDevErrorAndRequestCount(t *testing.T) { + for _, status := range []int{http.StatusBadRequest, http.StatusUnprocessableEntity} { + t.Run(fmt.Sprint(status), func(t *testing.T) { + const payload = `{"error":{"message":"invalid tool schema sentinel"}}` + requests := 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + if body["stream_options"] == nil { + t.Fatal("unrelated refusal disabled stream usage") + } + w.WriteHeader(status) + fmt.Fprint(w, payload) + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + _, err := client.CompleteWithMessages(context.Background(), userMessages("hello")) + if err == nil || err.Error() != apiError(status, []byte(payload)).Error() { + t.Fatalf("error=%v, want %v", err, apiError(status, []byte(payload))) + } + } + if requests != 2 { + t.Fatalf("requests=%d, want one per call as on dev", requests) + } + }) + } +} + +func TestK11RawStreamEndsAndBanksExactlyOnce(t *testing.T) { + for _, mode := range []string{"completed", "cut measured", "cut missing", "cancelled"} { + t.Run(mode, func(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, mode) + defer finish() + read := loggingTo(t) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + usage := `,"usage":{"prompt_tokens":10,"completion_tokens":3,"total_tokens":13,"cost":0.01}` + if mode == "cut missing" { + usage = "" + } + fmt.Fprintf(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]%s}\n\n", usage) + if mode == "completed" { + fmt.Fprint(w, "data: [DONE]\n\n") + } + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + bills, receipts := 0, 0 + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ctx = WithBilling(ctx, func(b Billed) { + bills++ + if b.PromptTokens != 10 || b.CompletionTokens != 3 || b.Cost != 0.01 { + t.Errorf("bill=%+v", b) + } + }) + ctx = WithReconcile(ctx, func(Reconciled) { receipts++ }) + before := telemetry.Snapshot() + chunks, errs := client.StreamComplete(ctx, "hello") + for range chunks { + if mode == "cancelled" { + cancel() + } + } + var failures []error + for err := range errs { + if err != nil { + failures = append(failures, err) + } + } + if mode == "completed" && len(failures) != 0 || mode != "completed" && len(failures) != 1 { + t.Fatalf("errors=%v", failures) + } + for _, err := range failures { + if strings.Contains(err.Error(), "panic") { + t.Fatal(err) + } + } + wantBills := 1 + if mode == "cut missing" { + wantBills = 0 + } + if bills != wantBills || receipts != 0 { + t.Fatalf("bills=%d receipts=%d", bills, receipts) + } + if calls := telemetry.Snapshot().ModelCalls - before.ModelCalls; calls != 1 { + t.Fatalf("model calls=%d, want one ending", calls) + } + rows := read() + ends := ended(rows) + if len(rows) != 2 || len(ends) != 1 || rows[0].ID != ends[0].ID { + t.Fatalf("rows=%+v", rows) + } + if rows := telemetry.SpoolContents(); len(rows) != 1 { + t.Fatalf("usage events=%d, want one ending", len(rows)) + } + }) + } +} + +func TestK9MissingUsageEstimatesDoNotEnterBillingOrCallLogCounts(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeTask, "estimate-only") + defer finish() + read := loggingTo(t) + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"observed answer\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + bills, receipts := 0, 0 + ctx := WithBilling(context.Background(), func(Billed) { bills++ }) + ctx = WithReconcile(WithUnmeteredReceipts(ctx), func(Reconciled) { receipts++ }) + response, err := client.CompleteWithMessages(ctx, userMessages("hello")) + if err != nil { + t.Fatal(err) + } + if response.Usage != nil || bills != 0 || receipts != 0 { + t.Fatalf("response usage=%+v bills=%d receipts=%d", response.Usage, bills, receipts) + } + rows := ended(read()) + if len(rows) != 1 || rows[0].PromptTokens != 0 || rows[0].CompletionTokens != 0 || rows[0].Cost != 0 { + t.Fatalf("call log=%+v", rows) + } + var event struct { + Props map[string]any `json:"props"` + } + spool := telemetry.SpoolContents() + if len(spool) != 1 { + t.Fatalf("events=%d", len(spool)) + } + if err := json.Unmarshal(spool[0], &event); err != nil { + t.Fatal(err) + } + if event.Props["usage_status"] != "estimated" || event.Props["estimated_total_tokens"] == nil || event.Props["total_tokens"] != nil { + t.Fatalf("usage=%+v", event.Props) + } +} + +func TestK7FinishedEmptyOutputIsObservedAsZero(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "partial") + defer finish() + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10}}\n\ndata: [DONE]\n\n") + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + chunks, errs := client.StreamComplete(context.Background(), "hello") + for range chunks { + } + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + rows := telemetry.SpoolContents() + if len(rows) != 1 { + t.Fatalf("events=%d", len(rows)) + } + var event struct { + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &event); err != nil { + t.Fatal(err) + } + p := event.Props + if p["usage_status"] != "mixed" || p["input_tokens"] != float64(10) || p["total_tokens"] != float64(10) || p["output_tokens"] != nil || p["estimated_output_tokens"] != float64(0) { + t.Fatalf("usage=%+v", p) + } +} + +func TestK11RawCutRecoversOneLateReceipt(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "raw-recovery") + defer finish() + read := loggingTo(t) + requests, lookups := 0, 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/generation") { + lookups++ + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"data":{"tokens_prompt":10,"tokens_completion":3,"total_cost":0.01}}`) + return + } + requests++ + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"id\":\"raw-recovery-private\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n") + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "https://openrouter.ai/api/v1", Model: "stub/model", HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + results := make(chan Reconciled, 2) + bills := 0 + ctx := WithReconcile(WithBilling(context.Background(), func(Billed) { bills++ }), func(r Reconciled) { results <- r }) + before := telemetry.Snapshot() + chunks, errs := client.StreamComplete(ctx, "hello") + for range chunks { + } + failures := 0 + for err := range errs { + if err != nil { + failures++ + } + } + result := receiptResult(t, results) + if !result.Found || result.PromptTokens != 10 || result.CompletionTokens != 3 || result.Cost != 0.01 || failures != 1 || bills != 0 || requests != 1 || lookups != 1 { + t.Fatalf("result=%+v failures=%d bills=%d requests=%d lookups=%d", result, failures, bills, requests, lookups) + } + select { + case extra := <-results: + t.Fatalf("duplicate settlement: %+v", extra) + default: + } + if rows := read(); len(rows) != 2 || len(ended(rows)) != 1 { + t.Fatalf("call log=%+v", rows) + } + if calls := telemetry.Snapshot().ModelCalls - before.ModelCalls; calls != 1 { + t.Fatalf("model calls=%d", calls) + } + rows := telemetry.SpoolContents() + if len(rows) != 2 { + t.Fatalf("usage events=%d", len(rows)) + } + var first, second struct { + ID string `json:"event_id"` + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &first); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(rows[1], &second); err != nil { + t.Fatal(err) + } + if first.Props["usage_status"] != "estimated" || first.Props["estimate_scope"] != "observed_partial" || second.Props["usage_status"] != "reported" || first.Props["usage_receipt"] != second.Props["usage_receipt"] || first.ID == second.ID { + t.Fatalf("receipts=%+v %+v", first, second) + } +} + +func TestK10RepeatedProviderGenerationIDsKeepDistinctAnswerReceipts(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "reused-provider-id") + defer finish() + calls := 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprintf(w, "data: {\"id\":\"reused-id\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":%d,\"completion_tokens\":%d}}\n\ndata: [DONE]\n\n", calls+9, calls+2) + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "stub/model", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + if _, err := client.CompleteWithMessages(context.Background(), userMessages("hello")); err != nil { + t.Fatal(err) + } + } + rows := telemetry.SpoolContents() + if len(rows) != 2 { + t.Fatalf("events=%d", len(rows)) + } + var first, second struct { + ID string `json:"event_id"` + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &first); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(rows[1], &second); err != nil { + t.Fatal(err) + } + if first.Props["usage_status"] != "reported" || second.Props["usage_status"] != "reported" || first.Props["input_tokens"] != float64(10) || second.Props["input_tokens"] != float64(11) || first.Props["usage_receipt"] == second.Props["usage_receipt"] || first.ID == second.ID { + t.Fatalf("receipts=%+v %+v", first, second) + } +} diff --git a/internal/provider/usage_estimate.go b/internal/provider/usage_estimate.go new file mode 100644 index 000000000..fd0b87dc8 --- /dev/null +++ b/internal/provider/usage_estimate.go @@ -0,0 +1,331 @@ +package provider + +import ( + "encoding/json" + "regexp" + "strings" + + "github.com/Agent-Field/agentfield/sdk/go/ai" + "github.com/Agent-Field/codeaf/internal/modelsource" + "github.com/Agent-Field/codeaf/internal/telemetry" +) + +// usageObservation retains counts and tool-name presence for one attempt. Estimates never +// enter the billing ledger, model response, price limits, or provider totals. +type usageObservation struct { + inputBytes int + inputTextOnly bool + outputBytes int + outputObserved bool + usageParsed bool + inputReported bool + outputReported bool + totalReported bool + toolOrder []int + toolNames map[int]bool +} + +// UnmarshalJSON preserves absent versus explicit zero counts. Cost-only usage +// and a one-sided receipt must not turn an unknown token count into zero. +func (u *usageWire) UnmarshalJSON(data []byte) error { + type plain usageWire + var decoded plain + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + *u = usageWire(decoded) + var fields map[string]json.RawMessage + if err := json.Unmarshal(data, &fields); err != nil { + return err + } + read := func(names ...string) (int, bool) { + for _, name := range names { + raw, exists := fields[name] + if !exists || string(raw) == "null" { + continue + } + var value int + if json.Unmarshal(raw, &value) == nil && value >= 0 { + return value, true + } + } + return 0, false + } + if value, known := read("prompt_tokens", "input_tokens"); known { + u.PromptTokens, u.inputReported = value, true + // Native Anthropic input excludes its separately reported cache work. + if _, canonical := fields["prompt_tokens"]; !canonical { + u.PromptTokens += u.CacheReadInputTokens + u.CacheCreationInputTokens + } + } + if value, known := read("completion_tokens", "output_tokens"); known { + u.CompletionTokens, u.outputReported = value, true + } + if value, known := read("total_tokens"); known { + u.TotalTokens, u.totalReported = value, true + } + return nil +} + +func (o *usageObservation) observeUsage(u *usageWire) { + o.usageParsed = true + if u != nil { + o.inputReported = o.inputReported || u.inputReported + o.outputReported = o.outputReported || u.outputReported + o.totalReported = o.totalReported || u.totalReported + } +} + +func usageInputBytes(messages, tools []json.RawMessage, source []ai.Message) (int, bool) { + for _, message := range source { + for _, part := range message.Content { + if part.Type != "text" { + return 0, false + } + } + } + size := 0 + for _, message := range messages { + var fields map[string]json.RawMessage + if json.Unmarshal(message, &fields) != nil { + return 0, false + } + // Opaque sidecars do not reveal how much input the provider consumed. + // Decline the component rather than count ciphertext as textual work. + for field := range fields { + switch field { + case "role", "content", "name", "tool_call_id", "tool_calls", "reasoning", "reasoning_content", "reasoning_text": + default: + return 0, false + } + } + size += len(message) + } + for _, tool := range tools { + size += len(tool) + } + return size, len(messages) > 0 +} + +func (o *usageObservation) observeDelta(delta streamDelta) { + before := o.outputBytes + o.outputBytes += len(delta.Content) + // Some gateways duplicate the same thought in multiple aliases. Use one + // textual channel, not the JSON wrapper or opaque encrypted signatures. + o.outputBytes += max(len(delta.Reasoning), len(delta.ReasoningContent), len(delta.ReasoningText)) + for _, call := range delta.ToolCalls { + index := toolCallIndex(call, o.toolOrder) + if o.toolNames == nil { + o.toolNames = make(map[int]bool) + } + if _, known := o.toolNames[index]; !known { + o.toolOrder = append(o.toolOrder, index) + } + if !o.toolNames[index] && call.Function.Name != "" { + o.outputBytes += len(call.Function.Name) + o.toolNames[index] = true + } else if _, known := o.toolNames[index]; !known { + o.toolNames[index] = false + } + o.outputBytes += len(call.Function.Arguments) + } + o.outputObserved = o.outputObserved || o.outputBytes > before +} + +func (o *usageObservation) observeWhole(payload []byte) { + var wire struct { + Choices []struct { + Message json.RawMessage `json:"message"` + FinishReason *string `json:"finish_reason"` + } `json:"choices"` + } + if json.Unmarshal(payload, &wire) != nil { + return + } + for _, choice := range wire.Choices { + var message ai.Message + if json.Unmarshal(choice.Message, &message) != nil { + continue + } + before := o.outputBytes + for _, part := range message.Content { + if part.Type == "text" { + o.outputBytes += len(part.Text) + } + } + for _, call := range message.ToolCalls { + o.outputBytes += len(call.Function.Name) + len(call.Function.Arguments) + } + // Decode reasoning separately so a valid array of text parts in content + // does not prevent its visible reasoning channel from being counted. + var thought struct { + Reasoning string `json:"reasoning"` + Content string `json:"reasoning_content"` + Text string `json:"reasoning_text"` + } + if json.Unmarshal(choice.Message, &thought) == nil { + o.outputBytes += max(len(thought.Reasoning), len(thought.Content), len(thought.Text)) + } + finished := choice.FinishReason != nil && strings.TrimSpace(*choice.FinishReason) != "" && !strings.EqualFold(strings.TrimSpace(*choice.FinishReason), "error") + o.outputObserved = o.outputObserved || o.outputBytes > before || finished + } +} + +// approximateUsageTokens is a versioned heuristic, not a tokenizer. UTF-8 byte +// counts include message and tool serialization overhead on input and observed +// content, tool arguments and visible thinking on output. Non-text input is +// unknown; hidden reasoning, unseen output, and media are never invented. +func approximateUsageTokens(size int) int { + if size <= 0 { + return 0 + } + return 1 + (size-1)/4 +} + +func (c *Client) usageDimensions(facts recordFacts, attemptID, model string) telemetry.UsageDimensions { + base := facts.baseFor(c.config.BaseURL) + generation := "" + if facts.response != nil { + generation = facts.response.ID + } + id := usageReceiptID(base, generation, attemptID) + d := telemetry.UsageDimensions{RoutingProvider: c.usageRoutingProvider(base), ModelFamily: telemetry.ModelFamily(model), AccountingSource: "provider", ReceiptID: id} + o := usageObservation{} + if facts.knobs.trace != nil { + o = facts.knobs.trace.usageObservation + } + reported := facts.response != nil && facts.response.Usage != nil + if o.usageParsed { + d.MissingInput, d.MissingOutput = !o.inputReported, !o.outputReported + } else { + d.MissingInput, d.MissingOutput = !reported, !reported + } + // A provider may report only a total. Keep that measurement intact and do + // not invent an input/output split or add estimates to a known total. + if reported && o.totalReported && d.MissingInput && d.MissingOutput { + total := facts.response.Usage.TotalTokens + d.ReportedTotal = &total + d.UsageStatus = "total_only" + return d + } + // A request that was rejected or never accepted proves no consumed work. + // An interrupted 2xx stream may still contribute observed output, explicitly + // marked partial; it is never extrapolated to a hypothetical completed reply. + accepted := facts.status >= 200 && facts.status < 300 && facts.response != nil + if accepted { + if d.MissingInput && o.inputTextOnly { + estimate := approximateUsageTokens(o.inputBytes) + d.EstimatedInput = &estimate + } + if d.MissingOutput && o.outputObserved { + estimate := approximateUsageTokens(o.outputBytes) + d.EstimatedOutput = &estimate + } + } + if d.EstimatedInput != nil || d.EstimatedOutput != nil { + d.EstimateScope = "text_only" + if !o.inputTextOnly { + d.EstimateScope = "multimodal_input" + } + if facts.err != nil { + d.EstimateScope = "observed_partial" + } + } + switch { + case !d.MissingInput && !d.MissingOutput: + d.UsageStatus = "reported" + case d.EstimatedInput != nil || d.EstimatedOutput != nil: + d.UsageStatus = "estimated" + if !d.MissingInput || !d.MissingOutput { + d.UsageStatus = "mixed" + } + case !d.MissingInput || !d.MissingOutput: + d.UsageStatus = "partial" + default: + d.UsageStatus = "missing" + } + return d +} + +// usageReceiptID pairs an observation with its later receipt, while retaining +// the physical attempt's identity. Compatible servers can reuse generation ids +// for different answers, so the provider's id alone cannot identify a receipt. +func usageReceiptID(base, generation, attempt string) string { + generation = canonicalGenerationID(generation) + if generation == "" { + return attempt + } + return attempt + ":generation:" + base + ":" + generation +} + +func canonicalGenerationID(generation string) string { + return strings.TrimSpace(generation) +} + +// usageEndpoint follows the actual attempt when a billing road changed. +func (c *Client) usageEndpoint(knobs callKnobs) string { + if knobs.trace != nil && knobs.trace.attemptBase != "" { + return knobs.trace.attemptBase + } + return c.config.BaseURL +} + +func (c *Client) streamUsageSupported(knobs callKnobs) bool { + base := c.usageEndpoint(knobs) + direct := c.config.Direct + if base != c.config.BaseURL { + direct = telemetry.RoutingProvider(base) != modelsource.DefaultID + } + // The router already receives usage.include. Only a connected service + // needs the compatible API's opt-in, including after a billing-road move. + if !direct { + return false + } + _, refused := c.streamUsageUnsupported.Load(base) + return !refused +} + +// Only a refusal naming the option itself permits this repair. An unrelated +// validation error must not silently turn off accounting or duplicate work. +func refusesStreamUsage(body []byte) bool { + var payload any + if json.Unmarshal(body, &payload) != nil { + return streamUsageRefusal.MatchString(strings.ToLower(string(body))) + } + var refused func(any) bool + refused = func(value any) bool { + switch value := value.(type) { + case string: + return streamUsageRefusal.MatchString(strings.ToLower(value)) + case []any: + for _, item := range value { + if refused(item) { + return true + } + } + case map[string]any: + // Pydantic puts the rejected field and its explanation in the same + // detail object. A sibling error's location must not qualify it. + if loc, ok := value["loc"].([]any); ok { + for _, field := range loc { + if field == "stream_options" || field == "include_usage" { + if msg, ok := value["msg"].(string); ok && strings.EqualFold(msg, "Extra inputs are not permitted") { + return true + } + } + } + } + for _, key := range []string{"error", "message", "detail", "msg"} { + if refused(value[key]) { + return true + } + } + } + return false + } + return refused(payload) +} + +// The refusal must qualify the option directly; unrelated clauses cannot +// turn a supported option off merely by mentioning another unknown field. +var streamUsageRefusal = regexp.MustCompile(`(?:\b(?:unsupported|unrecognized|unknown|unexpected)(?: (?:request|body|parameter|argument|field|keyword|option|input|supplied)s?)*[\s:='"` + "`" + `]+(?:stream_options|include_usage)\b|\b(?:stream_options|include_usage)['"` + "`" + `]*(?: (?:is|are|was))? (?:not supported|unsupported|not permitted|unrecognized|unknown|unexpected)\b|\bextra (?:inputs|fields) (?:are |is )?not permitted[\s:='"` + "`" + `]+(?:stream_options|include_usage)\b)`) diff --git a/internal/provider/usage_estimate_test.go b/internal/provider/usage_estimate_test.go new file mode 100644 index 000000000..42e42f560 --- /dev/null +++ b/internal/provider/usage_estimate_test.go @@ -0,0 +1,303 @@ +package provider + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + + "github.com/Agent-Field/agentfield/sdk/go/ai" + "github.com/Agent-Field/codeaf/internal/home" + "github.com/Agent-Field/codeaf/internal/telemetry" +) + +func TestUsageFallbackWorksAcrossCompatibleProviders(t *testing.T) { + for _, source := range []string{"ollama", "vllm", "lmstudio", "llamacpp", "custom", "openai", "openrouter"} { + for _, streamed := range []bool{true, false} { + t.Run(fmt.Sprintf("%s/stream=%t", source, streamed), func(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "fallback") + defer finish() + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if streamed { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, `data: {"id":"generation-private","choices":[{"index":0,"delta":{"content":"hello world"},"finish_reason":"stop"}]}`+"\n\ndata: [DONE]\n\n") + } else { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"id":"generation-private","choices":[{"index":0,"message":{"role":"assistant","content":"hello world","reasoning":"visible thought"},"finish_reason":"stop"}]}`) + } + }) + client, err := NewClient(Config{SourceID: source, APIKey: "test", BaseURL: "http://private-host/v1", Model: "qwen3", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if streamed { + ctx = WithStreamObserver(ctx, func(StreamEvent) {}) + } + if _, err := client.CompleteWithMessages(ctx, userMessages("hello")); err != nil { + t.Fatal(err) + } + rows := telemetry.SpoolContents() + if len(rows) != 1 { + t.Fatalf("rows=%d", len(rows)) + } + var event struct { + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &event); err != nil { + t.Fatal(err) + } + if event.Props["usage_status"] != "estimated" || event.Props["estimated_input_tokens"].(float64) <= 0 || event.Props["estimated_output_tokens"].(float64) <= 0 { + t.Fatal(event.Props) + } + if _, exists := event.Props["total_tokens"]; exists { + t.Fatal("estimate became measured usage") + } + for _, secret := range []string{"hello world", "private-host", "generation-private"} { + if strings.Contains(string(rows[0]), secret) { + t.Fatalf("private value escaped: %s", secret) + } + } + }) + } + } +} + +func TestUsagePresencePreservesZeroAliasesAndPartialReceipts(t *testing.T) { + for _, tc := range []struct { + body string + input, output bool + prompt, completion int + }{ + {`{}`, false, false, 0, 0}, {`{"cost":0.1}`, false, false, 0, 0}, + {`{"total_tokens":12}`, false, false, 0, 0}, + {`{"prompt_tokens":0,"completion_tokens":0}`, true, true, 0, 0}, + {`{"prompt_tokens":7}`, true, false, 7, 0}, {`{"output_tokens":3}`, false, true, 0, 3}, + {`{"input_tokens":7,"output_tokens":3,"cache_read_input_tokens":5,"cache_creation_input_tokens":2}`, true, true, 14, 3}, + {`{"prompt_tokens":7,"completion_tokens":3,"cache_read_input_tokens":5}`, true, true, 7, 3}, + {`{"prompt_tokens":null,"completion_tokens":2}`, false, true, 0, 2}, + } { + var usage usageWire + if err := json.Unmarshal([]byte(tc.body), &usage); err != nil { + t.Fatal(err) + } + if usage.inputReported != tc.input || usage.outputReported != tc.output || usage.PromptTokens != tc.prompt || usage.CompletionTokens != tc.completion { + t.Fatalf("%s: %+v", tc.body, usage) + } + } +} + +func TestUsageEstimatesKeepComponentsAndPartialScopeHonest(t *testing.T) { + client := &Client{config: Config{BaseURL: "http://local/v1"}} + for _, tc := range []struct { + name string + o usageObservation + usage *ai.Usage + err error + want string + estimateIn, estimateOut bool + }{ + {"partial", usageObservation{inputBytes: 80, inputTextOnly: true, outputBytes: 20, outputObserved: true, usageParsed: true, inputReported: true}, &ai.Usage{PromptTokens: 11}, nil, "mixed", false, true}, + {"cost_only", usageObservation{inputBytes: 80, inputTextOnly: true, outputBytes: 20, outputObserved: true, usageParsed: true}, &ai.Usage{}, nil, "estimated", true, true}, + {"zero", usageObservation{usageParsed: true, inputReported: true, outputReported: true}, &ai.Usage{}, nil, "reported", false, false}, + {"media", usageObservation{outputBytes: 20, outputObserved: true, usageParsed: true}, nil, nil, "estimated", false, true}, + {"cut", usageObservation{inputBytes: 80, inputTextOnly: true, outputBytes: 20, outputObserved: true, usageParsed: true}, nil, context.Canceled, "estimated", true, true}, + } { + t.Run(tc.name, func(t *testing.T) { + facts := recordFacts{status: 200, response: &ai.Response{Usage: tc.usage}, err: tc.err, knobs: callKnobs{trace: &callTrace{usageObservation: tc.o}}} + d := client.usageDimensions(facts, "attempt", "qwen3") + if d.UsageStatus != tc.want || (d.EstimatedInput != nil) != tc.estimateIn || (d.EstimatedOutput != nil) != tc.estimateOut { + t.Fatalf("%+v", d) + } + if tc.name == "cut" && d.EstimateScope != "observed_partial" { + t.Fatal(d) + } + if tc.name == "media" && d.EstimateScope != "multimodal_input" { + t.Fatal(d) + } + facts.status = 400 + d = client.usageDimensions(facts, "attempt", "qwen3") + if d.EstimatedInput != nil || d.EstimatedOutput != nil { + t.Fatal("refused request invented usage") + } + }) + } +} + +func TestUsageEstimateIncludesToolsAndVisibleReasoning(t *testing.T) { + o := usageObservation{} + o.observeDelta(streamDelta{Content: "abcd", Reasoning: "abcdefgh", ReasoningContent: "abcdefgh", ToolCalls: []toolCallDelta{{Function: toolFunctionDelta{Name: "read", Arguments: `{"x":1}`}}}}) + if o.outputBytes != 4+8+4+7 { + t.Fatalf("output bytes=%d", o.outputBytes) + } + messages := []json.RawMessage{json.RawMessage(`{"role":"user","content":"hello"}`)} + size, known := usageInputBytes(messages, []json.RawMessage{json.RawMessage(`{"name":"read","schema":{"type":"object"}}`)}, userMessages("hello")) + if !known || size <= len(messages[0]) { + t.Fatal("tool schema missing from estimate") + } + source := userMessages("hello") + source[0].Content = append(source[0].Content, ai.ContentPart{Type: "image_url", ImageURL: &ai.ImageURLData{URL: "data:image/png;base64,secret"}}) + if _, known := usageInputBytes(messages, nil, source); known { + t.Fatal("image data estimated as text") + } +} + +func TestRawStreamUsageReachesTelemetry(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "raw") + defer finish() + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, `data: {"id":"raw","choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}`+"\n\ndata: [DONE]\n\n") + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://local/v1", Model: "qwen3", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + chunks, errs := client.StreamComplete(context.Background(), "hello") + for range chunks { + } + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + if rows := telemetry.SpoolContents(); len(rows) != 1 { + t.Fatalf("raw stream receipts=%d", len(rows)) + } +} + +func TestLateProviderReceiptSupersedesEstimateWithSameIdentity(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "recovery") + defer finish() + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if strings.HasSuffix(r.URL.Path, "/generation") { + fmt.Fprint(w, `{"data":{"tokens_prompt":10,"tokens_completion":3,"total_cost":0.01}}`) + return + } + fmt.Fprint(w, `{"id":"late-private","choices":[{"index":0,"message":{"role":"assistant","content":"hello world"},"finish_reason":"stop"}]}`) + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "https://openrouter.ai/api/v1", Model: "qwen3", HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + result := make(chan Reconciled, 1) + ctx := WithUnmeteredReceipts(WithReconcile(context.Background(), func(r Reconciled) { result <- r })) + if _, err := client.CompleteWithMessages(ctx, userMessages("hello")); err != nil { + t.Fatal(err) + } + if answer := receiptResult(t, result); !answer.Found { + t.Fatal("receipt not recovered") + } + rows := telemetry.SpoolContents() + if len(rows) != 2 { + t.Fatalf("rows=%d", len(rows)) + } + var first, second struct { + ID string `json:"event_id"` + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &first); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(rows[1], &second); err != nil { + t.Fatal(err) + } + if first.Props["usage_status"] != "estimated" || second.Props["usage_status"] != "reported" { + t.Fatal(first.Props, second.Props) + } + if first.Props["usage_receipt"] != second.Props["usage_receipt"] || first.ID == second.ID { + t.Fatal("estimate and receipt cannot be reconciled") + } + if first.Props["total_tokens"] != nil || second.Props["total_tokens"] != float64(13) { + t.Fatal("measured totals double counted") + } +} + +func TestUsageOptionRefusalFallsBackWithoutBreakingOlderServers(t *testing.T) { + t.Setenv(home.EnvVar, t.TempDir()) + telemetry.EnableForTest(t, true) + telemetry.VersionForTest(t, "v0.7.0") + finish := telemetry.BeginUsageSession(telemetry.ModeChat, "legacy") + defer finish() + requested, accepted := 0, 0 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + if body["stream_options"] != nil { + requested++ + w.WriteHeader(400) + fmt.Fprint(w, `{"error":{"message":"unknown field stream_options"}}`) + return + } + accepted++ + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, `data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}`+"\n\ndata: [DONE]\n\n") + }) + client, err := NewClient(Config{APIKey: "test", BaseURL: "http://legacy/v1", Model: "qwen3", Direct: true, HTTPClient: handlerClient(handler)}) + if err != nil { + t.Fatal(err) + } + ctx := WithStreamObserver(context.Background(), func(StreamEvent) {}) + for i := 0; i < 2; i++ { + if _, err := client.CompleteWithMessages(ctx, userMessages("hello")); err != nil { + t.Fatal(err) + } + } + if requested != 1 || accepted != 2 { + t.Fatalf("opt-in refusals=%d, accepted=%d", requested, accepted) + } + rows := telemetry.SpoolContents() + if len(rows) != 3 { + t.Fatalf("receipts=%d", len(rows)) + } + var refused struct { + Props map[string]any `json:"props"` + } + if err := json.Unmarshal(rows[0], &refused); err != nil { + t.Fatal(err) + } + if refused.Props["estimated_total_tokens"] != nil { + t.Fatal("400 request counted as consumption") + } + if refusesStreamUsage([]byte(`{"error":"stream_options is valid but tool schema is invalid"}`)) { + t.Fatal("unrelated refusal disables usage") + } +} + +func TestUsageTotalOnlyReceiptIsKeptWithoutInventingASplit(t *testing.T) { + for _, total := range []int{0, 12} { + var wire usageWire + if err := json.Unmarshal([]byte(fmt.Sprintf(`{"total_tokens":%d}`, total)), &wire); err != nil { + t.Fatal(err) + } + o := usageObservation{inputBytes: 80, inputTextOnly: true, outputBytes: 20, outputObserved: true} + o.observeUsage(&wire) + client := &Client{config: Config{BaseURL: "http://local/v1"}} + d := client.usageDimensions(recordFacts{status: 200, response: &ai.Response{Usage: wire.usage()}, knobs: callKnobs{trace: &callTrace{usageObservation: o}}}, "attempt", "qwen3") + if d.UsageStatus != "total_only" || d.ReportedTotal == nil || *d.ReportedTotal != total || d.EstimatedInput != nil || d.EstimatedOutput != nil { + t.Fatal(d) + } + } +} + +func TestUsageWholeResponseCountsTextPartsWithoutCountingMedia(t *testing.T) { + o := usageObservation{} + o.observeWhole([]byte(`{"choices":[{"message":{"role":"assistant","content":[{"type":"text","text":"hello"},{"type":"image_url","image_url":{"url":"data:image/png;base64,not-tokens"}}],"reasoning":"thought","tool_calls":[{"id":"call","type":"function","function":{"name":"read","arguments":"{}"}}]}}]}`)) + if !o.outputObserved || o.outputBytes != 5+7+4+2 { + t.Fatalf("observation=%+v", o) + } +} diff --git a/internal/provider/wire.go b/internal/provider/wire.go index 48193e94c..517b9baa1 100644 --- a/internal/provider/wire.go +++ b/internal/provider/wire.go @@ -261,7 +261,7 @@ func refusesDisabledReasoning(payload []byte) bool { // wireRequest is the serialized body. It shadows the SDK's max_tokens so the // adapter, not the SDK, decides which output-limit field a given endpoint gets, -// and adds the three fields the SDK's Request type has no home for. +// and adds fields the SDK's Request type has no home for. type wireRequest struct { *requestAlias @@ -279,6 +279,10 @@ type wireRequest struct { Messages []json.RawMessage `json:"messages"` Tools []json.RawMessage `json:"tools,omitempty"` + // StreamOptions asks compatible servers for their terminal token receipt. + // OpenRouter's usage.include alone does not opt local APIs into this frame. + StreamOptions *streamOptions `json:"stream_options,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` MaxCompletionTokens *int `json:"max_completion_tokens,omitempty"` @@ -302,6 +306,10 @@ type wireRequest struct { Provider *providerPrefs `json:"provider,omitempty"` } +type streamOptions struct { + IncludeUsage bool `json:"include_usage"` +} + type requestAlias ai.Request // encodeRequest applies outbound hygiene and the economy fields, then @@ -385,6 +393,9 @@ func (c *Client) encodeRequest(request *ai.Request, knobs callKnobs) ([]byte, er Tools: tools, PromptCacheKey: knobs.cacheKey, } + if scrubbed.Stream && c.streamUsageSupported(knobs) { + wire.StreamOptions = &streamOptions{IncludeUsage: true} + } // The knob is decided here and written after the context budget below, // which may shrink the thinking budget to what the window leaves. sentEffort, thinking := EffortNone, 0 @@ -442,7 +453,13 @@ func (c *Client) encodeRequest(request *ai.Request, knobs callKnobs) ([]byte, er wire.MaxTokens = &ceiling } } - return json.Marshal(wire) + encoded, err := json.Marshal(wire) + if err == nil && knobs.trace != nil { + // The next attempt's input count is taken from the encoded messages and + // tools, after repairs and cache annotations, without retaining the body. + knobs.trace.usageObservation.inputBytes, knobs.trace.usageObservation.inputTextOnly = usageInputBytes(messages, tools, scrubbed.Messages) + } + return encoded, err } // effortAsked is what this call really asks for after the client's own pinned diff --git a/internal/telemetry/dimensions.go b/internal/telemetry/dimensions.go index 8f2acb967..20b0bd0c2 100644 --- a/internal/telemetry/dimensions.go +++ b/internal/telemetry/dimensions.go @@ -16,6 +16,12 @@ type UsageDimensions struct { UsageStatus string AccountingSource string ReceiptID string + MissingInput bool + MissingOutput bool + EstimatedInput *int + EstimatedOutput *int + EstimateScope string + ReportedTotal *int } // category refuses arbitrary caller strings even when a caller bypasses the @@ -84,23 +90,89 @@ func UsageReceipt(mode Mode, input, output int, sessionID string, now time.Time, event := UsageDelta(mode, input, output, sessionID, now) event.Props["routing_provider"] = category(dimensions.RoutingProvider, "openrouter openai anthropic google bedrock azure ollama other unknown") event.Props["model_family"] = category(dimensions.ModelFamily, "deepseek claude gpt gemini llama qwen kimi glm other unknown") - status := dimensions.UsageStatus - if status != "missing" { - status = "reported" - } - event.Props["usage_status"] = status source := dimensions.AccountingSource if source != "provider" { source = "engine" } event.Props["accounting_source"] = source - if status == "missing" { - delete(event.Props, "input_tokens") - delete(event.Props, "output_tokens") - delete(event.Props, "total_tokens") + // Status can describe absence for older callers, but the final status is + // derived from the safe components that actually survive serialization. + missingInput := dimensions.MissingInput || dimensions.UsageStatus == "missing" || dimensions.UsageStatus == "estimated" + missingOutput := dimensions.MissingOutput || dimensions.UsageStatus == "missing" || dimensions.UsageStatus == "estimated" + for _, key := range []string{"input_tokens", "output_tokens", "total_tokens"} { + delete(event.Props, key) + } + measured := [2]*int{} + if !missingInput && safeUsageCount(input) { + measured[0] = &input + } + if !missingOutput && safeUsageCount(output) { + measured[1] = &output + } + estimated := [2]*int{} + if dimensions.ReceiptID != "" { + for i, value := range [2]*int{dimensions.EstimatedInput, dimensions.EstimatedOutput} { + if measured[i] == nil && value != nil && safeUsageCount(*value) { + estimated[i] = value + } + } + } + // A pair whose sum is unsafe cannot be sent with a relay-valid total. + // Leave that group unknown rather than clamp either component to fiction. + measuredCount := emitUsageComponents(event.Props, "", measured) + estimatedCount := emitUsageComponents(event.Props, "estimated_", estimated) + status := "missing" + switch { + case measuredCount > 0 && estimatedCount > 0: + status = "mixed" + case measuredCount == 2: + status = "reported" + case measuredCount == 1: + status = "partial" + case estimatedCount > 0: + status = "estimated" + case dimensions.ReportedTotal != nil && safeUsageCount(*dimensions.ReportedTotal): + event.Props["total_tokens"] = *dimensions.ReportedTotal + status = "total_only" + } + event.Props["usage_status"] = status + if estimatedCount > 0 { + event.Props["estimate_method"] = "utf8_bytes_4_v1" + event.Props["estimate_scope"] = category(dimensions.EstimateScope, "text_only observed_partial multimodal_input") } if dimensions.ReceiptID != "" { - event.ID = hashHex("usage_delta|" + sessionID + "|" + dimensions.ReceiptID)[:32] + receipt := hashHex("usage_receipt|" + sessionID + "|" + dimensions.ReceiptID) + event.Props["usage_receipt"] = receipt + // An estimate and its later provider receipt are different events about + // one receipt. PostHog groups this hash and prefers measured components. + event.ID = hashHex("usage_delta|" + receipt + "|" + status)[:32] } return event } + +// safeUsageCount matches the relay's nonnegative JavaScript safe integers. +func safeUsageCount(value int) bool { + return value >= 0 && uint64(value) <= 1<<53-1 +} + +func emitUsageComponents(props map[string]any, prefix string, values [2]*int) int { + total, count := 0, 0 + for _, value := range values { + if value != nil { + if uint64(total)+uint64(*value) > 1<<53-1 { + return 0 + } + total += *value + count++ + } + } + if count > 0 { + for i, key := range []string{"input_tokens", "output_tokens"} { + if values[i] != nil { + props[prefix+key] = *values[i] + } + } + props[prefix+"total_tokens"] = total + } + return count +} diff --git a/internal/telemetry/dimensions_test.go b/internal/telemetry/dimensions_test.go index 4706e7fed..26a18057b 100644 --- a/internal/telemetry/dimensions_test.go +++ b/internal/telemetry/dimensions_test.go @@ -2,8 +2,11 @@ package telemetry import ( "encoding/json" + "os" + "reflect" "strings" "testing" + "time" ) func TestUsageReceiptMissingDoesNotInventZeroTokens(t *testing.T) { @@ -102,3 +105,92 @@ func TestLateUsageReceiptKeepsOriginalSessionAndCountsOnce(t *testing.T) { t.Fatal(row) } } + +func TestUsageEstimateAndRecoveredReceiptShareIdentity(t *testing.T) { + testHome(t) + estimate := 25 + initial := UsageReceipt(ModeChat, 0, 0, "session", testNow, UsageDimensions{UsageStatus: "estimated", EstimatedOutput: &estimate, EstimateScope: "observed_partial", ReceiptID: "generation:private"}) + recovered := UsageReceipt(ModeChat, 10, 3, "session", testNow, UsageDimensions{UsageStatus: "reported", ReceiptID: "generation:private"}) + if initial.Props["usage_receipt"] != recovered.Props["usage_receipt"] || initial.ID == recovered.ID { + t.Fatal("estimate and recovery cannot be reconciled") + } + if initial.Props["total_tokens"] != nil || initial.Props["estimated_total_tokens"] != 25 { + t.Fatal(initial.Props) + } +} + +func TestUsageRelayFixtureMatchesClientConstructors(t *testing.T) { + testHome(t) + data, err := os.ReadFile("testdata/usage-accounting.json") + if err != nil { + t.Fatal(err) + } + var fixture struct { + Events []jsonEvent `json:"events"` + } + if err := json.Unmarshal(data, &fixture); err != nil { + t.Fatal(err) + } + estimateIn, estimateOut := 100, 25 + dimensions := []UsageDimensions{ + {UsageStatus: "estimated", EstimatedInput: &estimateIn, EstimatedOutput: &estimateOut, EstimateScope: "text_only"}, + {UsageStatus: "reported"}, + {UsageStatus: "mixed", MissingOutput: true, EstimatedOutput: &estimateOut, EstimateScope: "text_only"}, + {UsageStatus: "partial", MissingOutput: true}, + {UsageStatus: "missing"}, + } + counts := [][2]int{{0, 0}, {90, 20}, {90, 0}, {90, 0}, {0, 0}} + for i, d := range dimensions { + want := fixture.Events[i] + d.ReceiptID = "generation:contract" + d.RoutingProvider = "other" + d.ModelFamily = "qwen" + d.AccountingSource = "provider" + at, err := time.Parse(time.RFC3339, want.EventTime) + if err != nil { + t.Fatal(err) + } + event := UsageReceipt(ModeChat, counts[i][0], counts[i][1], "contract", at, d) + event.InstallHash = want.InstallIDHash + for _, key := range commonPropNames { + event.Props[key] = want.Props[key] + } + actual, err := json.Marshal(event) + if err != nil { + t.Fatal(err) + } + var got map[string]any + var expected map[string]any + if err := json.Unmarshal(actual, &got); err != nil { + t.Fatal(err) + } + raw, err := json.Marshal(want) + if err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(raw, &expected); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(got, expected) { + t.Fatalf("%s: got %s, want %s", d.UsageStatus, actual, raw) + } + } +} + +func TestUsageEnrichedBatchFitsRelayBodyLimit(t *testing.T) { + testHome(t) + estimate := int(1<<53 - 1) + event := UsageReceipt(ModeChat, estimate, 0, "contract", testNow, UsageDimensions{UsageStatus: "mixed", MissingOutput: true, EstimatedOutput: &estimate, EstimateScope: "observed_partial", ReceiptID: "generation:contract", RoutingProvider: "anthropic", ModelFamily: "deepseek", AccountingSource: "provider"}) + event.Props["codeaf_version"] = strings.Repeat("v", 64) + events := make([]Event, MaxEventsPerPOST) + for i := range events { + events[i] = event + } + body, err := json.Marshal(map[string]any{"schema_version": 1, "events": events}) + if err != nil { + t.Fatal(err) + } + if len(body) > 64*1024 { + t.Fatalf("enriched batch %d bytes exceeds relay cap", len(body)) + } +} diff --git a/internal/telemetry/events.go b/internal/telemetry/events.go index 2c687294f..45751b07d 100644 --- a/internal/telemetry/events.go +++ b/internal/telemetry/events.go @@ -331,7 +331,7 @@ var allowedProps = map[string]map[string]bool{ "tool_calls", "tool_calls_failed", "cost_usd", "stop_reason", "exit_code", })), "usage_delta": merge(set(commonPropNames), set([]string{ - "mode", "input_tokens", "output_tokens", "total_tokens", "routing_provider", "model_family", "usage_status", "accounting_source", + "mode", "input_tokens", "output_tokens", "total_tokens", "routing_provider", "model_family", "usage_status", "accounting_source", "estimated_input_tokens", "estimated_output_tokens", "estimated_total_tokens", "estimate_method", "estimate_scope", "usage_receipt", })), "fault": merge(set(commonPropNames), set([]string{"mode", "scope", "fingerprint"})), } @@ -373,14 +373,20 @@ var propDocs = map[string]map[string]string{ "exit_code": "0 to 5", }, "usage_delta": { - "mode": "chat or task", - "input_tokens": "provider-reported input tokens since the preceding usage event", - "output_tokens": "provider-reported output tokens since the preceding usage event", - "total_tokens": "the sum of this event's input and output tokens", - "routing_provider": "openrouter, openai, anthropic, google, bedrock, azure, ollama, other, or unknown", - "model_family": "deepseek, claude, gpt, gemini, llama, qwen, kimi, glm, other, or unknown", - "usage_status": "reported or missing", - "accounting_source": "provider or engine", + "mode": "chat or task", + "input_tokens": "provider-reported input tokens since the preceding usage event", + "output_tokens": "provider-reported output tokens since the preceding usage event", + "total_tokens": "the sum of this event's measured components, or the provider total with no split for total_only", + "routing_provider": "openrouter, openai, anthropic, google, bedrock, azure, ollama, other, or unknown", + "model_family": "deepseek, claude, gpt, gemini, llama, qwen, kimi, glm, other, or unknown", + "usage_status": "reported, total_only, mixed, estimated, partial, or missing", + "accounting_source": "provider or engine", + "estimated_input_tokens": "estimated input tokens for a missing provider count", + "estimated_output_tokens": "estimated observed output tokens for a missing provider count", + "estimated_total_tokens": "the sum of this event's estimated input and output tokens", + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only, observed_partial, multimodal_input, or unknown", + "usage_receipt": "sha256 of the session and local receipt identity; groups estimates with later provider receipts", }, "fault": { "mode": "chat, task, or other", @@ -403,7 +409,7 @@ func EventPropNames(event string) []string { return []string{"mode", "duration", "turns", "model_calls", "model_calls_failed", "tool_calls", "tool_calls_failed", "cost_usd", "stop_reason", "exit_code"} case "usage_delta": - return []string{"mode", "input_tokens", "output_tokens", "total_tokens", "routing_provider", "model_family", "usage_status", "accounting_source"} + return []string{"mode", "input_tokens", "output_tokens", "total_tokens", "routing_provider", "model_family", "usage_status", "accounting_source", "estimated_input_tokens", "estimated_output_tokens", "estimated_total_tokens", "estimate_method", "estimate_scope", "usage_receipt"} case "fault": return []string{"mode", "scope", "fingerprint"} } @@ -455,14 +461,20 @@ var exampleProps = map[string]map[string]string{ "exit_code": "0", }, "usage_delta": { - "mode": string(ModeChat), - "input_tokens": "10000", - "output_tokens": "2500", - "total_tokens": "12500", - "routing_provider": modelsource.DefaultID, - "model_family": "deepseek", - "usage_status": "reported", - "accounting_source": "provider", + "mode": string(ModeChat), + "input_tokens": "10000", + "output_tokens": "2500", + "total_tokens": "12500", + "routing_provider": modelsource.DefaultID, + "model_family": "deepseek", + "usage_status": "reported", + "accounting_source": "provider", + "estimated_input_tokens": "100", + "estimated_output_tokens": "25", + "estimated_total_tokens": "125", + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only", + "usage_receipt": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", }, "fault": { "mode": string(ModeChat), diff --git a/internal/telemetry/events_test.go b/internal/telemetry/events_test.go index 752ee92f1..768480201 100644 --- a/internal/telemetry/events_test.go +++ b/internal/telemetry/events_test.go @@ -283,10 +283,25 @@ func TestAllowlistedPropsMatchTheConstructors(t *testing.T) { for _, name := range AllowlistedProps(event.Name) { allowed[name] = true } - if len(event.Props) != len(allowed) { - t.Errorf("%s: the constructor emits %d props, the allowlist names %d", event.Name, len(event.Props), len(allowed)) - } + emitted := map[string]bool{} for key := range event.Props { + emitted[key] = true + } + if event.Name == "usage_delta" { + estimate := 3 + optional := UsageReceipt(ModeChat, 10, 0, "session-1", testNow, UsageDimensions{UsageStatus: "mixed", MissingOutput: true, EstimatedOutput: &estimate, EstimateScope: "text_only", ReceiptID: "receipt"}) + for key := range optional.Props { + emitted[key] = true + } + optional = UsageReceipt(ModeChat, 0, 0, "session-1", testNow, UsageDimensions{UsageStatus: "estimated", EstimatedInput: &estimate, EstimateScope: "text_only", ReceiptID: "receipt"}) + for key := range optional.Props { + emitted[key] = true + } + } + if len(emitted) != len(allowed) { + t.Errorf("%s: the constructor emits %d props, the allowlist names %d", event.Name, len(emitted), len(allowed)) + } + for key := range emitted { if !allowed[key] { t.Errorf("%s: prop %q is not on the allowlist", event.Name, key) } diff --git a/internal/telemetry/testdata/usage-accounting.json b/internal/telemetry/testdata/usage-accounting.json new file mode 100644 index 000000000..edc2024a8 --- /dev/null +++ b/internal/telemetry/testdata/usage-accounting.json @@ -0,0 +1,556 @@ +{ + "schema_version": 1, + "events": [ + { + "event_name": "usage_delta", + "event_id": "854f13d5b621faef5441bddd81104162", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "estimated", + "accounting_source": "provider", + "usage_receipt": "ad9607bf6c6da3d5f760a755c629ce6f5b906f5d60bab2a6bc90abaab0abd8f7", + "estimated_input_tokens": 100, + "estimated_output_tokens": 25, + "estimated_total_tokens": 125, + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only" + } + }, + { + "event_name": "usage_delta", + "event_id": "3e6deae5146ab1328a8653f5f7ab78fb", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "reported", + "accounting_source": "provider", + "usage_receipt": "ad9607bf6c6da3d5f760a755c629ce6f5b906f5d60bab2a6bc90abaab0abd8f7", + "input_tokens": 90, + "output_tokens": 20, + "total_tokens": 110 + } + }, + { + "event_name": "usage_delta", + "event_id": "684797b3e1faaa72d6119d46a18e3b17", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "mixed", + "accounting_source": "provider", + "usage_receipt": "ad9607bf6c6da3d5f760a755c629ce6f5b906f5d60bab2a6bc90abaab0abd8f7", + "input_tokens": 90, + "total_tokens": 90, + "estimated_output_tokens": 25, + "estimated_total_tokens": 25, + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only" + } + }, + { + "event_name": "usage_delta", + "event_id": "cc04eea9220a21bc379d252644c0538b", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "partial", + "accounting_source": "provider", + "usage_receipt": "ad9607bf6c6da3d5f760a755c629ce6f5b906f5d60bab2a6bc90abaab0abd8f7", + "input_tokens": 90, + "total_tokens": 90 + } + }, + { + "event_name": "usage_delta", + "event_id": "4e30a1f1bde1dc4afd4c3e6f0043fa0d", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "ad9607bf6c6da3d5f760a755c629ce6f5b906f5d60bab2a6bc90abaab0abd8f7" + } + }, + { + "event_name": "usage_delta", + "event_id": "4da6fe4da9c42e811954aec6cff61714", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "reported", + "accounting_source": "provider", + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + "usage_receipt": "2c3f0188fe56f627853a8961719cc142db82759c4ba766ad6fa8419839e25b95" + } + }, + { + "event_name": "usage_delta", + "event_id": "99b01b9ff84e52798502c73ef57bfe26", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "partial", + "accounting_source": "provider", + "input_tokens": 0, + "total_tokens": 0, + "usage_receipt": "231f72fea8783a980b0f982d69c9c402ec2b2ea73c1c72e5240116576a394678" + } + }, + { + "event_name": "usage_delta", + "event_id": "624163f5959ab4268b620eb533e182b5", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "mixed", + "accounting_source": "provider", + "input_tokens": 0, + "estimated_output_tokens": 7, + "total_tokens": 0, + "estimated_total_tokens": 7, + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only", + "usage_receipt": "81c3762242dd5520fa48bfd2523e5ed61f7209f0a2a4291d2b6b14d15ff88eb0" + } + }, + { + "event_name": "usage_delta", + "event_id": "9a3a713cdb4e336fdcb203c9a2658274", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "estimated", + "accounting_source": "provider", + "estimated_output_tokens": 7, + "estimated_total_tokens": 7, + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only", + "usage_receipt": "604d0cd6aac37ac5d0803d7829b18e46b63fa3ed8a438dbed926a22de6b4362c" + } + }, + { + "event_name": "usage_delta", + "event_id": "5c8122049448fe1bf0e19ee5092e3d98", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "total_only", + "accounting_source": "provider", + "total_tokens": 0, + "usage_receipt": "c934ec9956c2807a429f885bac5c83820a5a81a6fc96a15e11331c40d748f335" + } + }, + { + "event_name": "usage_delta", + "event_id": "d055d68bfd611270dc946b64cfec3860", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "11b45a1add7de1722b849b058d5716665b33f03dc593eb7d75159420a2c05f1a" + } + }, + { + "event_name": "usage_delta", + "event_id": "9f0b8bc013adb319ceb44684d18dee74", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "79b0506666b335ea083c28157f99fdf6fbe8db68e6af01338aa8481d060aa431" + } + }, + { + "event_name": "usage_delta", + "event_id": "c1d7b96bc107d1e6af82f28f67fca257", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "partial", + "accounting_source": "provider", + "output_tokens": 7, + "total_tokens": 7, + "usage_receipt": "c73f3450d7c8451340e00ec6d8bfdfbfcbe52a31a55133c05db28802e6e1ec8f" + } + }, + { + "event_name": "usage_delta", + "event_id": "61394ab6a3fc89af69e603875b9ae7f2", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "c27ad57dda0f18089b916963d1db1169b2bc15b9dcc561e8c6af70bd50ad03f9" + } + }, + { + "event_name": "usage_delta", + "event_id": "4ff6d8484aa447441e2582c6dea1dfdc", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "783659bcba98f766acb4cd7250769f8573fcba39ed5a6cc57e823d300c8d6bd4" + } + }, + { + "event_name": "usage_delta", + "event_id": "1b71e721a90405ccb9fea619bd2cd2a5", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "partial", + "accounting_source": "provider", + "input_tokens": 7, + "total_tokens": 7, + "usage_receipt": "3a0faab1c5a034c784d1bcf5a538ba76742811c04437d3076639741b5f8b7f51" + } + }, + { + "event_name": "usage_delta", + "event_id": "174a5e478e0b0c65959524ee2b14de2e", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "estimated", + "accounting_source": "provider", + "estimated_output_tokens": 7, + "estimated_total_tokens": 7, + "estimate_method": "utf8_bytes_4_v1", + "estimate_scope": "text_only", + "usage_receipt": "97f510756fb3d558c764318d7980d2be8e3fb1cd8f2b84e2a71a9f3df909108f" + } + }, + { + "event_name": "usage_delta", + "event_id": "3592a8838038e95f48c225e88b0b6038", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "e50586ce2708da07b1eb74645b02bb71d20fc02c73f75d09023c709fe283ce1b" + } + }, + { + "event_name": "usage_delta", + "event_id": "8635d05eb7a8f3265444914d4a6c5e01", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "reported", + "accounting_source": "provider", + "input_tokens": 0, + "output_tokens": 7, + "total_tokens": 7, + "usage_receipt": "14de2b54d4965f8f6ee23b8527387acad5fc12f92a4f846d7a159ceda98ff619" + } + }, + { + "event_name": "usage_delta", + "event_id": "ffffffffffffffffffffffffffffffff", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider" + } + }, + { + "event_name": "usage_delta", + "event_id": "3a2d206cb579587e31b6c5e8e75663a8", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "ca6cd1c83a472d0a0dd29062eeeba3ace41f607aa6b0e81e4fc22fa1930e082b" + } + }, + { + "event_name": "usage_delta", + "event_id": "2842c8b282505f77045d98e6c5c3668b", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "missing", + "accounting_source": "provider", + "usage_receipt": "318d54203d43a2a72863fe3f6e6198b015dc7809b51fa777bdec3978c6bf34eb" + } + }, + { + "event_name": "usage_delta", + "event_id": "6126221bdd2c9a00fd9ef4b96e521153", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "reported", + "accounting_source": "provider", + "input_tokens": 9007199254740991, + "output_tokens": 0, + "total_tokens": 9007199254740991, + "usage_receipt": "fefd3795f6498d6efe089231c0d84c2b006c24965e608250d40217de31834bd4" + } + }, + { + "event_name": "usage_delta", + "event_id": "b936731df8a8652b1e74e7bbe4ae9cfe", + "install_id_hash": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "session_id_hash": "cc8321d6375c494d043fdd0260f21bc0ec51dacc9f6abb7f909cdcd3041b78bf", + "event_time": "2026-10-21T12:00:00Z", + "props": { + "codeaf_version": "v0.7.0", + "channel": "stable", + "os": "linux", + "arch": "amd64", + "usage_context": "local", + "install_method": "script", + "mode": "chat", + "routing_provider": "other", + "model_family": "qwen", + "usage_status": "partial", + "accounting_source": "provider", + "output_tokens": 7, + "total_tokens": 7, + "usage_receipt": "028affc987755cb6fc8fbd79ecc39c7c4f12d46830e4a5243b94adbcc1994066" + } + } + ] +} diff --git a/internal/telemetry/usage.go b/internal/telemetry/usage.go index 498f21afa..3839c2d2e 100644 --- a/internal/telemetry/usage.go +++ b/internal/telemetry/usage.go @@ -106,17 +106,28 @@ func StartPeriodicFlush(interval time.Duration) func() { // looking up the active session then would attribute its tokens to the wrong // work. This callback records at most once and still honors current opt-out. func CaptureUsageRecorder(dimensions UsageDimensions) func(int, int) { + record := CaptureUsageRecorderWithPresence(dimensions) + if record == nil { + return nil + } + return func(input, output int) { record(input, output, true, true) } +} + +// CaptureUsageRecorderWithPresence retains unknown components during recovery. +// A cost-only receipt updates money without replacing estimates with zeroes. +func CaptureUsageRecorderWithPresence(dimensions UsageDimensions) func(int, int, bool, bool) { session := activeUsageSession.Load() if session == nil || !enabledFor() { return nil } mode, id := session.mode, session.id var once sync.Once - return func(input, output int) { + return func(input, output int, inputKnown, outputKnown bool) { once.Do(func() { - if !enabledFor() { + if !enabledFor() || (!inputKnown && !outputKnown) { return } + dimensions.MissingInput, dimensions.MissingOutput = !inputKnown, !outputKnown event := UsageReceipt(mode, input, output, id, time.Now(), dimensions) if err := SpoolSync(event); err != nil { oneWarning("codeaf: usage counts could not be saved locally") diff --git a/internal/telemetry/usage_c3_contract_test.go b/internal/telemetry/usage_c3_contract_test.go new file mode 100644 index 000000000..044465cac --- /dev/null +++ b/internal/telemetry/usage_c3_contract_test.go @@ -0,0 +1,214 @@ +package telemetry + +import ( + "encoding/json" + "fmt" + "math" + "os" + "reflect" + "regexp" + "testing" + "time" +) + +func TestK13UsageReceiptEmitsOnlyRelayValidComponents(t *testing.T) { + testHome(t) + maxSafe, unsafe, maxInt := int(1<<53-1), int(1<<53), int(^uint(0)>>1) + zero, one, seven, negative := 0, 1, 7, -1 + cases := []struct { + name string + input, output int + d UsageDimensions + status string + measured, estimated [2]*int + total *int + }{ + {"reported_zero", 0, 0, UsageDimensions{UsageStatus: "reported"}, "reported", [2]*int{&zero, &zero}, [2]*int{}, nil}, + {"partial_zero", 0, 3, UsageDimensions{UsageStatus: "partial", MissingOutput: true}, "partial", [2]*int{&zero, nil}, [2]*int{}, nil}, + {"mixed_zero", 0, 0, UsageDimensions{UsageStatus: "mixed", MissingOutput: true, EstimatedOutput: &seven}, "mixed", [2]*int{&zero, nil}, [2]*int{nil, &seven}, nil}, + {"estimated_output", 0, 0, UsageDimensions{UsageStatus: "estimated", EstimatedOutput: &seven}, "estimated", [2]*int{}, [2]*int{nil, &seven}, nil}, + {"total_only_zero", 0, 0, UsageDimensions{UsageStatus: "total_only", MissingInput: true, MissingOutput: true, ReportedTotal: &zero}, "total_only", [2]*int{}, [2]*int{}, &zero}, + {"missing", 0, 0, UsageDimensions{UsageStatus: "missing"}, "missing", [2]*int{}, [2]*int{}, nil}, + {"negative_estimate", 0, 0, UsageDimensions{UsageStatus: "estimated", EstimatedOutput: &negative}, "missing", [2]*int{}, [2]*int{}, nil}, + {"negative_measurement", -1, 7, UsageDimensions{UsageStatus: "reported"}, "partial", [2]*int{nil, &seven}, [2]*int{}, nil}, + {"unsafe_estimate", 0, 0, UsageDimensions{UsageStatus: "estimated", EstimatedOutput: &unsafe}, "missing", [2]*int{}, [2]*int{}, nil}, + {"unsafe_measured_sum", maxSafe, 1, UsageDimensions{UsageStatus: "reported"}, "missing", [2]*int{}, [2]*int{}, nil}, + {"unknown_status_partial", 7, 0, UsageDimensions{UsageStatus: "bogus", MissingOutput: true}, "partial", [2]*int{&seven, nil}, [2]*int{}, nil}, + {"unknown_status_estimate", 0, 0, UsageDimensions{UsageStatus: "bogus", MissingInput: true, MissingOutput: true, EstimatedOutput: &seven}, "estimated", [2]*int{}, [2]*int{nil, &seven}, nil}, + {"overflowed_estimates", 0, 0, UsageDimensions{UsageStatus: "estimated", EstimatedInput: &maxInt, EstimatedOutput: &maxInt}, "missing", [2]*int{}, [2]*int{}, nil}, + {"malformed_total_only", 0, 7, UsageDimensions{UsageStatus: "total_only", ReportedTotal: &one}, "reported", [2]*int{&zero, &seven}, [2]*int{}, nil}, + {"estimate_without_receipt", 0, 0, UsageDimensions{UsageStatus: "estimated", EstimatedOutput: &seven}, "missing", [2]*int{}, [2]*int{}, nil}, + {"unsafe_estimated_sum", 0, 0, UsageDimensions{UsageStatus: "estimated", EstimatedInput: &maxSafe, EstimatedOutput: &one}, "missing", [2]*int{}, [2]*int{}, nil}, + {"unsafe_total_only", 0, 0, UsageDimensions{UsageStatus: "total_only", MissingInput: true, MissingOutput: true, ReportedTotal: &unsafe}, "missing", [2]*int{}, [2]*int{}, nil}, + {"safe_boundary", maxSafe, 0, UsageDimensions{UsageStatus: "reported"}, "reported", [2]*int{&maxSafe, &zero}, [2]*int{}, nil}, + {"unsafe_measurement", unsafe, 7, UsageDimensions{UsageStatus: "reported"}, "partial", [2]*int{nil, &seven}, [2]*int{}, nil}, + } + data, err := os.ReadFile("testdata/usage-accounting.json") + if err != nil { + t.Fatal(err) + } + var fixture struct { + Events []jsonEvent `json:"events"` + } + if err := json.Unmarshal(data, &fixture); err != nil { + t.Fatal(err) + } + if len(fixture.Events) != 5+len(cases) { + t.Fatalf("fixture events=%d", len(fixture.Events)) + } + for i, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + d := tc.d + d.ReceiptID = "generation:c3-" + tc.name + if tc.name == "estimate_without_receipt" { + d.ReceiptID = "" + } + d.EstimateScope = "text_only" + d.AccountingSource = "provider" + d.RoutingProvider = "other" + d.ModelFamily = "qwen" + at, err := time.Parse(time.RFC3339, fixture.Events[5+i].EventTime) + if err != nil { + t.Fatal(err) + } + event := UsageReceipt(ModeChat, tc.input, tc.output, "contract", at, d) + p := event.Props + if p["usage_status"] != tc.status { + t.Fatalf("status: %v, want %s", p, tc.status) + } + for j, side := range []string{"input_tokens", "output_tokens"} { + for _, group := range []struct { + prefix string + values [2]*int + }{{"", tc.measured}, {"estimated_", tc.estimated}} { + key := group.prefix + side + if group.values[j] == nil { + if _, ok := p[key]; ok { + t.Fatalf("invented %s: %v", key, p) + } + } else if p[key] != *group.values[j] { + t.Fatalf("%s: %v", key, p) + } + } + } + if tc.total != nil && p["total_tokens"] != *tc.total { + t.Fatal(p) + } + if err := relayUsageC3(p); err != nil { + t.Fatal(err) + } + // The shared fixture adds accepted constructor outputs without changing + // the envelope or the relay fixture test's first estimate/recovery pair. + if d.ReceiptID == "" { + event.ID = fixture.Events[5+i].EventID + } + event.InstallHash = fixture.Events[5+i].InstallIDHash + for _, key := range commonPropNames { + event.Props[key] = fixture.Events[5+i].Props[key] + } + got, _ := json.Marshal(event) + want, _ := json.Marshal(fixture.Events[5+i]) + var g, w any + json.Unmarshal(got, &g) + json.Unmarshal(want, &w) + if !reflect.DeepEqual(g, w) { + t.Fatalf("fixture: got %s want %s", got, want) + } + }) + } +} + +// relayUsageC3 mirrors validateTokenUsage and its scalar schemas in +// /home/abir/woi-82/site/src/lib/telemetry-usage.ts, including safe sums, +// presence-based status, disjoint components, and estimate receipt metadata. +func relayUsageC3(p map[string]any) error { + fail := func(s string) error { return fmt.Errorf("%s: %v", s, p) } + number := func(key string) (float64, bool) { + v, ok := p[key] + if !ok { + return 0, false + } + switch v := v.(type) { + case int: + return float64(v), true + case float64: + return v, true + default: + return math.NaN(), true + } + } + for _, key := range []string{"input_tokens", "output_tokens", "total_tokens", "estimated_input_tokens", "estimated_output_tokens", "estimated_total_tokens"} { + if v, ok := number(key); ok && (math.IsNaN(v) || math.IsInf(v, 0) || v < 0 || v > 1<<53-1 || math.Trunc(v) != v) { + return fail("unsafe count") + } + } + measured, estimated := 0, 0 + for _, side := range []string{"input_tokens", "output_tokens"} { + _, m := p[side] + _, e := p["estimated_"+side] + if m && e { + return fail("overlap") + } + if m { + measured++ + } + if e { + estimated++ + } + } + status := p["usage_status"] + _, hasTotal := p["total_tokens"] + totalOnly := status == "total_only" && measured == 0 && estimated == 0 && hasTotal + for _, prefix := range []string{"", "estimated_"} { + if prefix == "" && totalOnly { + continue + } + a, ak := number(prefix + "input_tokens") + b, bk := number(prefix + "output_tokens") + total, tk := number(prefix + "total_tokens") + if !ak && !bk { + if tk { + return fail("orphan total") + } + } else if !tk || a+b > 1<<53-1 || total != a+b { + return fail("inconsistent total") + } + } + valid := false + switch status { + case "total_only": + valid = totalOnly + case "reported": + valid = measured == 2 && estimated == 0 + case "estimated": + valid = measured == 0 && estimated > 0 + case "mixed": + valid = measured > 0 && estimated > 0 + case "partial": + valid = measured == 1 && estimated == 0 + case "missing": + valid = measured == 0 && estimated == 0 + } + if !valid { + return fail("status disagrees") + } + if receipt, exists := p["usage_receipt"]; exists { + s, ok := receipt.(string) + if !ok || !regexp.MustCompile(`^[a-f0-9]{64}$`).MatchString(s) { + return fail("bad receipt") + } + } + if estimated > 0 { + if p["usage_receipt"] == nil || p["estimate_method"] != "utf8_bytes_4_v1" { + return fail("missing estimate metadata") + } + switch p["estimate_scope"] { + case "text_only", "multimodal_input", "observed_partial", "unknown": + default: + return fail("bad scope") + } + } else if p["estimate_method"] != nil || p["estimate_scope"] != nil { + return fail("orphan metadata") + } + return nil +}