Feature/databricks gateway connector - #14
Conversation
`_strip_version` only matched a COMPACT trailing date (`-YYYYMMDD`), which is Anthropic's convention (`claude-sonnet-4-5-20250929`). OpenAI stamps a HYPHENATED one (`gpt-5-2025-08-07`, `gpt-4.1-2025-04-14`, `o3-2025-04-16`), and OpenRouter lists the BARE id (`openai/gpt-5`) — so a name we could not strip back to bare never matched. Because `resolve_model` prefers the response's own `model` over the requested one, `create(model="gpt-5")` resolves to `gpt-5-2025-08-07` and misses. Verified against the live OpenRouter table with this repo's own `lookup_openrouter`: gpt-4.1, gpt-4.1-mini, gpt-5, gpt-5-mini, o3 and o4-mini all fell through to token events, so anyone in price mode on a current OpenAI model was getting no cost at all. `gpt-4o` looked fine only by luck — OpenRouter happens to list `openai/gpt-4o-2024-08-06` verbatim. The pattern now accepts both shapes. All six resolve, the Anthropic compact cases still pass, and Workers AI ids (`@cf/meta/llama-3.3-70b-instruct-fp8-fast`) are left untouched.
…ovider hint Three changes to the native OpenAI adapter, all found while validating real provider responses. The drift contract did not hold one level down. `extras` swept only top-level usage keys, but `prompt_tokens_details` is itself a KNOWN top-level key, so nothing nested inside it was ever inspected. A live `gpt-5.6-sol` response carries `prompt_tokens_details.cache_write_tokens: 3022`, and those tokens were discarded with no error and no `on_error`. Every drift test passed, because none of them looked inside a details object. The sweep now recurses into the four `*_tokens_details` containers. Deliberately NOT mapped to `CanonicalUsage.cache_write`: for OpenAI these sit INSIDE `prompt_tokens` and bill at the plain input rate — cross-checked against Databricks' own metered spend, which charged exactly what billing all 3,025 as input produces. OpenRouter publishes a separate cache-write rate, so mapping the field would charge those 3,022 tokens twice, a 2.24x over-bill. `extras` keeps it visible without touching the money. Tokens in neither named bucket were silently dropped. For genuine OpenAI, `total_tokens` always equals prompt + completion — verified across every captured response, zero deltas. Behind an OpenAI-COMPATIBLE proxy fronting a thinking model it breaks: measured against Gemini through Google's compat layer, prompt 57 / completion 47 / total 1253, with 1,149 thinking tokens reported nowhere. A positive delta now folds into `output` as `extras["unaccounted_output_tokens"]`, minus any reasoning already broken out so an additive-reasoning provider is not billed for them twice. `extract_openai_native` also gains `provider_hint`, because two of Databricks' gateway surfaces use the same `openai.OpenAI` class but need different price tables, and the response body cannot tell them apart. Only the wrapper knows the `base_url`; the adapter stays the single place `provider` is decided.
Second entry in the `gateway/` namespace, alongside Cloudflare. `extract_databricks_log()` maps a `system.ai_gateway.usage` row to `CanonicalUsage`; `resolve_databricks_subscription()` reads Lago attribution from the caller's `Databricks-Ai-Gateway-Request-Tags` header. Verified against real rows read from a live workspace over the SQL Statement Execution API — 226 rows, all 36 columns (the public docs undercount at ~28), with 22 captured fixtures covering both destination types, cache read/write, reasoning, embeddings and all three failure shapes. Three mapping quirks a docs-only reading gets wrong, all caught by real rows: `destination_name` means DIFFERENT things per destination type — the model for a hosted row, the PROVIDER SERVICE (a Unity Catalog credential name) for BYOK. A single "model, falling back to name" rule bills a credential as the model on every BYOK row. `destination_model` is unstable for hosted models: the same `destination_name` reports both `gpt-oss-20b` and the display label `GPT OSS 20B`, depending on which of Databricks' two request aliases the caller used. Most hosted entities carry a second, INNER prefix — `system.ai.databricks-<model>`, on 38 of 48 distinct names. It is a serving-endpoint artefact, not part of the model id, but it cannot be stripped unconditionally because Databricks also publishes models genuinely named that way (`databricks-dbrx-instruct`). `destination_model` is the tie-breaker; disagreement keeps the raw name, since an ugly id is recoverable and a silently renamed model is not. `provider="databricks"` for hosted models is deliberately unmatchable in `_VENDOR_MAP`. Databricks bills them in DBUs against a rate card published only as HTML and present in no system table, while OpenRouter does list bare `openai/gpt-oss-20b` at 0.2-0.4x of Databricks' real rate — so being stamped "openai" would silently under-bill 2.5-5x. Same trap as Workers AI. This table's `input_tokens` INCLUDES cache_read and cache_write, the inverse of the providers' own response bodies. The adapter extracts faithfully and does not subtract; the module docstring records why, and why a computed fallback would need to correct per provider rather than uniformly. The barrel now exports gateway-scoped names, so neither gateway is the implicit default.
`gateway/databricks.py` is the one piece of gateway code that does I/O, and the
adapter beside it stays pure. Cloudflare's read is a single paginated GET and
rightly lives in its example notebook; Databricks needs a SQL warehouse, the
Statement Execution API, columnar-to-dict zipping, chunked result fetching, a
statement poll, and two tables reconciled against each other. Hand-rolled that is
~100 lines in which four money-losing mistakes are easy, and the first version of
the demo notebook made three of them:
Silent truncation — only chunk 0 arrives inline, so a window wide enough to span
`total_chunk_count > 1` bills a fraction of itself with no error.
Double billing — a BYOK call appears in BOTH `ai_gateway.usage` and
`external_model_spend`.
Unscoped idempotency keys — `transaction_id` is unique account-wide, so a key
built from the source row alone blocks that row from ever reaching a second
subscription. And the subscription billed is not always the one on the row, since
an untagged row falls back to the caller's default, so `event_id_for()` builds
the key from the resolved value.
Lost rows — a row with NULL ids, or an id a driver hands back as a non-string,
collapsed to an empty key, so every such row in the window shared one
transaction_id and only the first was ever billed. Falls back to a content hash,
which stays deterministic so re-runs remain idempotent.
`LagoSDK.backfill_databricks(source, "7 days")` bills a whole window and returns
`{"cost": n, "tokens": n, "skipped": n}`. It also accepts an already-read
iterable of rows, because a SQL warehouse costs roughly 1,500x the model-serving
usage it reports on — reading the window twice to print a summary first doubles
the expensive half and lets the summary disagree with what was billed.
A BYOK bucket with no spend row is billed by neither path, which the spend
table's ~19h lag makes routine for the newest hour, so it warns rather than
vanishing. The window is validated rather than escaped, since it reaches SQL by
interpolation. Deliberately absent: scheduler, cursor store, credential store.
In price mode a hosted call logged `lago pricing failed: no price for provider='databricks' model='meta-llama-4-maverick-040225'` and routed it to `on_error` on EVERY request. That description is wrong: nothing failed. Databricks bills hosted models in DBUs at a per-model rate that exists on an HTML page and in no system table — verified across every column of all 88 of them — so token counts are the complete answer for them, not a degraded fallback, and no refresh could ever supply the missing rate. New `TOKEN_BILLED_PROVIDERS` names the providers this applies to. `emit()` skips the lookup for them, emits token counts, and states the reason once per model at info level instead of warning once per call. Deliberately a narrow exception to "never silently under-bill". That invariant exists so a price miss cannot pass unnoticed, and it still holds for every miss a customer could act on — a cold table, an unmatched model name, a mistyped provider all still raise `PricingUnavailableError`. This covers only the case where the miss is structural and permanent. The reason to make it is that an alarm which always fires is one nobody reads: leaving it in place taught the reader to ignore `on_error`, which is precisely how a real miss gets missed. It keys on the PROVIDER, so it covers Databricks-hosted traffic only. BYOK through the same gateway is stamped openai/anthropic and prices normally — verified exact against Databricks' own metered spend on 38 of 38 buckets. The OpenAI wrapper supplies the `provider_hint` that makes this reachable, reading `base_url` once at wrap time. It must key on `/ai-gateway/mlflow/`, not `/ai-gateway/`, or the OpenAI BYOK path gets mis-stamped and priced against the wrong table.
README gains a `## Databricks AI Gateway` section covering all live ingress paths and the backfill, plus the gotchas customers would otherwise report as SDK bugs: `gpt-oss` inflates input by ~100 tokens from a server-injected preamble, `claude-opus-4-5` does not cache through this gateway at all, hosted models report three different name strings, and running the live path and the backfill over the same traffic emits token events twice. Corrects a claim that had nothing behind it: the "What gets billed" table said hosted backfill produced a dollar cost from `system.billing.usage` × `list_prices`, while the paragraph two lines below said the opposite, and neither table appears anywhere in the source tree. Hosted bills token counts on both paths. Those dollars do exist — `list_prices`, or `account_prices` for an account's contract rate — so "hosted USD is impossible" was also wrong; that applies only to the tokens→DBU rate. They are not billed from because they come from a different Databricks screen than the gateway view: no `request_tags`, so per-subscription splits would be ours rather than Databricks', and ~19h of lag. Every number this connector sends is one you can find on a Databricks *gateway* page, which is the property that makes it checkable. Each backfilled event carries the grouping key of the surface it came from — `endpoint_name` for hosted, `bucket` for BYOK — so grouping Lago the way the Databricks page groups puts the two side by side. Without it the comparison fails on naming alone, since our `model` is normalized and the page's is not. CONTRIBUTING gains an "Adding a gateway" recipe, the bar a read must clear to belong in the SDK rather than a notebook, and the two rules that keep a connector comparable against the gateway's own dashboard. `examples/databricks_gateway_demo.ipynb` demonstrates both halves and was re-run against a live workspace: 107 billable rows over 7 days, 60 dollar-cost events plus 88 token events, all transaction_ids unique.
…rker The test set `max_batch_size` equal to `max_buffer_size`, and `push` sets `_wake` whenever `len(buffer) >= max_batch_size`. So the overflowing push both dropped i=0 AND woke the background worker, which then drained all 10,000 events through `_take_batch`. When that landed before the next line read the buffer, `buf` came back empty and the assertion read `assert 0 == 10000`. Failed in CI on a loaded runner; reproduced deterministically by sleeping 50ms in that window, which is all the scheduler needs to do for free. Fixed the same way `test_repeated_overflow_keeps_window_sliding` already was: keep the batch size ABOVE the buffer cap so the buffer can never reach it, and the worker only runs once shutdown() releases the sender. Nothing in the test depends on batch size — every assertion is about buffer CONTENTS. Verified over 150 consecutive runs. Pre-existing; unrelated to the Databricks connector, but it is what turns this branch's CI red.
`mistral` was missing from `_INPUT_INCLUDES_CACHE_READ`, so in price mode the cached portion of a prompt was billed twice: once at the full input rate because `input` was never reduced, and again at the cache-read rate. Mistral's API is OpenAI-shaped and reports `prompt_tokens_details.cached_tokens` as a SUBSET of `prompt_tokens`. Its own documented example is unambiguous — prompt_tokens=1013, cached_tokens=1008, total_tokens=1043=prompt+completion, which only reconciles if the cached tokens sit inside the prompt count. Mistral bills them at 10% of the input rate. Measured 6.15x over-bill on that payload. 13 of 18 Mistral models on OpenRouter publish a cache-read rate, so the wrong path was reachable for most of them, including Mistral routed through a Cloudflare gateway (the gateway adapter leaves provider="mistral" unmapped). Token mode was unaffected — only the price computation was wrong. money_golden.json gains a `mistral` case built from Mistral's documented payload; removing the fix makes it produce 0.0006019 against the expected 0.0000979 and fail. This is the second provider missing from that set after `workers-ai`. The set is still hand-maintained; a completeness check over every provider slug the SDK can emit remains the real fix.
ancorcruz
left a comment
ancorcruz
left a comment
There was a problem hiding this comment.
Review notes
Ten items. The first five I'd fix before merge — four of them are silent-failure paths in the money-reading half, which is the one place a mistake doesn't announce itself. The rest is hardening, one scope question, and a cleanup.
Four were reproduced rather than reasoned about: the chunk-fetch truncation, the empty-columns collapse, the false "NOT billed" warning (the repo's own 22 fixtures produce 4 phantom rows), and the prefix-strip non-determinism (fixtures hosted_chat.json / hosted_chat_1.json carry the same destination_name with destination_model of llama-4-maverick vs Llama 4 Maverick).
Two notes that aren't inline because the lines aren't in this diff:
sdk.py's precomputed branch reportsunit = usage.input + usage.output. I flagged that on the Cloudflare PR already, but this connector aggravates it: perdatabricks_gateway.py's own billing-hazard note, this table'sinputincludes cache_read/cache_write, so onbyok_anthropic_cache_write_1.json(input=1825,cache_read=1812) theunitpublished alongside Databricks' dollar figure overstates by ~3x. The notebook's reconciliation block sumsr.usage.inputthe same way.- Considered and set aside: mixing a row ordinal into
_row_id's hash fallback to break ties between two id-less identical rows.ORDER BY event_timeisn't a unique sort, so an ordinal isn't stable across runs and would trade a rare collision for a broken idempotency guarantee. Leaving it as-is looks like the better trade.
|
|
||
| total_chunks = int(manifest.get("total_chunk_count") or 1) | ||
| statement_id = body.get("statement_id") | ||
| for index in range(1, total_chunks): |
There was a problem hiding this comment.
Blocker: the chunk loop reintroduces the silent truncation this module's docstring says it exists to prevent.
Neither the POST at line 209, the poll GET at line 268, nor this chunk GET checks HTTP status — every one goes straight to .json(). So when a chunk fetch fails (503, expired statement, revoked token mid-read), the error body has no data_array, or [] swallows it, and query() returns a partial row set with no exception. The total_chunks > 1 log at line 236 then cheerfully reports "spanned 3 chunks" for a read that got 2.
That is verbatim the first bullet of the module docstring: "A naive reader works on a small window and quietly bills a fraction of a large one, with no error."
Worth noting the codebase already has the convention — pricing.py calls raise_for_status() at all five of its fetch sites. This is the one HTTP path where a silent failure costs billing rows, and it's the one without the check. test_query_zips_columns_and_follows_every_chunk covers only the success path.
A raise_for_status() on all three calls is enough; letting it raise is right here, since backfill_databricks aborting loudly is far better than under-billing quietly.
There was a problem hiding this comment.
Confirmed, reproduced live, and fixed in 9e22961.
Forcing a real two-chunk statement (9,000 rows, chunk 0 = 6,750) and breaking chunk 1:
| before | after | |
|---|---|---|
| 403 / 404 / 503 on chunk 1 | 6,750 of 9,000 rows returned, no exception, 3/3 | all three raise, naming the chunk and the API's own cause |
| baseline (unbroken) | 9,000 | 9,000 |
So the loss was exactly the 25% you predicted, and query() reported success for it.
One correction, because it changes where the hole actually was. "Neither the POST at 209, the poll GET at 268, nor this chunk GET checks HTTP status" — the first two already failed loudly. Neither reaches .json() usefully on an error body: state is absent, so both raise RuntimeError("Databricks statement None: {…}") with the real cause embedded. Only the chunk loop was silent. I fixed the misleading statement None: prefix anyway, since it reads like a missing id rather than a failed call.
Beyond the status check, two things your comment implies but doesn't ask for:
total_row_countis now asserted against the assembled set. That catches a short read even when every individual chunk returns 200 — which the status check alone would not.if not columns: raise, which is the cheap half of yourmanifest.schemathread.
Four tests each side, all confirmed failing with the fix reverted. Worth noting why this survived review: the test doubles had no status_code / ok attributes at all, so no test could have expressed a failing chunk. They have them now — that gap is the more reusable finding here.
Note the branch was showing the pre-review tree when you looked; it's pushed now.
| if total_chunks > 1: | ||
| logger.info("lago: databricks result spanned %d chunks (%d rows)", total_chunks, len(arrays)) | ||
|
|
||
| return [dict(zip(columns, row, strict=False)) for row in arrays] |
There was a problem hiding this comment.
Blocker: a missing/short manifest.schema makes the whole window read as zero usage and report success.
columns is built from manifest.schema.columns with .get("columns", []) on line 223, so a SUCCEEDED body without a schema yields columns == []. With strict=False, dict(zip([], row)) is {} — so query() returns [{}, {}, ...], one empty dict per real row, no exception.
Downstream every one of those degrades cleanly and wrongly: extract_databricks_log({}) gives all-zero usage, the hosted loop skips them on nonzero_numeric(), the BYOK index builds all-zero buckets keyed ('', '', '', '{}'), and backfill_databricks returns {"cost": 0, "tokens": 0, "skipped": 0} for a window that had real traffic. Every defensive layer does its job and the result is a confident zero.
strict=True would turn a column/row length mismatch into an error, but it won't catch the empty case (zip([], row) is simply empty) — so an explicit if not columns: raise RuntimeError(...) is the part that matters.
There was a problem hiding this comment.
Traced this end to end and it is not reachable — downgrading from Blocker, with the guard added anyway.
Your downstream analysis is exactly right if columns is ever empty against a non-empty data_array. Measured on the live warehouse, that pairing does not occur:
- Every
SELECTreturns a full schema — 36 columns, including a zero-row result. - The only statement shape that yields
manifest.schema.columns == []isUSE CATALOG, and it also returns 0 rows. Sodict(zip([], row))never runs: there is no row to zip.
Which means the [{}, {}, …] shape — one empty dict per real row, reporting success — has no live path to it. A SUCCEEDED body carrying rows but no schema would be a Databricks API contract violation rather than an edge we can hit.
The guard is in regardless (9e22961, folded into the chunk-truncation fix): if not columns: raise. It costs one branch, and your reasoning about what happens downstream is sound enough that I'd rather fail loudly than rely on the API keeping that contract. But it's hardening, not a bug fix, and the changelog says so rather than claiming a live defect.
| model = model[len(_HOSTED_NAME_PREFIX) :] | ||
| if model.startswith(_HOSTED_ENDPOINT_PREFIX): | ||
| shed = model[len(_HOSTED_ENDPOINT_PREFIX) :] | ||
| if shed == _safe_str(row.get("destination_model")): |
There was a problem hiding this comment.
Blocker: exact-equality against a column this module documents as unstable, so one hosted model can bill under two different ids.
The comment above (lines 88-101) is right that the prefix can't be stripped unconditionally and that destination_model is the signal for telling an artefact from a real name. The problem is the exactness of the comparison, given what the module docstring says two paragraphs earlier: destination_model for hosted models flips between a slug and a human display label — measured, gpt-oss-20b / GPT OSS 20B.
The fixtures show the flip directly: hosted_chat.json and hosted_chat_1.json carry the same destination_name (system.ai.llama-4-maverick) with destination_model of llama-4-maverick and 'Llama 4 Maverick'.
llama-4-maverick is unaffected because it carries no inner prefix. But for one of the 38-of-48 prefixed names, the flip decides the outcome: a slug row emits qwen35-122b-a10b while a display-label row emits databricks-qwen35-122b-a10b. Same model, two Lago rows — the exact split the comment says the rule prevents, just triggered by row-level label variance instead of by a bad strip.
Comparing normalized forms on both sides (_alnum, or lowercase + spaces→hyphens) makes the decision stable whichever label the row carried, and still refuses to strip when the two columns genuinely disagree — Qwen35 122B A10B normalizes onto qwen35-122b-a10b, while databricks-dbrx-instruct vs dbrx-instruct still does not.
There was a problem hiding this comment.
Checked this against the live workspace and the blocker does not fire. The mechanism you describe is real; what makes it harmless is which names it can land on.
Measured over 47 distinct hosted destination_names:
| count | display-label flips | |
|---|---|---|
carry the inner databricks- prefix |
37 | 0 |
do not (system.ai.gpt-oss-20b, system.ai.llama-4-maverick, …) |
10 | all of them |
The exact-equality comparison sits inside the strip block, which is only entered for a name carrying the inner prefix. The flip only ever happens on the names without it — where the block is never entered and the comparison is never evaluated. Zero of the 37 prefixed names ever disagree, so there is no name for which one model can bill under two ids.
The fixture you cited is the harmless case, which is why it shows the flip so clearly.
The flip also isn't per-request variance, which is what would make it dangerous: it tracks the row generation. service_type IS NULL → 4 rows, all display labels, 0 prefixed. service_type = 'MODEL_SERVICE' → 90 rows, 44 prefixed, 0 display labels. So it's a legacy-writer artefact, stable per row, not something that alternates for the same model.
Not fixed, deliberately: it would be hardening against a case with no live path. Happy to make the comparison tolerant (normalise both sides before comparing) as its own pair if you'd rather close the shape off — it's cheap, I just didn't want to spend the PR's narrowness on something measured at zero.
| # Index token counts by the spend table's own grouping key, so a BYOK event | ||
| # can carry real counts alongside Databricks' dollar figure. | ||
| tokens: dict[tuple[Any, ...], CanonicalUsage] = {} | ||
| for row, u in extracted: |
There was a problem hiding this comment.
Blocker (severity is trust, not money): failed rows enter the join index, so the "NOT billed" warning is permanently wrong.
This loop filters only on u.provider == "databricks". The hosted loop at line 384 additionally guards if not u.nonzero_numeric(): continue; this one doesn't, so failed calls — which the adapter docstring says arrive with NULL tokens — become entries in tokens, and then show up in set(tokens) - billed_keys at line 369.
Ran the repo's own 22 fixtures through extract_databricks_log: 4 rows are non-hosted with entirely zero usage — gemini_broken.json, gemini_broken_1.json, unmanaged_path.json, unmanaged_path_1.json. Those are calls that never reached a provider, so external_model_spend will never have a row for them, and the warning's advice ("re-run this window later to bill them") is false forever, not just this window. Any workspace with a broken Gemini connection gets a permanent warning on every run, plus an inflated bucket count.
Which matters more than a cosmetic log: this warning is the only signal that real BYOK rows went unbilled (see the backfill_databricks comment). An alarm that always fires is one nobody reads — which is the argument TOKEN_BILLED_PROVIDERS makes in this same PR.
Same if not u.nonzero_numeric(): continue guard here fixes it.
There was a problem hiding this comment.
Confirmed, and the operator-trust damage is worse than "permanently wrong". Fixed in 15d1433.
Live, before:
- 29 unbilled buckets reported — 28 phantoms, 1 genuine.
- 83 of 151 BYOK rows carry zero usage, and every single one is a 4xx/5xx (500 / 403 / 404 / 502 / 400). Not one successful call among them, exactly as the adapter docstring predicts.
- The one example bucket the warning surfaces to the operator was itself a phantom, and it printed
model=— empty, because a failed row has nodestination_model. So the single most actionable line in the message was pointing at a call that never consumed anything.
After: the real reader warns nothing on that window. And a positive control on the same live rows with one hour's spend withheld — the shape of genuine lag — warns for exactly 54 buckets in that hour, not inflated by the 33 failed calls sitting in the same hour, and names a real model (claude-haiku-4-5) instead of the empty string.
The fix is the one-line guard you pointed at, placed before the key is built rather than at the top of the body. That ordering matters for a second reason: _merge_usage keeps the first row's non-numeric fields, so a zero row arriving first in a bucket would become its representative and donate its empty destination_model to the whole bucket. I checked whether that was live — 0 of 94 buckets hold both real and zero rows, because a failed call's empty model gives it its own key — so it isn't firing today, and the changelog says "closed before the join key changes" rather than claiming a live fix.
2 tests, both confirmed failing with the guard reverted; the revert's captured log reproduced the live model= symptom, which is what convinced me the ordering was worth stating.
| duplicates rather than double-bill. Does not flush — call ``flush()`` when | ||
| you want to block on delivery. | ||
| """ | ||
| counts = {"cost": 0, "tokens": 0, "skipped": 0} |
There was a problem hiding this comment.
Blocker: the one-call entrypoint returns success while knowingly under-billing.
counts reports cost / tokens / skipped, but skipped only counts rows this loop received and couldn't attribute. The BYOK token buckets that read_usage drops — the documented external_model_spend aggregation lag, so routinely the window's most recent hour — are never yielded, so they land in no counter. They're reported only to a module logger, never through config.on_error.
Net effect: backfill_databricks returns {"cost": 60, "tokens": 47, "skipped": 0} for a window where real usage went unbilled, and a caller reconciling on the return value or an on_error hook sees nothing. That's in tension with the invariant emit() is otherwise careful about — the CHANGELOG's "don't silently under-bill" and the PricingUnavailableError report on a price miss.
Two options, either works: have read_usage surface the count (return it, or yield the unbilled buckets as a sentinel) so counts can carry a deferred key; or route the warning through on_error so it reaches the same hook every other billing gap uses. The docstring's promise that the return value is "counts of what it emitted" stays true either way — the gap is that what it didn't emit is invisible.
There was a problem hiding this comment.
Confirmed, and your framing of the net effect is precisely what the live run shows. Fixed in a124962.
Against real Databricks rows and a real Lago, replaying a window with one hour's spend withheld (the shape of genuine aggregation lag):
| before | after | |
|---|---|---|
| returned counts | {'cost': 12, 'tokens': 54, 'skipped': 0} |
{'cost': 12, 'tokens': 54, 'skipped': 0, 'deferred': 54} |
| buckets provably unbilled | 54 | 54 |
on_error calls |
0 | 1, naming hour=2026-08-07T13 provider=anthropic model=claude-haiku-4-5 |
The part that makes this a trust bug rather than a reporting nit: the only caller-visible difference from a clean run was cost 66 → 12, which is indistinguishable from a quieter window. There was no signal to reconcile against.
skipped was routed through on_error in the same change, because it's the same defect in the same function — counted and returned, but never reported, and it never reaches emit() where every other dropped event is reported from.
Two design notes:
deferred_bucketsis rewritten per read, not accumulated. A healthy window read after a lagging one must not restate the older window's gap — that's the phantom shape from your join-index thread arriving through a different door.- The already-read-list path returns
deferred: 0by construction, and the docstring says why: only the reader knows about a bucket it never yielded.
Accepted cost: the gap now logs twice — the reader's own warning plus _report_error's line. Kept deliberately, because a direct read_usage caller never reaches the second one. The backfill's wording moved to the run's side ("this run left N … unbilled") so the two read as one gap seen from two layers rather than as two separate gaps.
Control: the untouched window returns deferred: 0 and fires nothing, so it isn't a per-run false alarm. 3 tests, all confirmed failing on revert.
| billed_keys: set[tuple[Any, ...]] = set() | ||
|
|
||
| for row in spend: | ||
| usd = _safe_float(row.get("usage_quantity")) |
There was a problem hiding this comment.
A negative usage_quantity passes this filter, then floors to $0 and burns the row's idempotency key.
if not usd catches 0.0 but not -0.0042. A credit or correction row therefore yields a DatabricksUsageRow(usd_cost=-0.0042), and compute_precomputed_cost does _parse_price(usd_cost) or Decimal(0) — _parse_price returns None for negatives (pricing.py:196), so base is 0.
The floor-to-zero itself is deliberate and documented at pricing.py:346. The problem is what happens here: a real llm_cost event is pushed with precise_total_amount_cents="0" under a real transaction_id. The credit is lost, nothing is logged, and the id is now consumed — so if the row is later corrected, the re-run is rejected as a duplicate and can't fix it.
if usd <= 0: continue (with a warning when it's negative) keeps the id unused, which is the recoverable state. Whether Databricks ever emits negative usage_quantity on this table I can't confirm — but the cost of guarding is one comparison, and the cost of not guarding is unrecoverable.
There was a problem hiding this comment.
Confirmed — and it is worse than "billed at zero", which I only found by reading the event back out of Lago rather than checking what we sent. Fixed in 74b93b9.
Repro used the window's single most expensive real spend row (record_id=6a7cfc05…, gpt-5.6, $0.015245) restated negative:
| before | after | |
|---|---|---|
| events for the negated row | 1, value: "0" |
0 |
| what Lago holds for that id | {"value": "0"} |
HTTP 404 — nothing |
| counts | cost 66, deferred 0, 0 on_error |
cost 65, deferred 1, 1 on_error naming the hour |
| the correction, same window re-run | still "0" — 422 value_already_exist |
"0.015245" |
That last row is the whole finding: the $0 event isn't just wrong, it makes the right number unbillable forever, because the record_id-derived transaction_id is spent. Control: before the fix, the same corrected row billed fine under a fresh event-id prefix (0.015245) — which proves the money was billable all along and the burnt id is what lost it.
Guard is usd <= 0 rather than a negative-specific test, so it also covers the 0.0 case the old if not usd already caught, in one predicate.
Two deliberate choices worth flagging:
- The negative branch logs, carrying the figure, model and hour. A credit or restatement is a real Databricks event this connector has no way to represent, and skipping it leaves the customer billed more than Databricks metered — so it must not vanish quietly.
- The bucket now surfaces in the
deferredreport. That's the honest reading (its tokens went unbilled) even though "re-run later" won't resolve this cause. Documented in-comment rather than special-cased, because a caller reconciling ondeferredshould see it.
| WHERE usage_start_time >= {window} | ||
| """) | ||
|
|
||
| usage = self.query(f""" |
There was a problem hiding this comment.
SELECT * plus an unused ORDER BY, on the one resource this module says costs ~1500x the usage it prices.
Two separate costs:
Warehouse time. SELECT * pulls all 36 columns including endpoint_metadata, routing_information, invocation_metadata and the service_* set, while the adapter reads about twelve. The ORDER BY event_time forces a sort over the whole window, and the rows go straight into extracted and a dict index where order is never used — the spend loop iterates spend, the hosted loop iterates extracted, and neither depends on ordering.
A correctness edge. _row_id's fallback hashes the entire row (line 421), so for a row with NULL invocation_id and request_id, an unused column changes the transaction_id. invocation_metadata or a latency field differing between two reads of the same window would produce a different id for the same row — breaking the idempotency the docstring promises. Naming the needed columns removes that coupling as a side effect.
Given the module explicitly tells operators "read one wide window per run; never poll in a tight loop", narrowing the projection seems worth it here.
There was a problem hiding this comment.
Fixed in a6dc422 — 14 named columns via _USAGE_COLUMNS, ORDER BY dropped. But the two halves of your comment landed very differently once measured, and the one you flagged as the smaller point is the real one.
The projection cost, measured on the raw API payload over 247 real rows:
| bytes/row | rows before the 25 MiB INLINE cap | |
|---|---|---|
SELECT * (36 cols) |
1,411 | ~18.6k |
| the 14 actually read | 435 | ~60.3k |
3.24x, and the framing that matters is the second column, not warehouse time. Exceeding disposition=INLINE doesn't degrade — it hard-fails the statement. So for a module whose own guidance is "read one wide window per run", SELECT * was a 3.2x smaller maximum window. That's the argument I'd have led with, and it isn't in your comment.
Your correctness edge is unreachable. The SELECT * → row-hash id path: 0 of 241 rows have both invocation_id and request_id null, so the hash fallback never runs and the column set can't influence an id. Real, just not live.
Re-verified the billing output is untouched by re-reading the same 2026-08-06 window both ways through the real reader: 120 billable events, identical as a multiset, 120 unique transaction ids, deferred buckets identical (0 each) — only the row order differs.
Residual I deliberately did not fix (tell me if you want it). Dropping ORDER BY makes which row represents a merged spend bucket warehouse-dependent, because _merge_usage keeps the first row's _BUCKET_INVARIANT_EXTRAS. What bills is unaffected — numerics are summed. Measured 0 disagreements over 66 live BYOK buckets, but only 2 of those held more than one row, so that's unobserved rather than impossible. Two UC connections to the same vendor + model + hour + tags would collide. The real fix is a bucket refusing to state a field its rows disagree on, which is a semantics change well beyond this.
Your third bullet — "consider documenting/handling the INLINE ceiling" — is documented, not handled: the cap and both row ceilings are stated on _USAGE_COLUMNS and in the changelog. Handling it (EXTERNAL_LINKS, or splitting a window by time when the manifest is large) is a bigger change and isn't scheduled.
2 tests, both confirmed failing when the narrowing is reverted — and the coverage one also fails if a needed column is dropped from the list, so it guards over-narrowing in the other direction too.
| # A no-op for real OpenAI either way: total always equals prompt + completion. | ||
| declared_total = _safe_int(usage.get("total_tokens")) | ||
| if declared_total: | ||
| unaccounted = declared_total - (input_tokens + output_tokens + reasoning) |
There was a problem hiding this comment.
Scope question on the total_tokens guard: it subtracts reasoning, but nothing else that can be additive.
The reasoning subtraction is well argued, and the Gemini-behind-a-compat-layer case it's built for is real. What I'm less sure of is that reasoning is the only field that can inflate total_tokens relative to input + output.
This same diff documents (lines 52-53) that Anthropic's cache_creation_input_tokens sits OUTSIDE input_tokens, and both gateways in this SDK front Anthropic models behind OpenAI-shaped surfaces. On a payload like {prompt_tokens: 13, completion_tokens: 4, total_tokens: 1829, prompt_tokens_details: {cache_write_tokens: 1812}}, unaccounted is 1812 and those cache-write tokens get folded into output — billed at the output rate, and simultaneously surfaced in extras["prompt_tokens_details.cache_write_tokens"].
I can't point to a live payload with that shape, so this may be unreachable in practice — but the guard is unconditional on provider, and the argument for it ("no completion_tokens_details to recover them from") is specifically about a payload with no breakdown. Gating on that — only fold in the remainder when the details sub-objects are absent or empty — would keep the measured case working while making the assumption explicit. If instead you've confirmed no fronted provider reports cache-creation additively in total_tokens, a line saying so would be enough.
There was a problem hiding this comment.
Good question, and the honest answer is not changed — but I measured the delta it would fold on every surface we have data for, and it's zero on all three.
Your premise is right: reasoning is not a priori the only field that can inflate total_tokens, and Anthropic's cache_creation_input_tokens sitting outside input_tokens is exactly the shape that would do it. What decides it is whether any surface actually reports that way.
| surface | measured | unaccounted |
|---|---|---|
| native OpenAI | total == prompt + completion, zero deltas across every capture |
0 |
| Databricks gateway | total_tokens == input + output for 112 of 112 rows, cache_read and cache_write inside input, reasoning inside output — for every vendor, Anthropic included (i=1822, cr=1812, t=1826) |
0 |
| Cloudflare gateway | preserves each vendor's native shape: a real Anthropic entry reads input=10, output=4, total_tokens=14 with input_cached_tokens=3429 sitting outside that total |
0 |
The Cloudflare row is the interesting one for your argument: cache tokens really do sit outside input there, exactly as you describe — but they sit outside total_tokens too, so they don't inflate it relative to input + output and the guard never fires.
So the hazard needs a proxy that folds cache_write into total_tokens while reporting it in neither named bucket. None of the three surfaces this adapter serves does that today.
Not fixed for two reasons, both about blast radius rather than disagreement: it's the money path under the byte-identical cross-port invariant, and subtracting a second field changes the Gemini-compat case the guard exists for (there, reasoning is 0 and the whole delta is real output — a cache_write subtraction would have to be provably inapplicable there before it's safe). Happy to do it as its own pair with the same treatment the reasoning subtraction got, if you'd rather close the asymmetry than rely on the measurement.
Related and in scope: the Databricks row above was a live over-bill for a different reason — the SDK keyed cache de-overlap on the vendor name while that table normalises every vendor to OpenAI shape, so Anthropic billed at 1.570x. Fixed on the JS thread's twin (5d97549 here), now 0.998.
| _DATABRICKS_HOSTED_PATH = "/ai-gateway/mlflow/" | ||
|
|
||
|
|
||
| def _provider_hint_for(client: Any) -> str: |
There was a problem hiding this comment.
The base_url→provider rule now exists in three places, and this one only covers the OpenAI wrapper.
The reasoning is right and the hazard is real — my concern is placement. base_url is currently inspected to decide a provider in three spots that must stay in agreement: here, sdk.py::_auto_prime_pricing_for (its own read, for Cloudflare), and _infer_provider's model-string variant. Only this one produces a provider_hint, and only wrappers/openai.py passes it.
Concretely: Databricks hosts Anthropic models (system.ai.databricks-claude-sonnet-4-5). Whether a customer can reach those through anthropic.Anthropic(base_url=...) depends on whether Databricks exposes a native Anthropic-shaped surface for hosted entities — if it does, that client gets provider="anthropic", strips to claude-sonnet-4-5, hits OpenRouter, and prices a DBU-billed model at Anthropic's rate: the 2.5-5x mispricing this docstring says it makes impossible. Worth confirming one way or the other, since the answer decides whether this is a gap or just an asymmetry.
Either way, there's a bonus in consolidating: a single base_url→provider resolver consumed by wrap() for every client kind would also close the Cloudflare /compat streaming hole I flagged on the base PR — where a workers-ai/@cf/... streaming call is stamped openai because the response carries no model and the string check wants a bare @cf/. The wrapper knows the base_url there too; it just has no way to say so today.
There was a problem hiding this comment.
Your concrete question has a definite answer, and it resolves in the code's favour — I probed it live rather than reasoning about it.
No, a customer cannot reach a hosted DBU-billed Claude model through /ai-gateway/anthropic. Against the real workspace:
- Hosted names (
system.ai.databricks-claude-*) → 403,"not in the allowed models list". - No
Databricks-Model-Provider-Serviceheader at all → 501NOT_IMPLEMENTED.
That surface requires a header naming a Unity Catalog connection, i.e. the customer's own credential. It is BYOK-only by construction, so provider="anthropic" is correct there and there is no path by which the hint mislabels a DBU-billed call. An asymmetry with the mlflow surface, not a gap.
On the placement concern — the Python side has nothing to forget. I went looking for the mirror of the JS defect and it doesn't exist here: the stream wrapper is nested inside wrap_openai_client and reaches the adapter only through _emit_from, which closes over provider_hint. There is no argument a future call site could omit. 7acad14 records that as a comment so the next person doesn't "fix" it by adding a parameter.
The JS twin was the real one — a module-level async function* taking the hint as a defaulted 5th positional parameter — and it is now required, so omission is TS2554 rather than a silent revert to provider: "openai". That's the thread on #29. Reproduced live before fixing: dropping the argument typechecked clean and a real hosted stream of system.ai.gpt-oss-20b emitted provider: "openai".
Consolidating the three base_url readers: not done. You're right that they must stay in agreement and that only one produces a hint. I left it because the three answer different questions (which price table to warm, which provider to bill, which surface a model string implies) and a single helper would need to serve all three without becoming a switch on its caller. Worth doing, but it's a refactor across both ports and I'd rather it not ride along in this PR. Say the word and I'll scope it separately.
One correction to something adjacent, since it affects how much the hint is worth: provider_hint_for's docstring names meta-llama/llama-4-maverick as reachable, and that's overstated. It's true of the OpenRouter catalogue but not of the lookup — VENDOR_MAP sends it to openai/…, so it misses under the openai vendor prefix in every spelling. Flagged, not corrected here.
| from typing import Any | ||
|
|
||
| from ..canonical import CanonicalUsage | ||
| from .adapters.databricks_gateway import ( |
There was a problem hiding this comment.
Cleanup (low priority): third copy of the _safe_* helpers, and a private symbol crossing a module boundary.
_safe_dict / _safe_int / _safe_str now exist in adapters/openai_native.py, gateway/adapters/cloudflare_gateway.py and gateway/adapters/databricks_gateway.py. The bodies have already diverged — only the Databricks copy tolerates a JSON-string dict, which is a genuinely useful improvement the other two don't get.
This import also makes _safe_str part of a cross-module contract despite the leading underscore. adapters/_common.py is the established home for shared adapter helpers; a gateway/_common.py mirroring it would fit the existing layout.
Also minor, same file: json is imported at line 35 and again inside _canonical_tags at line 471.
There was a problem hiding this comment.
Agreed on the direction, and the count is higher than three — which is also why I didn't do it here.
It's the 7th–8th copy, not the third. _safe_dict / _safe_int / _safe_str are defined locally in 7 of the 8 adapter files. So this is an established repo-wide pattern rather than something this PR introduced, and consolidating it means touching 8 files in both ports under the line-by-line mirroring invariant — a much bigger diff than the connector it would be riding on.
Your two observations are both correct and are the strongest case for doing it:
- The bodies have already diverged. Only the Databricks copy tolerates a JSON-string dict, which is a genuinely useful improvement the other six don't get. Divergence-in-place is exactly what duplication costs, and it's already happened.
_safe_strcrossing a module boundary under a leading underscore is a real contract smell.adapters/_common.pyis the right home.
Not fixed, and I'd rather not fold it in: the user has asked repeatedly for narrow PRs, and a helper consolidation across 16 files is the kind of change that makes a connector review unreviewable. Happy to take it as its own mirrored pair — and if we do, the JSON-string tolerance should become the shared behaviour rather than being dropped back to the lowest common denominator.
Brings the merged Cloudflare AI Gateway connector (#13) and the review-fix stack (#19) under the Databricks work. Four conflicts: - `pricing.py` version-strip: main split the 3-digit `-002` arm out into an OpenRouter-only regex (correctly — the shared helper also builds the AWS/Bedrock price keys, where a shortened key overwrites another model's rate). This branch had widened the SHARED regex with the hyphenated `-YYYY-MM-DD` date that every current OpenAI model reports. Kept both, and added the hyphenated arm to the OpenRouter regex as well: the OpenRouter lookup now goes through the scoped helper, so leaving it out silently re-broke price mode for gpt-4.1 / gpt-5 / o3 / o4-mini. Confirmed by reverting the widening — the six dated-model tests fail. - `test_buffer_overflow.py`: both sides independently fixed the same worker/assertion race. Kept one comment, mirroring the JS twin's wording. - `test_pricing.py` and `CHANGELOG.md` were additive on both sides; kept both. Verified on the merged tree: ruff, ruff format, mypy strict, 641 unit tests.
…g pricing Two JS-port behaviours the Python queue never got. Both measured against the real `EventQueue` on a driver rather than reasoned about. **Unbounded respin after isolating a batch.** A permanent batch failure (422) whose isolated sends then fail transiently (429) re-queued the survivors and continued straight into re-taking them with the backoff reset to 0 — no delay anywhere in the cycle. Measured: **280,388 HTTP requests in 1.2s**, aimed at the server that had just asked us to slow down. `_send_individually` now returns how many it re-queued; a non-zero count arms the normal 1->2->4->...->60s backoff, and the immediate-continue is kept only for the case it was written for — isolation fully resolved the batch, so the buffer really did shrink. After: 8 calls, backoff 2.0s. The exit drain gets `requeue_transient=False` for the same reason: there is no later retry to re-queue TO, so an event failing there is reported lost rather than handed back to a buffer this same loop immediately re-takes. **Refresh ran ahead of the drain.** `maybe_refresh()` does HTTP — up to 10s per source — and ran before the buffer was drained, on every tick. Measured: a 600ms refresh pushed first delivery to 629ms; now 37ms. Drain, then refresh, then drain again so anything pushed during the fetch does not wait out another flush interval. Both the refresh and the trailing drain are skipped once `_stopping` is set, so a shutdown landing mid-tick no longer spends the caller's budget on a table nothing will read. Four tests mirroring the JS names; each confirmed to fail with its own fix reverted. ruff, mypy strict, 645 unit tests green.
…e captures
Three machine-swept directories held captures that re-asserted the same thing.
The sweep tests assert DISPATCH and non-zero usage, so two captures sharing a
usage shape, an adapter family and a pricing provider exercise one code path
twice. Reduced to one per distinct key:
bedrock/converse 39 -> 12 shape x pricing provider
bedrock/invoke 39 -> 14 shape x invoke family x pricing provider
mistral_native/all_models 73 -> 12 shape x model family x vision flag
Every fixture a dedicated test names is kept as its group's representative, so
nothing referenced by name disappeared — verified by extracting every literal
`*.json` in both repos' tests first. The sweeps' own coverage assertions are
unchanged and still pass: all 7 InvokeModel families, all 8 Mistral families,
4 vision captures.
A COVERAGE.md per directory lists every model that was captured and which
committed fixture now stands for it, so the breadth of the live verification
stays reviewable as a table instead of as 113 near-identical JSON blobs.
Hand-curated directories are untouched. Several Databricks pairs that look
redundant are not: hosted_chat.json / hosted_chat_1.json differ only in
destination_model ("llama-4-maverick" vs "Llama 4 Maverick"), which is the
captured evidence that the column is unstable.
Separately — personal data removed from all 22 Databricks fixtures.
`system.ai_gateway.usage` logs the caller's account email and source IP on
every row, and no adapter reads either, so both repos were publishing a
personal Gmail address, a residential IP, a real workspace subdomain and one
live Lago subscription id to a public index. Replaced with RFC 2606/5737
reserved values; no credentials were ever committed. The gateway fixtures have
no capture script to hang a scrub step off, so a hygiene test is the durable
fix — confirmed to fail on all three when the originals are put back.
The Statement Execution API returns only chunk 0 inline; query() fetched the
rest but never checked their HTTP status. A failed chunk fetch returns a JSON
error body with no `data_array`, so `or []` appended zero rows, the loop moved
to the next index, and query() returned a PARTIAL result reporting success —
verbatim the "silent truncation" this module's docstring leads with.
Measured against a live warehouse on a genuinely chunked read (9,000 rows over
2 chunks): a 403, 404 and 503 on chunk 1 each returned 6,750 rows with no
exception. 25% of the window billed as if it were all of it, three times out of
three.
Chunk fetches now raise through a new _raise_for_api_error, which reads the
API's body rather than calling raise_for_status(): Databricks puts the cause
there and requests shows only the status line. The assembled row count is also
asserted against manifest.total_row_count, which catches truncation no
per-request status check can see — a chunk returning HTTP 200 with fewer rows
than promised, or a manifest/chunk disagreement.
The other two calls were already loud, contrary to how this was reported: a
non-OK submission or poll has no status.state, so _await_statement's own guard
already raised. What changes for them is only legibility — the real cause
instead of "Databricks statement None: {...}". The 403 "does not have required
scopes: sql" this class warns operators about is the error most likely to hit a
first-time setup, so it is the one that had to read clearly.
Rows arriving with no manifest.schema.columns now raise rather than zipping to
{} each, which every layer downstream degrades cleanly and wrongly on, ending
in a confident {"cost": 0, "tokens": 0} for a window that had real traffic. Not
observed on this API — every SELECT returns a full schema, zero-row reads
included — so this guards the decode, not a known bug. strict=False is kept on
the zip deliberately: a length mismatch is an API-contract violation the
row-count check already catches, and strict=True would newly reject reads that
work today.
_FakeResponse gained status_code, which a real requests response always
carries. Its absence is why this gap survived: no test could express a non-OK
response.
system.ai_gateway.usage re-reports every vendor in the OpenAI shape. Measured across 246 rows and 6 vendors: total_tokens == input + output for EVERY group, with cache_read AND cache_write inside input and reasoning inside output. The two billing paths decided that from the provider name, which is right for a native call and wrong for a table row. Anthropic's own API reports cache additively — measured live, cache_read=3962 against input=9 — so an Anthropic row read from this table had its cached tokens counted twice. On a real backfill over 2026-08-06: 48,798 tokens billed against 31,091 consumed, 1.570x. The same window now reports 31,018, the same ~0.2% lag-only shortfall openai already had. The api is the honest key. A gateway row reuses the live vendor names, so nothing in the name separates a table row from a direct call. workers-ai stays a provider entry by contrast: it names a vendor reachable through exactly one surface, so there the name is sufficient. This also adds the first cache_write de-overlap the SDK has had. It is surface-only by design — Anthropic is the one vendor whose native API bills cache writes at all, and it reports them additively, so no native response needs the correction. Databricks-hosted models were a latent 1.991x rather than a live one: they bill as token counts and 0 of 96 hosted rows carry cache today. Fixed before it fires. Cloudflare AI Gateway is deliberately excluded and now pinned by a test. Its logs preserve each vendor's native shape — a real Anthropic entry reads input=10, output=4, total=14 with input_cached_tokens=3429 outside that total — so adding it would under-bill the cached portion. Five cases added to money_golden.json, byte-identical in both repos, including controls that an openai row is subtracted exactly once (correcting it twice was a measured 13% under-bill) and that the native Anthropic path is untouched. All five new assertions confirmed failing with the fix reverted.
emit() read time.time() at each of its three push sites, so backfill_databricks
billed a whole window into whatever period the script happened to run in.
Measured on a live backfill over 2026-08-06: 128 events read off rows spanning
2026-08-06 to 2026-08-11 all carried one timestamp, the run's — up to 13.9 days
of drift. Once the event is in Lago nothing can tell which period the usage
actually belonged to. The same window now emits 25 distinct timestamps, one per
distinct source time, 128 of 128 matching their own row.
emit() takes a new timestamp= (a datetime, a naive one read as UTC, or epoch
seconds) and backfill_databricks passes each row's own time through a new
DatabricksUsageRow.occurred_at: event_time for a usage row, and for a spend row
the bucket hour START. The start is the only instant certain to sit inside the
hour that row aggregates; the hour's end would push a bucket closing exactly on
a period boundary into the following period.
Resolved once per call, ahead of every branch, rather than at each push site. A
price-lookup miss falls through to the token path, so one usage row can reach
two of those sites, and two separate clock reads there let a call straddling a
period boundary land half in each period.
A value that cannot be read is reported through on_error (where="timestamp")
and the call still bills, at now. Stamping the wrong period is a reconciliation
problem the operator can see and fix; dropping the event is revenue that never
appears at all.
An ISO-8601 string is deliberately not accepted, and neither is a numeric one.
Python 3.10 is still supported and its fromisoformat rejects the trailing "Z"
that gateway APIs emit, while the JS port's new Date() accepts it — a string
would parse in one repo and fail in the other. int("1786112523") would likewise
coerce where the JS port's typeof check refuses. Connectors parse their own
source column instead, where the shapes it really returns are known and tested:
_epoch takes both the API's "…Z" strings and databricks-sql-connector's
datetime objects, and reads an offset-less stamp as UTC so the JS port's Date
cannot treat it as local time. All five real column shapes verified to produce
byte-identical epochs in both repos.
The live wrap() path passes no timestamp and stamps now exactly as before,
pinned by its own test.
11 tests added; the 10 behavioural ones confirmed failing with the fix
reverted. 555 pass, ruff + format + mypy clean.
_interval_sql returned a SQL string, so current_timestamp() - INTERVAL 1 DAY
was re-evaluated per statement — 5.1s of drift measured between the spend read
and the usage read on a warm warehouse. Spend runs first, so the usage window
was the narrower one, and a Databricks-hosted row landing in that gap was read
by neither statement: hosted bills from ai_gateway.usage alone, so the call was
simply never billed. Both statements now carry one pair of literals, resolved
once here rather than twice in SQL.
The window is floored to the hour. external_model_spend is an hourly aggregate
whose usage_start_time is always the hour START — 65 of 65 live rows, none
unaligned — so a mid-hour bound failed the predicate for the hour containing it
while ai_gateway.usage happily returned that same hour's rows. Measured live: a
since of 13:30 read 11 of 65 spend rows, dropping $0.1256 of $0.1723 (73%) of
the window's metered dollars, while still reading 35 BYOK usage rows (31,815
tokens) from inside the hour it had dropped — tokens that then tripped the "no
spend row" warning as if the table were lagging. The same read now returns all
65, both statements carrying the identical floored pair.
Flooring the lower bound can read rows slightly older than asked for. That is
safe in the only direction that matters: every transaction_id is derived from
the source row, so a row already billed is rejected as a duplicate, and one not
billed yet should be. Under-reading is what loses money.
The still-aggregating hour is excluded. A spend row cannot be complete before
its hour closes: the 08:00-09:00 row appeared ~7 min after 09:00, so the wait is
not a fixed lag but however long is left in the hour — ~9 min for a call at :58,
~66 min for one at :01. Billing the open hour bills a fraction of it under that
hour's record_id, and the corrected re-run is then rejected by Lago as a
duplicate transaction_id, so the remainder is never billed at all. The bound
applies to both tables, because a window whose halves cover different hours is
the first bug over again. The caller-visible consequence is that the newest hour
arrives on the next run; this reader keeps no cursor, so pass a window
comfortably wider than the run interval.
A since resolving entirely inside the excluded hour ("30 minutes") now warns and
reads nothing, rather than spending warehouse time to return zero rows that say
nothing about whether there was traffic.
Closes the session-timezone dependency too. The bound no longer goes through
current_timestamp(), and the literal is rendered zone-explicit
(TIMESTAMP '... +00:00'): a bare literal is parsed in the warehouse's own
spark.sql.session.timeZone, so on a workspace set to anything but UTC the
identical literal named a different instant and the whole window slid by that
offset. Verified live that the suffixed form is accepted and resolves to the
same epoch under this warehouse's Etc/UTC.
The interval string no longer reaches SQL at all, but stays validated strictly:
a window quietly read as something other than what the caller wrote under-reads.
Re-verified live end-to-end: the same backfill over 2026-08-06 still bills
anthropic 31,018 against 31,091 consumed and openai 15,589 against 15,637 — the
byte-identical figures from the cache de-overlap fix, so the money path is
untouched.
A rejected external call is logged with NULL tokens and an empty `destination_model`, so its key can never match an `external_model_spend` row — the call bought nothing and Databricks meters no dollars for it. The BYOK indexing loop had no zero-usage guard, so those rows became buckets that fell through to the "no spend row yet ... re-run this window later to bill them" warning, which for them is advice that can never work. Measured live over a fully-aggregated window: 29 buckets reported, 28 of them phantoms (all 83 zero-usage rows in the window were 4xx/5xx), and the one example bucket the warning names for the operator was a phantom — so the single genuinely lagging bucket was the least visible thing in the message. The hosted loop below already applied this guard. Also fixes a wall-clock-dependent test shipped with the window work: it drove the open-hour guard through the real clock, so "30 minutes" spanned two hours at :34 and one at :04 and the assertion passed or failed on the current minute. The collapse is now pinned against the frozen `_NOW`, and the guard itself is driven by a `since` half an hour ahead of the clock.
`backfill_databricks` returned {"cost": …, "tokens": …, "skipped": 0} while BYOK
buckets whose external_model_spend row had not landed went unbilled, and
`on_error` never fired. Measured live with one hour's spend withheld — the shape
of real spend-table lag — 54 buckets were unbilled and the only visible
difference was `cost` falling 66 -> 12, which a caller cannot tell from a
quieter window.
The return value now carries `deferred`, and both it and `skipped` are reported
through `on_error` (where="backfill"), the hook every other billing gap uses.
`DatabricksSource.deferred_buckets` exposes the same buckets to a caller who
reads the window itself, since an already-read list cannot report a bucket the
reader never yielded.
`read_usage` ran `SELECT * FROM system.ai_gateway.usage` — 36 columns to bill off 14. The Statement Execution API's default `disposition=INLINE` FAILS a statement whose response exceeds 25 MiB rather than paginating past it, so the width of the projection sets the largest window this reader can handle, for a module whose own guidance is "read one wide window per run". Measured live over 247 real rows: 1,411 bytes/row for `SELECT *` against 435 for the columns actually read — a ceiling of ~18k rows where it should be ~60k. The dropped columns are the wide ones nothing bills off (`routing_information`, `endpoint_metadata`, `url`, `user_agent`). Re-read the same window both ways: 120 billable events, identical as a multiset, 120 unique transaction ids, same deferred buckets. The coupling runs the other way too, which is why this needs a test: a column missing from the projection reaches the adapter as absent, and every field degrades to zero/empty rather than raising — a silently under-billed event. So the test feeds the canned rows THROUGH the statement's own column list and asserts the events match an unprojected read. `ORDER BY event_time` dropped with it. Nothing downstream reads the row order: the BYOK join is keyed, the unbilled-bucket report is sorted, and each event's `transaction_id` derives from the row's own ids.
The Databricks drift sweep stopped at the row's columns. `token_details` is a STRUCT read field-by-field, so a key the adapter does not name reached neither a CanonicalUsage metric nor `extras`. Measured against the live table with the struct evolved by two fields (`cache_read_5m_input_tokens: 77`, `output_audio_tokens: 42`): 119 real tokens vanished with no error and no on_error — the exact failure the drift contract exists to prevent, and every drift test passed because none of them looked inside the struct. Latent today: the live struct has exactly the three fields the adapter maps, verified with DESCRIBE. Fixed anyway because this is the only column on the table that breaks tokens out by kind, so a new cache tier or output modality can land nowhere else — and this table's schema does evolve (`service_type`, `mcp_metadata`, `invocation_metadata` are later additions, and old rows still read `service_type = NULL`). Dotted key, not the container swept whole under `extras["token_details"]`: three of its keys ARE mapped, so publishing the container would re-emit counts already billed. Same shape openai_native already uses for its `*_tokens_details` containers. Both directions are pinned — one test fails if the sweep goes, another if it stops excluding the mapped keys — plus one for the JSON-string path, which is the one the backfill actually uses: the Statement Execution API serializes every STRUCT column as a string, measured, never as a dict. Scope, measured over a real window: the swept key reaches every hosted event (54 of 54) and no BYOK event (0 of 66), because a BYOK event is an hourly spend aggregate whose per-request extras are dropped by design. Both tables share this struct, so a new field still surfaces wherever there is hosted traffic.
The JS port's stream wrapper is a module-level generator taking `providerHint` as an argument, which was defaulted to `""` — a legitimate value, so an omission silently billed a Databricks-HOSTED call as `provider: "openai"`. That argument is now required there. This wrapper's equivalent is nested inside `wrap_openai_client` and reaches the adapter only through `_emit_from`, so the hint is closed over and there is nothing to forget. Comment only; the four end-to-end tests that pin the stamp already live in `test_wrapper_openai.py`.
Two guards for failures that are cheap to prevent and unrecoverable once they
happen.
A negative `usage_quantity` billed a $0 event and burnt the row's idempotency
key with it. The spend loop guarded on `if not usd`, which skips 0.0 but passes
-0.0042 straight through. `_parse_price` rejects a negative, so
`compute_precomputed_cost` floors the event to $0 — and that $0 event still
consumes the `record_id`-derived `transaction_id`, which Lago enforces unique
account-wide.
Driven end to end against real Lago on a real spend row: a $0.015245 row
restated negative was stored as `value: "0"`, and re-running the window once
Databricks had corrected it came back `422 value_already_exist` with Lago still
holding "0" — the same figure billed fine only under a different event-id
prefix. After the fix the negative row yields nothing, Lago 404s on that id, and
the corrected re-run stores `value: "0.015245"`.
The row is logged with its figure, model and hour rather than dropped quietly:
it means Databricks issued a credit or restatement, which this connector cannot
represent as an event, so skipping it leaves the customer billed more than
Databricks metered. Its bucket then surfaces in the deferred report — the same
window returns `{"cost": 65, ..., "deferred": 1}` with one on_error naming the
hour, where before it returned `deferred: 0` and looked complete.
Unobserved on the live table: 0 of 64 spend rows are negative, minimum
$0.0000036. Guarded anyway — one comparison against a failure with no recovery.
`LagoSDK(cfg)` sent every event to production Lago with an unusable key. The
first positional parameter is `api_key`, so the config becomes the bearer token
while `config` stays None and a fresh default replaces every field the caller
set. Driven live with a config naming a local Lago: the SDK posted to
api.getlago.com, every event 401'd, `flush()` still returned True, and the
caller's own on_error was never invoked — it was one of the discarded fields.
The only trace was a WARNING per event. Now a TypeError at construction, naming
the correct call.
Also drops the duplicate `import json` in `_canonical_tags`; the module has
imported it at the top since it grew a row-hash fallback.
The two ports disagreed about which HTTP failures destroy events, and measuring
both showed neither was right. `_PERMANENT_STATUSES` routes a batch to
`_send_individually`, where each event that fails again is logged and dropped
for good. Python listed {400,401,402,403,404,409,413,415,422}; JS listed
{400,404,409,422}.
Driven over a real socket at a server returning each status
(`probes/t11_status_matrix`), the two conditional cases point opposite ways:
a key rotated back after 3s PY destroyed all 5 events inside the first
second, none ever reached Lago. JS held them
and delivered all 5 when it healed.
nginx-style 413 above a byte PY's split path delivered all 5. JS held the
limit, 200 below it oversized batch, delivered 0, and stalled at
the backoff ceiling forever.
Both ports now use {400, 409, 413, 422} and the whole matrix is identical
row-for-row. The rule is written down rather than enumerated: is what makes
this fail a property of the BATCH? A different payload is the only fix ->
permanent, and splitting saves the good events. An out-of-band fix — a key
restored, an invoice paid, a URL corrected, a proxy reconfigured -> transient,
because dropping is unrecoverable while holding is bounded by
`max_buffer_size`, oldest-first and reported.
Per status, against a real Lago instance rather than from the RFCs:
404 -> transient. Lago answers 404 `resource_not_found` for a wrong PATH,
i.e. a mistyped api_url. Neither port held it, so a typo destroyed
every event. Same class as the 405/410 both already held.
415 -> transient. Splitting provably cannot help: this client always sends
application/json, so every isolated send fails identically (measured,
all 5 dropped). Lago answers 422 to a bad content-type, so a 415 only
comes from a proxy someone can fix.
413 -> stays permanent, and it was missing here on the JS side. Lago answers
422 `too_many_events` to an oversized batch (probed at 20k events /
3.5 MiB), so a 413 only comes from something like nginx's
`client_max_body_size` — exactly where splitting recovers the events.
402 -> transient. Payment required is a property of the account and stops
being true the moment someone pays. Measured: 5 in, 6 HTTP calls out,
0 recoverable, one on_error for the lot.
401
403 -> transient. Already correct in JS, never ported here.
Unchanged and re-confirmed live: a replayed transaction_id is 422 and stays
permanent, so one bad id still does not take its batch down with it.
Five new held-until-it-heals cases, confirmed to fail with the old set. The
402 half of this supersedes PR #21, which is now redundant on that point only —
its empty-api_url report and its Cloudflare usage_metadata drift sweep are NOT
on this branch.
…ing TTL
Four fixes from the review round that never reached this branch. All four were
verified live against the real code before and after, not reasoned about.
An explicitly-passed falsy `api_url` silently resolved to PRODUCTION. Preferring
the config value over "" is right — `requests` raises MissingSchema, which is
not a LagoApiError, so the queue classified it transient, re-prepended the batch
and retried at the 60s ceiling forever, stopping all billing with a growing
buffer as the only symptom. But LagoConfig's default is the production URL, so
`api_url=os.environ.get("LAGO_API_URL", "")` with the var unset resolved there
with no on_error and no log:
resolved api_url https://api.getlago.com/api/v1
client POST target https://api.getlago.com/api/v1/events/batch
on_error invocations 0
For a CI job or a developer holding a real production key that writes live
billing data, and ingested events cannot be un-ingested. The fallback is
unchanged, so the original config-clobber bug stays fixed; it is now reported
under `config.api_url`. An unpassed api_url stays silent — None means the caller
never mentioned it, and reporting the common case would train customers to
ignore the channel this fix depends on.
`usage_metadata` from the Cloudflare gateway got no drift sweep, and had already
lost two counters. `extras` was a fixed three-key dict, so any counter the
adapter does not map vanished with no error and no on_error — the one place
violating the contract test_drift.py enforces for the native adapters. Replaying
the 14 captured fixtures through the adapter:
neurons dropped in 4 entries Cloudflare's Workers AI billing unit
input_text_tokens dropped in 1 entry
and a live Logs API pull also returns `units`, a cost quantity that appears in
no fixture at all — the hand-maintained enumeration in the module docstring had
already drifted past reality, which is exactly the failure mode a snapshot
invites. Unmapped keys now sweep into extras["usage_metadata"] against an
explicit _MAPPED_USAGE_KEYS set. Deliberately NESTED rather than merged flat:
the poller reads extras["cached"] to decide whether to skip billing a request
Cloudflare served for free, so a future usage_metadata key called `cached` or
`step` must not be able to shadow it. The regression test iterates the fixture
directory rather than a fixed key list, so a recapture that introduces a new
counter fails it with no test edit. Closes #16.
`prime()` re-downloaded the ~400-model OpenRouter catalogue regardless of the
TTL. It set `_openrouter_stale` unconditionally, and it is reached from
`_auto_prime_pricing_for` on a matching wrap() and from warm_pricing() — both of
which a server can run per request — so pricing_ttl_seconds never applied on
that path at all, on the thread the queue drains events from. With the shipped
1-hour TTL: 4 prime()+maybe_refresh() cycles produced 4 full downloads where 1
was correct; now 1. Gated on the same "no table, or past the TTL" test lookup()
uses, so priming and looking up cannot disagree — and a table that genuinely
ages out is still re-primed, so prices do not freeze at the first fetch.
A failed pricing fetch retried every tick, forever, with no backoff, ahead of
the drain. Only the success path cleared a source's stale flag, so a bad
credential re-attempted on every tick, each attempt costing up to the 10s
`_get_json` timeout, all of it before the drain. Measured with a failing
Cloudflare fetch: 5 ticks produced 5 real requests and 5 on_error reports; now
1, and it still recovers once the window expires — a backoff, not a permanent
give-up, the same reasoning that makes a 401 transient in the queue. Per-source
1->2->4->...->60s, matching the queue's own send backoff, so one bad credential
cannot delay the three healthy tables; a success clears the window rather than
letting it keep doubling across unrelated outages.
The four fetches deliberately stay sequential here: this refresh runs on a
blocking daemon thread, where a thread pool for four fetches is a larger change
than the problem warrants, and the per-source backoff already removes the harm.
JS additionally parallelises them under Promise.allSettled, which its event loop
makes free — language-inherent, like os.register_at_fork vs AsyncLocalStorage.
Every fix has a test that fails when the fix alone is reverted (verified by
reverting each one in turn). 601 unit tests, coverage 91.89%, ruff and mypy
strict clean.
ancorcruz
left a comment
ancorcruz
left a comment
There was a problem hiding this comment.
Review — Databricks connector
Scoped to what this PR adds or changes. A couple of things I found sit in code this PR doesn't touch; I've left those out and am tracking them separately rather than growing this thread.
CI green on all 6 jobs, 601 unit tests pass locally, and the branch is rebased on merged main. The live-workspace grounding is visible throughout and it's what makes the connector reviewable — the docstrings record measured numbers rather than assumptions, and several of them let me confirm a defect from the comment alone.
One blocker corrupts billing amounts (sdk.py:411) — hosted rows are billed roughly twice for cached tokens. Two more lose or mis-price money on real Databricks surfaces (databricks_gateway.py:167, wrappers/openai.py:107), and one is a regression in the prime() change (pricing.py:1173).
Worth stating plainly because the PR's reconciliation claim is otherwise solid: "11 of 11 hosted models match token-for-token" held only because no hosted fixture exercises caching or reasoning. I ran all 22 — not one hosted row has a non-zero cache_read, cache_write or reasoning. The measurement was honest; it just never touched the path that's wrong.
Three cleanups not worth their own threads:
pricing.py:1183—_in_backoffhas no production caller.maybe_refreshinlines the same predicate as a local_ready(with a comment saying why), and the four real callers are all intest_pricing.py. The canonical documentation for the backoff feature currently sits on a test-only accessor.sdk.py:495and:547—now = at if at is not None else int(time.time())can't take the else branch:emit()resolvesat = self._event_time(timestamp)before every branch and_event_timealways returns anint. Both docstrings say the fallback covers nothing today. Making the parameterat: intlets mypy enforce what the comment asks reviewers to remember.databricks.py:799—_merge_usagedoes_as_bucket(a), which copies all eleven numeric fields, then immediately overwrites every one witha+b. Two fullsetattrpasses per merge, once per extra request in every BYOK bucket, on the wide windows this module tells operators to use.
|
|
||
| if usd_cost is not None: | ||
| breakdown = compute_precomputed_cost(usd_cost, markup_value) | ||
| elif usage.provider in TOKEN_BILLED_PROVIDERS: |
There was a problem hiding this comment.
Blocker: this routes hosted rows into an emitter that doesn't de-overlap, so cached tokens are billed twice.
The reasoning for TOKEN_BILLED_PROVIDERS is right — Databricks publishes no per-token rate, so token counts are the honest answer. The problem is the destination. _emit_token_events bills straight from nonzero_numeric() with no overlap removal, and this connector's own adapter docstring (databricks_gateway.py:45-54) states that this table's input_tokens includes both cache_read and cache_write, and warns that using these counts unsubtracted "would over-bill 3.04x".
Take the shape the captured fixtures show (input_tokens=1651, cache_creation_input_tokens=1642, output_tokens=4, total_tokens=1655) on a hosted cache-capable entity — system.ai.databricks-claude-sonnet-4-5, which this module names as exactly that. Three events are emitted: input=1651, cache_write=1642, output=4. That's 3,297 tokens billed for 1,655 consumed. Reasoning does the same on a hosted thinking model (output=200 + reasoning=150 → 350 for 200).
What makes this worth blocking on is that the machinery to prevent it is in this PR: _token_semantics, deoverlapped_token_total, and _OPENAI_SHAPED_APIS — which literally contains "databricks_gateway". Price mode consumes all of it; the token path consumes none.
It also reaches the live wrap() path, where provider_hint="databricks" arrives with api="chat_completions", matching neither _OPENAI_SHAPED_APIS nor either _INCLUDES_ set — so the semantics are unknown there rather than merely unapplied.
Note the un-de-overlapped emitter itself predates this PR, so the general case is out of scope here and I'm raising it separately. The in-scope decision is this branch: either apply the same de-overlap before emitting, or give hosted rows a path that knows this table's semantics. A hosted fixture with non-zero cache tokens would pin it either way — there isn't one today, which is why the reconciliation passed.
| destination_type = _safe_str(row.get("destination_type")) | ||
| destination_name = _safe_str(row.get("destination_name")) | ||
|
|
||
| if destination_type == _HOSTED_DESTINATION_TYPE: |
There was a problem hiding this comment.
Blocker: the hosted/BYOK fallback points the money-losing way.
The split is a single == against PAY_PER_TOKEN_FOUNDATION_MODEL, with everything else falling through to BYOK. The docstring justifies that with the two types observed live (PAY_PER_TOKEN_FOUNDATION_MODEL, EXTERNAL_FOUNDATION_MODEL, plus NULL on pre-routing rejects), which is fair for what was measured — but the AI Gateway also fronts provisioned-throughput and custom registered-model endpoints.
For those, destination_type is neither value, so provider becomes api_type.split("/")[0] — "mlflow". read_usage then files them in the BYOK tokens index (they aren't provider == "databricks", so the hosted loop skips them), and external_model_spend is external-model-only by definition, so no spend row ever matches. Result: every such request sits in deferred_buckets permanently, gets reported as "the spend table lags; re-run this window later", and is never billed at all.
Re-running can't fix it, so this is the same shape as the hour-boundary issue below: a false gap that never closes, and a deferred counter that never returns to 0 — which the docstring says is the signal that the whole window billed.
Inverting the default fixes it: an unrecognized destination_type should be treated as hosted/token-billed. Token counts for something Databricks meters differently is a recoverable under-report; a row that no loop bills is silent lost revenue. Worth a _report_error on the unrecognized value too, so a new endpoint class shows up as a question rather than as a bucket that quietly never resolves.
| # Matching `/ai-gateway/mlflow/` specifically, NOT `/ai-gateway/`, is the whole | ||
| # point: the openai and anthropic surfaces live under the same prefix and must | ||
| # keep their real vendor provider so they price against OpenRouter. | ||
| _DATABRICKS_HOSTED_PATH = "/ai-gateway/mlflow/" |
There was a problem hiding this comment.
Blocker: matching only /ai-gateway/mlflow/ leaves Databricks' own documented base_url unhinted.
The docstring argues — correctly — that stamping "databricks" turns a possible 2.5–5x under-bill into a guaranteed honest miss, because OpenRouter lists bare openai/gpt-oss-20b and meta-llama/llama-4-maverick at a fraction of Databricks' DBU rate. But OpenAI(base_url=f"{HOST}/serving-endpoints") with model="databricks-meta-llama-3-3-70b-instruct" is Databricks' canonical documented way to call a pay-per-token foundation model, and it doesn't contain /ai-gateway/mlflow/.
On that path _provider_hint_for returns "", _infer_provider gives "openai", the row misses TOKEN_BILLED_PROVIDERS, and lookup_openrouter can match the bare listing — the exact outcome this function exists to make structurally impossible. The widened date-stripping in _strip_version makes an accidental match more likely, not less.
Adding /serving-endpoints to the match closes it. Worth checking the same question for the Anthropic wrapper, which takes no hint at all — Databricks hosts Claude as system.ai.databricks-claude-sonnet-4-5, and if that's reachable via an Anthropic-shaped client it would price against Anthropic's own rate.
There was a problem hiding this comment.
Won't fix. /ai-gateway/mlflow/v1 is the surface we cover, now stated in the README.
Two reasons. It can't bill wrong today: all 8 live hosted endpoints echo a 6-digit-dated name (gpt-oss-20b-080525) and we strip only 8-digit dates, so 0 of 8 price and every call falls back to token events. And that surface isn't billable at all — system.ai_gateway.usage holds only AI-Gateway traffic, so backfill_databricks can't see it either. The hint would fix attribution on usage we can't bill.
Real today: provider="openai" plus a spurious on_error per call. Noted, not fixed.
Anthropic wrapper stays hintless — that path is BYOK-only, so provider="anthropic" is correct.
| self._cloudflare_stale = True | ||
| elif key == "mistral": | ||
| self._mistral_stale = True | ||
| if self._is_cold(self._mistral_aliases, self._mistral_fetched): |
There was a problem hiding this comment.
Regression: gating prime() on _is_cold removes the only mechanism that re-fetched after a credential arrived.
TTL-gating the flags is the right fix for the eager-refetch problem. But it also silently disables the learned-key path. fetch_mistral_aliases returns {} when no key is configured (line 1059), and that empty table gets stamped _mistral_fetched = time.time().
Sequence: a provider="mistral" row arrives (Mistral traffic backfilled through a Cloudflare gateway keeps provider="mistral", as this module notes) → lookup() sets _mistral_stale → refresh runs with no key → _mistral_aliases = {}, marked fresh. Then sdk.wrap(mistral_client) calls learn_mistral_api_key(key) followed by prime(["mistral"]). Previously that set the flag unconditionally and the next tick fetched with the learned key. Now _is_cold({}, now) is False — the table isn't None and isn't stale — so the flag is never set and alias resolution stays dead for a full pricing_ttl_seconds (3600s by default). Every Mistral alias lookup misses and degrades to token events for the first hour.
Two options: have learn_mistral_api_key set _mistral_stale = True itself when it actually adopts a key (it knows a credential just became available, which is new information the TTL can't represent), or treat an empty table as cold. The first is narrower. A test asserting that learn_mistral_api_key + prime after an empty fetch produces a refetch would pin it.
|
|
||
| for row in spend: | ||
| usd = _safe_float(row.get("usage_quantity")) | ||
| if usd <= 0: |
There was a problem hiding this comment.
An unparseable usage_quantity is skipped in total silence, then misreported as a bucket whose spend row hasn't landed yet.
_safe_float deliberately returns 0.0 on TypeError/ValueError so one bad row can't abort the generator — right call. But usd <= 0 then continues, and only the usd < 0 sub-branch logs anything. A value the driver hands back in a form float() rejects therefore produces no log line, no on_error, and no event.
The compounding part: billed_keys.add(key) sits below the continue, so that bucket also lands in unbilled. Both read_usage and backfill_databricks then tell the operator "no external_model_spend row yet — re-run this window later to bill them". The spend row is right there; re-running produces the identical silent skip. Real metered dollars are lost behind an instruction that guarantees they stay lost.
Distinguishing "parsed to zero" from "failed to parse" is enough — _safe_float could return None on failure, and this branch could report it through on_error the way the negative-spend guard already does.
| ) | ||
|
|
||
|
|
||
| def test_no_real_databricks_workspace_hosts() -> None: |
There was a problem hiding this comment.
The scrubber misses the two identifiers that are actually in all 22 fixtures.
Good addition, and the threat model in the docstring is the right one. It just doesn't cover what's there. Verified on the committed fixtures:
"account_id": "73357fbe-b495-49ad-9eec-ad226318df4f" (22 of 22 files)
"workspace_id": "7474648573314045" (22 of 22 files)
Neither column is read by any adapter, so they're pure capture residue — and they're squarely inside the docstring's own framing ("whatever the provider chose to log about the operator"). This package publishes to PyPI, so they become permanent public identifiers for your Databricks account. Not credentials on their own, but they're the account coordinates an attacker pairs with anything else they find.
Two smaller gaps in the same test: the host pattern only recognises the AWS dbc-… form, so an Azure adb-<workspace-id>.N.azuredatabricks.net capture would pass unscrubbed; and examples/.env.example:39 names a concrete-looking warehouse id (a292ad231ac2d202) in its "e.g." comment, though the variable itself is correctly left empty.
Separately and worth confirming rather than assuming: was any real PAT ever pasted into a capture or notebook during this work? If so it wants rotating regardless of what's committed now, since git history keeps it.
| """ | ||
| counts = {"cost": 0, "tokens": 0, "skipped": 0, "deferred": 0} | ||
| reader = source if hasattr(source, "read_usage") else None | ||
| rows = reader.read_usage(since, event_id_prefix=event_id_prefix) if reader else source |
There was a problem hiding this comment.
since and event_id_prefix are silently dropped on the already-read-iterable path.
rows = reader.read_usage(since, event_id_prefix=event_id_prefix) if reader else source — when source is an iterable of DatabricksUsageRow, which the docstring explicitly recommends so the expensive window is read once, both arguments vanish with no error.
since being ignored is arguably fine (the caller already chose the window). event_id_prefix is not: the rows carry whatever prefix they were built with, so a caller passing event_id_prefix="dbx_v2" to re-bill a window Databricks has since restated gets the old prefix baked into every event_id_for(), and Lago rejects the entire re-run as duplicate transaction_ids. That's the precise failure event_id_for's own docstring and the negative-spend guard were both written to avoid, and the prefix argument is the documented way to escape it.
Either raise when a non-default event_id_prefix is passed with a pre-read iterable, or re-stamp prefix on the rows as they go by.
| text = str(value).strip() | ||
| if text: | ||
| return text | ||
| digest = hashlib.sha256( |
There was a problem hiding this comment.
The content-hash fallback hashes the SELECT projection, so editing _USAGE_COLUMNS re-bills every id-less row.
The docstring promises "the content hash keeps the key deterministic, so re-running the same window is still idempotent". But the dict being hashed is exactly what _USAGE_COLUMNS projects — and line 71 documents that tuple as something maintainers must edit ("Keep this in sync with extract_databricks_log").
So adding one column changes sha… for every row where invocation_id and request_id are both NULL. The next re-run of an already-billed window double-bills those rows instead of being rejected as duplicates, and the trigger is a routine maintenance edit with no visible connection to billing.
Hashing a deliberately chosen stable subset — event_time + endpoint_id + the token counts — keeps the key stable across projection changes and still distinguishes rows. Worth a comment saying the subset is load-bearing, so it isn't "tidied" later.
| # Drift sweep one level down, into the *_tokens_details sub-objects. Without | ||
| # this, an unrecognized nested field is silently dropped (see | ||
| # _MAPPED_DETAIL_FIELDS) because its container is a known top-level key. | ||
| for container, mapped in _MAPPED_DETAIL_FIELDS.items(): |
There was a problem hiding this comment.
The drift sweep writes to usage.extras, which nothing reads — so the contract it enforces is satisfied only inside the tests.
The reasoning for sweeping one level into *_tokens_details is sound, and the cache_write_tokens analysis (inside prompt_tokens, so mapping it would double-bill 2.24x) is exactly the kind of thing worth writing down.
But extras has no consumer. _emit_token_events builds properties from value/model/provider/api plus dimensions; _push_cost_event from model/provider/api/price_source/markup/unit/value/base_cost plus dimensions. Neither includes extras, and nothing passes it to _report_error. So the live prompt_tokens_details.cache_write_tokens: 3022 the comment says "vanished with no error" now lands in an in-memory dict that's discarded when emit() returns — externally identical to before the fix.
The same applies to the sweeps in cloudflare_gateway.py and databricks_gateway.py; the Cloudflare one fires on essentially every Workers AI row, since neurons and input_text_tokens are unmapped.
The contract belongs where it's claimed: one _report_error(..., "drift") in emit() when extras holds an unmapped-count key would make all three sweeps real, and would replace three per-adapter dictionaries feeding a dead field. As it stands, _as_bucket's _BUCKET_INVARIANT_EXTRAS filter is work done on data nothing consumes. (The emitters themselves predate this PR — flagging it here because this is where the guarantee is asserted.)
| source = DatabricksSource.from_env() # DATABRICKS_HOST / _TOKEN / _WAREHOUSE_ID | ||
| print(sdk.backfill_databricks(source, "7 days", default_subscription="sub_default")) | ||
| sdk.flush() | ||
| # {'cost': 60, 'tokens': 47, 'skipped': 0} |
There was a problem hiding this comment.
The documented return shape omits deferred, which this PR makes the primary billing-gap signal.
The snippet prints # {'cost': 60, 'tokens': 47, 'skipped': 0} while the implementation returns a fourth key, and backfill_databricks's own docstring says "a run with both at 0 is the only one that billed the whole window".
A caller writing the automated reconciliation check this counter exists for would assert on the three documented keys and never look at the fourth — which is precisely the "cost alone dropping from 66 to 12 is not something an automated caller can read as a gap" failure deferred was added to fix.
|
Fixed, both ports. Each test verified failing with its fix reverted.
Two found while measuring yours. Not fixed — measured, nothing can reach them.
On the bucket key: Four we disagree on.
Duplicate helpers: 11 / 7 / 2 / 1. Second time on this one. What you didn't find. Four of your round-2 findings are in both repos and were filed against one. And Databricks documents external model spend as unsupported for the Custom provider, so those rows can never get a spend row — the never-billed class is documented, not a guess. Adding a cross-port parity gate and a fixture matrix so absence stops reading as coverage. |
2 participants
Add Databricks AI Gateway connector
Databricks is the second gateway connector, after Cloudflare. Everything below was
established against a live workspace, not inferred from docs: 25 models exercised,
~100 calls, 226 real usage rows read over the SQL Statement Execution API.
What it covers
Databricks differs from Cloudflare in two ways that shape the whole design:
/compatendpointfronting every provider; Databricks makes each provider reachable only through its
own native surface — and two of those use the same
openai.OpenAIclass while needingdifferent price tables, so
base_urldiscrimination is load-bearing rather thancosmetic.
base_urlwrap()/ai-gateway/openai/v1external_model_spend/ai-gateway/anthropicexternal_model_spend/ai-gateway/mlflow/v1/ai-gateway/gemini/v1betaThe decision that shapes everything: gateway parity
What makes the Cloudflare connector trustworthy is that you can put Cloudflare's own
dashboard beside Lago and see the same numbers. Databricks makes that harder, because
hosted traffic appears on two surfaces in two units:
request_tags?custom_tagsis{}This connector mirrors the gateway's own surfaces. Hosted bills token counts
(matching the AI Gateway usage page); BYOK bills Databricks' own metered USD (matching
the external-model-spend view).
Hosted money is deliberately not billed from, and this is the part most worth
challenging in review. It is obtainable —
system.billing.usage×list_prices, oraccount_pricesfor an account's contract rate — so "hosted USD is impossible" would bewrong; that applies only to the tokens→DBU rate, which is published on an HTML page and
exists in no system table (verified by searching every column of all 88 of them). The
reason not to use it is a product one: it comes from a different Databricks screen than
the gateway view, carries no attribution (so per-subscription splits would be ours rather
than Databricks'), and lags ~19h — measured,
max(usage_start_time)inbilling.usagewas
2026-08-10T17:00whilemax(event_time)inai_gateway.usagewas2026-08-11T10:09. Every number this connector sends is one you can find on a Databricksgateway page.
Two earlier approaches to hosted pricing were built and reverted, both recorded in
CHANGELOG.md: vendoring the 18-model DBU rate card (hand-maintained price data in acodebase whose every other source is live HTTP), and solving the rate from the customer's
own tables (worked — all 6 solvable endpoints recovered the published rate to three
decimals — but the SQL warehouse needed to run the solve costs ~1,500× the usage it
prices: $6.54 of SQL against $0.0043 of MODEL_SERVING in this account).
Evidence
captured response body was run through the real pipeline (
extract_*_native→compute_costat live OpenRouter rates) and compared against Databricks' ownusage_quantity, joined on Databricks' own grouping key. 13 models, both cacheconventions, four reasoning models, costs from $0.0000036 to $0.015245.
byte-identical
transaction_ids across a re-run.(3,640 in / 4,344 out), confirmed against a real Lago instance.
qwen35-122b-a10bat 48 in /204 out, which is exactly what falls out of dividing that endpoint's DBU rows by the
published rate card — the tokens billed agree with what Databricks charged DBUs for,
via a completely separate table.
Commits — the first two are not Databricks-specific
Commits 1–2 fix bugs that affect every price-mode user, found while validating this
connector, and each passes the full suite on its own (462 → 474 → 559 tests). If you
would rather land them separately, say so and I will split them into their own PR
underneath this one — no re-work needed.
_strip_versiononly matched Anthropic's compact date (-20250929), not OpenAI'shyphenated one (
gpt-5-2025-08-07). Sinceresolve_modelprefers the response's ownname,
create(model="gpt-5")resolved to a name that matched nothing and fell throughto token events. Verified against the live OpenRouter table:
gpt-4.1,gpt-4.1-mini,gpt-5,gpt-5-mini,o3,o4-miniall missed.gpt-4olooked fine only by luck.extrasswept onlytop-level usage keys, and
prompt_tokens_detailsis itself a known key, so nothingnested was ever inspected: a live
gpt-5.6-solresponse'sprompt_tokens_details.cache_write_tokens: 3022was discarded with no error. Everydrift test passed, because none looked inside a details object.
Where to look closely
gateway/databricks.py— the money paths. Four ways this read loses money silently(chunk-0 truncation, double billing across the two tables, unscoped idempotency keys,
and rows whose ids collapse to an empty string) are each guarded and each has a
regression test naming the failure.
gateway/adapters/databricks_gateway.py— three naming quirks that a docs-only readinggets wrong, all caught by real rows. In particular
destination_namemeans the modelfor hosted rows but a credential name for BYOK, so a single fallback rule bills
workspace.default.anthropickeyas the model.sdk.pyTOKEN_BILLED_PROVIDERS— a deliberate, narrow exception to "never silentlyunder-bill". Reasoning is in commit 5's message; the invariant still holds for every
miss a customer could act on.
input_tokensnote in the adapter docstring — this table's input count includescache tokens, the inverse of the providers' own response bodies. Nothing computes from
it today, and the docstring records why a computed fallback would need a per-provider
correction rather than a uniform one.
Gates
ruff check+ruff format --check+mypy --strictclean; 559 tests, 90% coverage(
gateway/databricks.pyat 100%). 22 fixtures are real captured rows — neverhand-written — and a sweep test iterates the whole directory so a capture no named test
mentions still asserts something.
Deliberately out of scope
this does not remember where it got to.
but every request past it returns
500with an empty body — including Databricks'own documented code sample, and including
:countTokens, which involves no inference.A genuine upstream failure on this gateway is richly wrapped; these carry none of that,
so the gateway throws while constructing the call. Needs a Databricks support ticket,
not SDK work. Reproduction ids and five saved 500-responses are recorded.
Known gaps
system.ai_gateway.usagestores the requested alias (gpt-5.6) while OpenRouter listsonly the resolved name (
gpt-5.6-sol), so that one model prices live and misses onbackfill — falling back to token events, never mispriced. Needs an alias step.
api="chat_completions"; the shape detector only knows ChatCompletions vs Responses. Numerically correct, mislabelled.