Workers AI client + partner-model pricing (typesafe/jev)

CHANGELOG.md

@@ -2,6 +2,17 @@
 
 All notable changes to this project will be documented here. Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow [SemVer](https://semver.org).
 
+## [Unreleased]
+
+### Added
+
+- **Workers AI client: `sdk.workers_ai(account_id, api_token, gateway_id=..., gateway_auth=...)`.** The one provider with no client to wrap: the official `cloudflare` package percent-encodes the slash in every model name and "No route for that URI" is the answer for all of them, and partner models such as `typesafe/jev` reject the `messages` array the gateway's `/compat` endpoint requires. The SDK now ships a one-method client, `run(model, input, extra_lago=...)`, that reaches every Workers AI model and bills the response's own `usage`. Verified live on 2026-09-21: `typesafe/jev` answered through Cloudflare under BYOK and its 446 input / 73 output tokens landed in Lago.
+  - **Two routes, chosen per model id — measured, not chosen.** `@cf/...` ids go to the gateway host's path route, which answers with `cf-aig-cache-status` (a HIT is not billed) and `cf-aig-log-id` (emitted as the `cf_log_id` dimension for reconciliation against the Logs API); its model-in-body variant logs `model: "run"`, so it is not used for them. Partner ids (`typesafe/jev`) go to the unified `api.cloudflare.com/.../ai/run` path with the gateway named in `cf-aig-gateway-id` — the only route where the partner key stored under the gateway's BYOK is consulted; the gateway host answers 402 "Insufficient balance" for the same call even with the key stored. That path returns no `cf-aig-*` headers, so a cached replay is indistinguishable from a fresh call and the client sends `cf-aig-skip-cache: true` there rather than bill the same answer twice.
+  - **New adapter `extract_workers_ai_native`** for the two usage vocabularies the one endpoint returns: OpenAI-shaped `prompt_tokens`/`completion_tokens` plus Cloudflare's `neurons` (kept in `extras`, never a metric) for catalog models, and `input_tokens`/`output_tokens` one level deeper, under `result.result`, for partner models. Eight captured fixtures, all successful responses, including a real gateway cache MISS/HIT pair and Jev both through Cloudflare and straight from TypeSafe's API.
+  - **Partner models are priced from AI Gateway's own cost table.** `GET .../ai-gateway/costs` is the rate the gateway itself bills by — 2,839 rows across every provider it fronts, `token_pricing` in USD per million — and it lists the partner models the Workers AI catalog omits. A bare `vendor/model` id that misses the catalog is fetched from it one row at a time on the queue's next tick, so the first call to such a model in a process bills tokens and reports the miss, and every later call bills dollars; an id the gateway does not price is remembered as a miss for the TTL rather than re-queried every tick. `warm_pricing(workers_ai_models=[...])` fetches named partner rows up front so even that first call prices, and a fetched row keeps serving past its TTL while it refreshes. Verified on `typesafe/jev`: 446 input tokens at $0.042 per million is 0.000018732 USD, the `cost` Cloudflare stamped on the live log entry, now a `money_golden.json` row in both repos. The row's `cost_in`/`cost_out` per-token fields are ignored — they are 0 on rows whose `token_pricing` is not.
+  - **The requested id is the billing key for Workers AI, not the served name.** Cloudflare's price catalog is keyed by the id you request, and the served name drifts past the catalog's version-strip fallback: `@cf/mistralai/mistral-small-3.1-24b-instruct` answers as `...-24b-v2`, which misses the catalog while the requested id prices. The served name is kept in `extras["served_model"]`.
+  - **Cloudflare gateway log entries now surface `byok` in `extras`.** A row served under BYOK still carries Cloudflare's list-price `cost` (measured: 446 in × $0.042/M on `typesafe/jev`) although Cloudflare charged nothing — the partner bills the customer directly. A backfill that bills `cost` for such a row double-charges; this is the field that lets it not.
+
 ## [0.3.1] - 2026-09-07
 
 ### Added

CONTRIBUTING.md

@@ -82,6 +82,25 @@ uv lock --upgrade-package X  # bump a single package
 5. Update `sdk.py::wrap()` to dispatch to the new wrapper.
 6. Add unit tests against the captured fixtures.
 
+## Adding a first-party client
+
+Some models have no client library to wrap — Cloudflare Workers AI is the case that forced this:
+the official `cloudflare` package cannot address any Workers AI model, and partner models such as
+`typesafe/jev` are not chat-shaped, so the OpenAI-compatible route cannot carry them either. For
+these the SDK ships its own minimal client (`src/lago_agent_sdk/workers_ai.py`, built by
+`LagoSDK.workers_ai()`). The rules are the wrappers' rules, applied to code we own:
+
+1. Capture real fixtures exactly as for a provider — successful responses only, one file per
+   shape, with a capture script alongside them.
+2. Keep the adapter a pure function in `adapters/`; the client only does HTTP and calls `emit()`.
+3. Raise the provider's own error before any instrumentation, so a failed call bills nothing and
+   reaches the caller unchanged. Wrap everything after the response in the same catch-all the
+   wrappers use — instrumentation never breaks the customer's call.
+4. Write down, in the module docstring, every routing decision that was *measured* rather than
+   chosen (which route logs the model, which header the gateway honours, what a cache hit looks
+   like). A first-party client has no upstream to blame for its choices.
+5. Mirror it in the JS repo in the same PR pair, tests and fixtures included.
+
 ## Adding a gateway
 
 `gateway/` is a **second front door** into the same kernel, separate from the provider-native

README.md

@@ -148,6 +148,14 @@ client = sdk.wrap(Anthropic(
 ))
 ```
 
+Models with no client to wrap — every Workers AI model, including partner models like `typesafe/jev` that are not chat-shaped — go through the SDK's own one-method client. Store the partner's key under the gateway's **Provider Keys** (BYOK) and name the gateway; Cloudflare then bills nothing and the partner bills you:
+
+```python
+ai = sdk.workers_ai(account_id, cf_api_token, gateway_id=gateway_id, gateway_auth=gateway_auth)
+out = ai.run("typesafe/jev", {"state": ticket_text, "questions": {...}}, extra_lago={"subscription": "sub_acme"})
+out = ai.run("@cf/meta/llama-3.2-3b-instruct", {"messages": [{"role": "user", "content": "Hello"}]})
+```
+
 Full guide, including backfill from the gateway's Logs API: [docs/cloudflare.md](docs/cloudflare.md).
 
 ### Databricks AI Gateway

docs/cloudflare.md

@@ -29,6 +29,12 @@ from lago_agent_sdk.gateway.adapters import extract_cloudflare_log, resolve_subs
 for entry in fetch_gateway_logs():  # GET .../ai-gateway/gateways/{id}/logs
     usage = extract_cloudflare_log(entry)
     sub = resolve_subscription(entry) or "sub_default"  # from the call's cf-aig-metadata, if set
+    if usage.extras.get("byok"):
+        # Served with the customer's own provider key (BYOK): Cloudflare charged nothing and the
+        # partner bills them directly, yet `cost` still carries Cloudflare's list price. Bill the
+        # tokens, never that number.
+        sdk.emit(usage, subscription=sub, mode="tokens", event_id=f"cf_{entry['id']}")
+        continue
     sdk.emit(usage, subscription=sub, mode="price", usd_cost=entry.get("cost") or 0, event_id=f"cf_{entry['id']}")
 sdk.flush()
 ```
@@ -38,3 +44,43 @@ This page is the complete picture; runnable notebooks are kept out of the repo (
 **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.
+
+## Workers AI models with no client to wrap
+
+Cloudflare-hosted models can be reached through the OpenAI-compatible `/compat` endpoint above, but only when they are chat-shaped. Partner models are not: `typesafe/jev` takes `{state, questions}` and refuses a `messages` array, and the official `cloudflare` package cannot address any Workers AI model at all (it percent-encodes the slash in the model name). For these the SDK ships its own one-method client:
+
+```python
+ai = sdk.workers_ai(
+    account_id,
+    cf_api_token,                       # a Cloudflare API token with Workers AI access
+    gateway_id=gateway_id,              # optional — see what it buys below
+    gateway_auth=gateway_auth,          # the gateway's cf-aig-authorization token, if authentication is on
+    subscription="sub_acme",            # default for every call; extra_lago={"subscription": ...} per call
+)
+
+# a partner model: the partner's key must be stored under the gateway's Provider Keys (BYOK)
+out = ai.run(
+    "typesafe/jev",
+    {
+        "state": "I was charged twice and need the duplicate refunded before Friday.",
+        "questions": {
+            "department": {"type": "choice", "instructions": "Which team should handle this?",
+                           "criteria": {"billing": "Payments, refunds", "technical": "Bugs, outages"}},
+        },
+    },
+    extra_lago={"dimensions": {"ticket": "T-4821"}},
+)
+out["result"]["result"]["answers"]["department"]["choice"]   # "billing"
+
+# a catalog model: same client, same billing
+ai.run("@cf/meta/llama-3.2-3b-instruct", {"messages": [{"role": "user", "content": "Hello"}], "max_tokens": 50})
+sdk.flush()
+```
+
+`run()` returns Cloudflare's full response envelope unchanged and raises `WorkersAIError` on an error status (a 402 for a partner model whose key is not stored, a 403 for a model not on your Workers plan), before anything is billed.
+
+**What the gateway buys.** With `gateway_id` set, `@cf/...` calls go through the gateway host: a cache hit (`cf-aig-cache-status: HIT`) is not billed, and every event carries a `cf_log_id` dimension that matches the entry's `id` in the Logs API. Partner models go through the unified `api.cloudflare.com/.../ai/run` path with the gateway named in a header, which is the only route where the gateway's stored partner key is consulted; that path returns no cache header, so the client asks the gateway to skip its cache for those calls rather than risk billing a cached replay. Pass `extra_headers={"cf-aig-skip-cache": "false"}` to opt back in.
+
+**Billing.** Token events from the response's own `usage`: `prompt_tokens`/`completion_tokens` for catalog models, `input_tokens`/`output_tokens` for partner models. Catalog models price from Cloudflare's published Workers AI rates in price mode as usual. Partner models are not in that catalog; their rates come from AI Gateway's own cost table (`GET .../ai-gateway/costs`), fetched per model in the background the first time one is seen. Left alone, the very first call to a partner model in a process bills tokens and reports the miss via `on_error`, and every call after it bills dollars; name the ids up front with `sdk.warm_pricing(["workers-ai"], workers_ai_models=["typesafe/jev"])` and even the first call prices. A fetched rate keeps serving past its TTL while it refreshes, so an expiry never bills a call as tokens. Jev lists input at $0.042 per million with free output, the same rate the gateway stamps as `cost` on its log entries. Under BYOK the gateway's Logs API still reports a `cost` for the partner call at Cloudflare's list price although Cloudflare charged nothing — the entry's `byok` field (surfaced in `extras["byok"]` by `extract_cloudflare_log`) is what tells a backfill not to bill it. The id billed is the one you requested, because that is what Cloudflare's price catalog is keyed by; the name the model reports (`jev-1.13.0`, `...-24b-v2`) is kept in `extras["served_model"]`.
+
+Streaming is not supported by this client; use the `/compat` endpoint through a wrapped OpenAI client for streamed chat.

src/lago_agent_sdk/__init__.py

@@ -17,9 +17,12 @@
     compute_cost,
 )
 from .sdk import LagoSDK
+from .workers_ai import WorkersAI, WorkersAIError
 
 __all__ = [
     "LagoSDK",
+    "WorkersAI",
+    "WorkersAIError",
     "LagoConfig",
     "CanonicalUsage",
     "LagoApiError",

src/lago_agent_sdk/adapters/__init__.py

@@ -4,6 +4,7 @@
 from .gemini_native import extract_gemini_native
 from .mistral_native import extract_mistral_native
 from .openai_native import extract_openai_native
+from .workers_ai_native import extract_workers_ai_native
 
 __all__ = [
     "extract_anthropic_native",
@@ -13,4 +14,5 @@
     "extract_gemini_native",
     "extract_mistral_native",
     "extract_openai_native",
+    "extract_workers_ai_native",
 ]

src/lago_agent_sdk/adapters/workers_ai_native.py

@@ -0,0 +1,123 @@
+"""Workers AI `/ai/run` adapter — maps a run response to CanonicalUsage.
+
+Verified against real captures (fixtures/workers_ai/, 2026-09-21) of the model-in-body
+route `POST /accounts/{id}/ai/run {"model": ..., "input": {...}}`. That route matters
+because it is the only one that reaches every Workers AI model: partner models such as
+`typesafe/jev` have no `@cf/` prefix and the path-style `/ai/run/{model}` answers "No
+route for that URI" for them, as does the gateway's `/compat` endpoint, which requires
+a `messages` array the model rejects.
+
+Two usage vocabularies come back from the one endpoint:
+
+  chat models    result.usage.prompt_tokens / completion_tokens / total_tokens
+                 result.usage.prompt_tokens_details.cached_tokens
+                 result.usage.neurons                               (01-04, 06, 07)
+  typesafe/jev   result.result.usage.input_tokens / output_tokens   (05: one level
+                 deeper — a partner model's answer is wrapped as
+                 `result: {state, result: {model, answers, usage}, gatewayMetadata}`;
+                 08 is the same object straight from TypeSafe's API, unwrapped)
+
+The REQUESTED model id is the one carried, not the served name — the opposite of the
+native adapters' rule, for a measured reason. Cloudflare's price catalog is keyed by the
+id you request, and the served name drifts from it in ways the catalog's version-strip
+fallback does not cover: `@cf/mistralai/mistral-small-3.1-24b-instruct` answers as
+`...-24b-v2` (04) and that name MISSES the catalog while the requested id prices;
+`typesafe/jev` answers as `jev-1.13.0` (05, 08). The served name is kept in
+`extras["served_model"]` when it differs, so nothing is lost — only the billing key stays
+the one Cloudflare itself bills by.
+
+`neurons` is Cloudflare's own billing unit, not a token count: it lands in `extras`
+and is never a metric. `gatewayMetadata` (`keySource: "BYOK"` on 05) says whose key paid
+for the call and lands in `extras` too — under BYOK Cloudflare charged nothing and the
+partner bills the customer directly, which matters to anyone reconciling against the
+Cloudflare dashboard. A reasoning model (02, deepseek-r1-distill) bundles its thinking
+into `completion_tokens` with no separate field, so `reasoning` stays 0 — the same
+shape Magistral has on Mistral. `cached_tokens` is a subset of `prompt_tokens`, which
+is why "workers-ai" sits in `INPUT_INCLUDES_CACHE_READ`.
+
+A failure body (402 no credits / 403 not on plan / 400 no such model — seen live, not
+kept as fixtures) carries `result: {}` and `errors: [...]`. The adapter yields an all-zero
+usage for it rather than raising — it is a pure function and cannot know the HTTP status;
+the client decides what a failure means.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ..canonical import CanonicalUsage
+
+# Every `usage` key this adapter maps or deliberately ignores. Anything else is drift
+# and is swept into `extras["usage"]` — never silently dropped, never miscounted.
+_KNOWN_USAGE_KEYS = frozenset(
+    {
+        # chat models
+        "prompt_tokens",
+        "completion_tokens",
+        "total_tokens",
+        "prompt_tokens_details",
+        # typesafe/jev
+        "input_tokens",
+        "output_tokens",
+        # Cloudflare's billing unit — kept in extras, not a metric
+        "neurons",
+    }
+)
+_KNOWN_DETAIL_KEYS = frozenset({"cached_tokens"})
+
+
+def _safe_dict(v: Any) -> dict[str, Any]:
+    return v if isinstance(v, dict) else {}
+
+
+def _safe_int(v: Any) -> int:
+    try:
+        return max(0, int(v or 0))
+    except (TypeError, ValueError):
+        return 0
+
+
+def extract_workers_ai_native(response: Any, model_id: str = "") -> CanonicalUsage:
+    """Translate a Workers AI `/ai/run` response body → CanonicalUsage.
+
+    Accepts the full envelope (`{"result": {...}, "success": true, ...}`) or the bare
+    `result` object. `model_id` is the model the caller requested and is the id carried
+    (see the module docstring); the served name, when it differs, lands in
+    `extras["served_model"]`.
+    """
+    payload = _safe_dict(response)
+    result = _safe_dict(payload.get("result")) or payload
+    gateway_meta = _safe_dict(result.get("gatewayMetadata"))
+    inner = _safe_dict(result.get("result"))
+    if inner and ("usage" in inner or "answers" in inner):
+        # Partner-model envelope (05): the model's own object sits one level down.
+        result = inner
+    # The envelope may also carry `usage` at the top level; check both so a shape change
+    # moves nothing to zero.
+    usage = _safe_dict(result.get("usage")) or _safe_dict(payload.get("usage"))
+    details = _safe_dict(usage.get("prompt_tokens_details"))
+
+    extras: dict[str, Any] = {}
+    if gateway_meta:
+        extras["gateway_metadata"] = gateway_meta
+    served = result.get("model")
+    if isinstance(served, str) and served and served != model_id:
+        extras["served_model"] = served
+    if "neurons" in usage:
+        extras["neurons"] = usage["neurons"]
+    drift = {k: v for k, v in usage.items() if k not in _KNOWN_USAGE_KEYS}
+    detail_drift = {k: v for k, v in details.items() if k not in _KNOWN_DETAIL_KEYS}
+    if detail_drift:
+        drift["prompt_tokens_details"] = detail_drift
+    if drift:
+        extras["usage"] = drift
+
+    return CanonicalUsage(
+        input=_safe_int(usage.get("prompt_tokens")) or _safe_int(usage.get("input_tokens")),
+        output=_safe_int(usage.get("completion_tokens")) or _safe_int(usage.get("output_tokens")),
+        cache_read=_safe_int(details.get("cached_tokens")),
+        model=model_id or (served if isinstance(served, str) else ""),
+        provider="workers-ai",
+        api="workers_ai_run",
+        extras=extras,
+    )

src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py

@@ -220,6 +220,12 @@ def extract_cloudflare_log(entry: dict[str, Any]) -> CanonicalUsage:
             "cached": entry.get("cached"),
             "step": entry.get("step"),
             "log_id": entry.get("id"),
+            # Which key paid: None (Cloudflare credits / the customer's own header key) or the
+            # BYOK alias (`"default"`). Measured 2026-09-21 on a `typesafe/jev` row served under
+            # BYOK: Cloudflare still fills `cost` with its list price (446 in × $0.042/M) although
+            # it charged nothing — the partner bills the customer directly. A backfill that bills
+            # `cost` for such a row double-charges; this is the field that lets it not.
+            "byok": entry.get("byok"),
             # Drift sweep — the same contract `adapters/openai_native.py` enforces, and
             # for the same reason: a counter this adapter does not map must not vanish
             # without an error or an on_error. `extras` used to be exactly the three