Feature/databricks gateway connector

CONTRIBUTING.md

@@ -70,6 +70,7 @@ uv lock --upgrade-package X  # bump a single package
 - `src/lago_agent_sdk/canonical.py` — the normalized usage shape sent to Lago
 - `src/lago_agent_sdk/queue.py` — async event queue with backoff
 - `src/lago_agent_sdk/lago_client.py` — thin HTTP client to `/events/batch`
+- `src/lago_agent_sdk/gateway/` — second front door: gateway usage logs → `CanonicalUsage`, for backfill
 - `tests/unit/` — unit tests, organized to mirror `src/`
 - `tests/unit/adapters/fixtures/` — captured real provider responses, used by adapter tests
 - `tests/integration/` — live tests, gated on credential env vars
@@ -84,6 +85,83 @@ uv lock --upgrade-package X  # bump a single package
 6. Add unit tests against the captured fixtures.
 7. Add a live integration test gated on the provider's API key env var.
 
+## Adding a gateway
+
+`gateway/` is a **second front door** into the same kernel, separate from the provider-native
+`adapters/` used by `wrap()`. A gateway connector reads a gateway's own usage log and maps it into
+`CanonicalUsage` for backfill; there is no client to patch. Two exist: Cloudflare and Databricks.
+
+1. Capture real rows/entries from a live gateway into
+   `tests/unit/gateway/adapters/fixtures/<gateway>/`, one file per scenario. Cover both success and
+   every failure shape you can produce — failed calls are where the surprises live.
+2. Write `src/lago_agent_sdk/gateway/adapters/<gateway>.py` exporting
+   `extract_<gateway>_log(entry) -> CanonicalUsage` and `resolve_<gateway>_subscription(entry) -> str | None`.
+   Keep it a **pure function**: no HTTP, no SDK state.
+3. Export both from `gateway/adapters/__init__.py` under explicitly gateway-scoped names, so no
+   gateway is the implicit default.
+4. Add `tests/unit/gateway/adapters/test_<gateway>.py` against the captured fixtures.
+5. Add a `## <Gateway> AI Gateway` README section and a `CHANGELOG.md` entry.
+6. Add `examples/<gateway>_gateway_demo.ipynb` showing backfill and live calls.
+
+### A connector is only as good as the comparison
+
+The reason the Cloudflare connector reads well is that you can put the gateway's own
+dashboard beside Lago and see the same numbers. Two rules protect that, and both were
+learned the hard way on Databricks:
+
+- **Emit the gateway's own grouping key as a dimension.** Our `model` is normalized; the
+  gateway's page is not. Group Lago by one and the dashboard by the other and the
+  comparison fails on naming alone, before any number is even wrong. Attach the key the
+  gateway's surface aggregates by, and only keys that are true of the whole row — a
+  per-request field on an hourly aggregate is one sampled value dressed up as a property
+  of the bucket.
+- **Never bill from a surface the gateway UI doesn't show.** Databricks does expose exact
+  dollars for its hosted models, in `system.billing.usage` x `list_prices` — on a
+  different screen, with no attribution tags, about a day behind. Billing from it would
+  produce a number the customer cannot find anywhere, which costs more trust than the
+  feature adds. Hosted therefore bills token counts, matching the page they do look at.
+
+### Does the read itself belong in the SDK?
+
+Default: **no.** The adapters stay pure and the fetching lives in the example notebook, as Cloudflare's
+does — its whole read is one paginated GET, and an SDK wrapper around that would be indirection for
+nothing.
+
+Databricks earned the exception, in `gateway/databricks.py` (a sibling module, so the adapter stays
+pure). The bar it cleared, and the one to hold a third gateway to: the read is long enough that a
+customer will reimplement it wrong, and the ways it goes wrong lose money silently. Databricks needs a
+SQL warehouse, the Statement Execution API, columnar-to-dict zipping, chunked result fetching and two
+tables reconciled against each other — and the first hand-rolled version in the demo notebook truncated
+at chunk 0, which bills a fraction of a wide window with no error at all. If a gateway's read is a loop
+over one endpoint, leave it in the notebook.
+
+When it does clear the bar: name it `gateway/<gateway>.py`, expose a `<Gateway>Source` with an explicit
+window and a `read_usage()` that yields rows already shaped for `emit()`, add the `backfill_<gateway>()`
+one-liner to `LagoSDK`, and use a dependency that is already core (`requests` / `undici`). No scheduler,
+no cursor store, no credential store — that is the poller, and it stays out of the SDK.
+
+### Things both existing connectors had to get right
+
+These are the traps, and every one of them cost real debugging:
+
+- **Which cost is authoritative.** Gateway traffic bills from the *gateway's* metered cost, not one we
+  compute — it keeps Lago reconcilable against the dashboard the customer looks at. Note the gateway
+  may under-report: Cloudflare's `cost` omits additive reasoning tokens, measured at 22.8x low on a
+  real call.
+- **Token semantics are per-gateway, not per-vendor.** Cloudflare passes Anthropic's cache counts
+  through *additively*; Databricks' table folds them *into* `input_tokens`. Same provider, opposite
+  conventions. Never assume the vendor's own convention survives the gateway.
+- **`provider` must be unmatchable when you cannot price it.** If a gateway bills on its own rate card,
+  stamp a provider that hits nothing in `_VENDOR_MAP` so the lookup misses honestly. Stamping a real
+  vendor name lets a near-miss model string match at 2.5-5x the wrong rate, silently.
+- **Idempotency keys must be subscription-scoped.** `transaction_id` is unique org-wide, so
+  `f"{prefix}_{subscription}_{row_id}"` — an unscoped id silently blocks a row from ever reaching a
+  second subscription.
+- **Failed calls appear in the log.** Extract them to all-zero so `nonzero_numeric()` is empty and
+  nothing is emitted, rather than billing zeros.
+- **Drift.** An unrecognized field must reach `extras`, including one level down inside nested
+  `*_details` objects. `test_drift.py` pins it.
+
 ## Pull request checklist
 
 - [ ] Unit tests cover the change

README.md

@@ -166,6 +166,95 @@ sdk.flush()
 
 See [`examples/cloudflare_gateway_demo.ipynb`](examples/cloudflare_gateway_demo.ipynb) for a runnable end-to-end version of both.
 
+**Gateway-routed calls are billed at the gateway's metered cost.** Cloudflare reports its own `cost` per log entry and the backfill passes that straight through, so Lago reconciles against the dashboard you actually look at. One measured consequence to be aware of: that field excludes additive *reasoning* tokens, so a thinking-heavy Gemini call bills about 4% of what Google charges (verified live at 22.8x on one call, 39.6x on another — the ratio tracks each prompt's thinking-to-output ratio). Cloudflare is exact on input, output, cache-read and cache-write.
+
+**If you hand-roll a poller, don't use `urllib`.** `gateway.ai.cloudflare.com` returns `403` with body `error code: 1010` to `Python-urllib` — its bot-signature check. Any other User-Agent passes, and `requests` (which this SDK uses) is fine. The failure looks like an auth error because the body is otherwise empty.
+
+## Databricks AI Gateway
+
+Unlike Cloudflare, Databricks has **no unified endpoint** — each provider is reachable only through its own native surface, and two of them use the same `openai.OpenAI` class. Which `base_url` you point at decides how the call is priced.
+
+**Databricks-hosted foundation models** (`system.ai.*`) — billed by Databricks in DBUs:
+
+```python
+from openai import OpenAI
+from lago_agent_sdk import LagoSDK
+
+sdk = LagoSDK(api_key="...", default_subscription_id="sub_acme")
+client = sdk.wrap(OpenAI(
+    api_key=DATABRICKS_TOKEN,
+    base_url=f"{DATABRICKS_HOST}/ai-gateway/mlflow/v1",
+    default_headers={"Databricks-Ai-Gateway-Request-Tags": json.dumps({"lago_subscription": "sub_acme"})},
+))
+client.chat.completions.create(model="system.ai.llama-4-maverick", messages=[{"role": "user", "content": "Hi"}])
+```
+
+**Your own vendor key (BYOK)** — Anthropic via its native passthrough, note `api_key="unused"` because the real credential goes in `Authorization`, and the Unity Catalog connection holding your Anthropic key is named in `Databricks-Model-Provider-Service`:
+
+```python
+from anthropic import Anthropic
+client = sdk.wrap(Anthropic(
+    api_key="unused",
+    base_url=f"{DATABRICKS_HOST}/ai-gateway/anthropic",
+    default_headers={
+        "Authorization": f"Bearer {DATABRICKS_TOKEN}",
+        "Databricks-Model-Provider-Service": "workspace.default.anthropickey",
+    },
+))
+```
+
+OpenAI BYOK is the same `OpenAI` class as the hosted example, against `/ai-gateway/openai/v1` with its own `Databricks-Model-Provider-Service`.
+
+### What gets billed
+
+| Path | Live `wrap()` | Backfill |
+|---|---|---|
+| BYOK (OpenAI / Anthropic) | **dollar cost**, priced from the vendor's published rates | dollar cost from Databricks' own `external_model_spend` |
+| Hosted (`system.ai.*`) | **token counts** | **token counts** |
+
+BYOK prices live because you pay the vendor directly, so the vendor's rate *is* your cost — verified against Databricks' own metered spend on 38 of 38 real buckets, exactly. Hosted models bill in DBUs against a rate card published only as HTML and present in no system table, so there is no rate to look up: those calls emit token counts instead of a dollar cost. That is the complete answer for them, not a degraded one, so it is **not** reported as an error — `TOKEN_BILLED_PROVIDERS` lists the providers this applies to, and the SDK notes it once per model at info level rather than warning on every call. A genuine price miss — a cold table, an unmatched model name — still reports through `on_error` as before.
+
+**Hosted dollars exist, and are deliberately not billed from.** `system.billing.usage` × `list_prices` (or `account_prices` for your contract rate) does yield exact USD per hour and endpoint. It is not used because it comes from a *different Databricks screen* than the gateway view: it carries no `request_tags`, so per-subscription splits would be ours rather than Databricks', and it lags the gateway by roughly a day — measured at ~19h on a live workspace. Every number this connector sends is one you can find on a Databricks **gateway** page, which is the property that makes it checkable.
+
+**Grouping matches the Databricks page.** Each backfilled event carries the grouping key of the surface it came from — `endpoint_name` for hosted, `bucket` (the hour) for BYOK. Group Lago by `endpoint_name` and you get the AI Gateway → Usage table row for row. Pass `dimensions={...}` to add your own keys; yours win on a name collision.
+
+**Don't run the live path and the backfill over the same hosted traffic.** Both emit token events, with different `transaction_id`s, so Lago accepts both and the counts double. Pick one per traffic stream: `wrap()` for real time, the backfill for completeness.
+
+`Databricks-Ai-Gateway-Request-Tags` is what makes attribution work. It lands in `request_tags` on `system.ai_gateway.usage` **and** is a first-class aggregation dimension on `external_model_spend`, so tagging `lago_subscription` means BYOK cost arrives already split per subscription — no apportioning needed.
+
+### Backfill — give it a window, it does the rest
+
+```python
+from lago_agent_sdk.gateway.databricks import DatabricksSource
+
+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}
+```
+
+Pass a `datetime` instead of `"7 days"` for an exact lower bound, and `unified=True` to bill the whole window to `default_subscription` regardless of per-call tags.
+
+Unlike Cloudflare's single paginated GET, this one is worth having in the SDK — hand-rolling it is ~100 lines with three money-losing traps in them. The Statement Execution API returns only **chunk 0** inline, so a wide window silently truncates and bills a fraction of it with no error. A BYOK call appears in **both** `ai_gateway.usage` and `external_model_spend`, so billing both charges twice. And `transaction_id` is unique account-wide, so an unscoped row id blocks that row from ever reaching a second subscription.
+
+To inspect a window before billing it, or to route rows yourself, read them directly — each row is already shaped for `emit()`:
+
+```python
+for row in source.read_usage("7 days"):
+    print(row.usage.model, row.subscription, row.usd_cost)  # usd_cost is None for hosted
+```
+
+Reading the system tables needs a PAT with the **`sql`** scope plus a SQL warehouse — the live calls above need neither. `examples/databricks_gateway_demo.ipynb` is a complete worked example of both halves. The pure `extract_databricks_log(row)` / `resolve_databricks_subscription(row)` functions stay available from `lago_agent_sdk.gateway.adapters` if you already have rows from `databricks-sql-connector` or your own warehouse job.
+
+**One cost note:** a SQL warehouse is a real cost centre. Measured on a test workspace, the warehouse queries cost roughly 1,500× the model-serving usage they were reporting on. Run the backfill as one query over a wide window, never as a tight polling loop.
+
+### Gotchas worth knowing
+
+- **`gpt-oss` models inflate input by ~100 tokens** from a server-injected preamble — a 2-character prompt bills 102. Not an SDK error.
+- **`claude-opus-4-5` does not cache through this gateway**: reproducibly `cache_read`/`cache_write` of 0 with the full prompt billed as input, on a request shape where `claude-sonnet-4-5` caches fine. An opus customer silently gets no cache discount.
+- **Hosted models report three different name strings.** `system.ai.llama-4-maverick` and `databricks-llama-4-maverick` both work as requests, and the response echoes a third (`meta-llama-4-maverick-040225`). Pricing keys off the resolved name, so reconciling by requested id will not line up.
+- **Embeddings** work on `/ai-gateway/mlflow/v1/embeddings` and report input only — no `completion_tokens` at all.
+
 ## Multi-tenant — pick a subscription per call
 
 Three ways to set the `external_subscription_id`, in priority order:

examples/.env.example

@@ -1,6 +1,9 @@
 # Copy this file to examples/.env and fill in real values.
 # examples/.env is gitignored — never commit real credentials.
 
+# ---------------------------------------------------------------------------
+# cloudflare_gateway_demo.ipynb
+# ---------------------------------------------------------------------------
 CF_ACCOUNT_ID=
 CF_GATEWAY_ID=
 CF_LOGS_TOKEN=
@@ -21,3 +24,22 @@ ANTHROPIC_API_KEY=
 # resolution too (see LagoConfig.mistral_api_key) — no separate credential
 # needed for that.
 MISTRAL_API_KEY=
+
+# ---------------------------------------------------------------------------
+# databricks_gateway_demo.ipynb
+# ---------------------------------------------------------------------------
+DATABRICKS_HOST=https://dbc-xxxxxxxx-xxxx.cloud.databricks.com
+# Part 2 (live calls) works with any token that has gateway inference access.
+# Part 1 (backfill) additionally needs the `sql` scope to read
+# system.ai_gateway.usage — without it every warehouse route returns
+# 403 "does not have required scopes: sql", including the Thrift path the
+# databricks-sql-connector uses, so changing client library does not help.
+DATABRICKS_TOKEN=
+# Backfill only. SQL Warehouses -> your warehouse -> Connection details.
+# Just the id here (the notebook builds the rest), e.g. a292ad231ac2d202.
+DATABRICKS_WAREHOUSE_ID=
+# Unity Catalog connections holding your own vendor keys (BYOK). Only needed for
+# whichever provider you actually call in Part 2. Three-level UC names, e.g.
+# workspace.default.anthropickey.
+DATABRICKS_PROVIDER_SERVICE_ANTHROPIC=
+DATABRICKS_PROVIDER_SERVICE_OPENAI=