tokenmeter-sdk 0.2.0__py3-none-any.whl

tokenmeter/_wrap.py

@@ -0,0 +1,807 @@
+"""tokenmeter.wrap(client) -- OpenAI (chat/completions + responses) and Anthropic
+(messages), streaming and non-streaming, sync and async.
+
+Design contract (non-negotiable, tested):
+  - the SDK never throws into host code: only OUR bookkeeping is try/except'd;
+    the real provider call and its real exceptions pass through untouched.
+  - the SDK never adds latency: no network I/O on this path, ever -- capture is
+    post-hoc from the response object, emit() is an in-memory append.
+
+Detection is duck-typed (no `import openai` / `import anthropic` at module load
+time) so wrap() works against whichever SDK the host app already has installed,
+without tokenmeter declaring a hard dependency on either.
+"""
+from __future__ import annotations
+
+import inspect
+import logging
+import time
+from datetime import datetime, timezone
+from typing import Any, Callable, Optional
+
+import json as _json
+
+from ._client import get_state
+from ._preflight import active_preflight_id
+from ._context import current
+from ._ids import uuid7_str
+from ._optimization import optimization_tags, enrich_tags
+from ._normalize import (
+    normalize_anthropic_usage,
+    normalize_azure_openai_usage,
+    normalize_bedrock_converse_usage,
+    normalize_bedrock_invoke_model,
+    normalize_gemini_usage,
+    normalize_openai_usage,
+)
+
+logger = logging.getLogger("tokenmeter")
+
+_ZERO_USAGE = {
+    "input_tokens": 0,
+    "output_tokens": 0,
+    "cache_read_tokens": 0,
+    "cache_write_tokens": 0,
+    "reasoning_tokens": 0,
+}
+_NULL_USAGE = {
+    "input_tokens": None,
+    "output_tokens": None,
+    "cache_read_tokens": None,
+    "cache_write_tokens": None,
+    "reasoning_tokens": None,
+}
+
+
+def _to_dict(obj: Any) -> Optional[dict]:
+    if obj is None:
+        return None
+    if isinstance(obj, dict):
+        return obj
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump()
+    if hasattr(obj, "dict"):
+        return obj.dict()
+    try:
+        return {k: v for k, v in vars(obj).items() if not k.startswith("_")}
+    except TypeError:
+        return None
+
+
+def _safe(fn: Callable, *a, **kw) -> None:
+    try:
+        fn(*a, **kw)
+    except Exception:
+        logger.debug("tokenmeter: instrumentation error (ignored)", exc_info=True)
+
+
+def _emit(event: dict) -> None:
+    state = get_state()
+    if state is None:
+        return
+    state.emit(event)
+
+
+def _snapshot_optimization(params: dict, ctx) -> dict:
+    try:
+        state = get_state()
+        if state is None:
+            return {}
+        options = {**state.optimization, **(ctx.optimization or {} if ctx else {})}
+        return optimization_tags(state.transport.ingest_key, params, options)
+    except Exception:
+        return {}
+
+
+def _build_event(*, provider: str, model_raw: Optional[str], kwargs: dict, latency_ms: int,
+                  usage_complete: bool, norm: dict, error_code: Optional[str] = None,
+                  ctx: Any = "_unset", model_map: Optional[dict] = None, otel: Any = "_unset",
+                  optimization: Optional[dict] = None) -> dict:
+    # ctx is captured by the CALLER at the moment the provider call was made
+    # (see _wrap_method) and threaded through, rather than re-read here via
+    # current(). Streaming responses are generators: current() read lazily
+    # inside a generator's `finally` block reflects whatever context is
+    # active at CONSUMPTION time, not at creation time (generators do not
+    # snapshot contextvars the way asyncio.Task does) -- a stream created
+    # inside action() and handed off to be iterated later (e.g. a framework's
+    # StreamingResponse) would otherwise silently lose all attribution. The
+    # non-streaming path still defaults to current() since it emits
+    # synchronously, in the same context the call was made in.
+    if ctx == "_unset":
+        ctx = current()
+    if provider == "bedrock" and kwargs.get("model"):
+        # Anthropic's Bedrock client: the request's `model` is the Bedrock
+        # modelId / inference profile that Cost Explorer bills; the response
+        # may carry Anthropic's own model name instead.
+        model_raw = kwargs["model"]
+    model_raw = model_raw or kwargs.get("model") or kwargs.get("modelId")
+    if model_map:
+        # Azure returns the deployment name, not the canonical model id, on
+        # some APIs (spike §2.2) -- resolve it here so §7.1's model_id
+        # aliasing/pricing catalog lookup still works. Unmapped deployments
+        # pass through unresolved and land in the unpriced quarantine, which
+        # already exists for exactly this case.
+        model_raw = model_map.get(model_raw, model_raw)
+    event = {
+        "event_id": uuid7_str(),
+        # Required by ingest's validateWireEvent (missing => quarantined as
+        # invalid_occurred_at). ISO 8601 UTC at capture time, same as the
+        # TypeScript SDK's capture.ts.
+        "occurred_at": datetime.now(timezone.utc).isoformat(),
+        "provider": provider,
+        "model_raw": model_raw,
+        "input_tokens": norm.get("input_tokens"),
+        "output_tokens": norm.get("output_tokens"),
+        "cache_read_tokens": norm.get("cache_read_tokens"),
+        "cache_write_tokens": norm.get("cache_write_tokens"),
+        "reasoning_tokens": norm.get("reasoning_tokens"),
+        "is_batch": False,
+        "usage_complete": usage_complete,
+        "latency_ms": latency_ms,
+        "error_code": error_code,
+        "trace_id": ctx.trace_id if ctx else None,
+        "parent_span_id": ctx.parent_span_id if ctx else None,
+        "action": ctx.action if ctx else None,
+        "customer_ref": ctx.customer_ref if ctx else None,
+        "feature": ctx.feature if ctx else None,
+        "team": ctx.team if ctx else None,
+        "environment": ctx.environment if ctx else None,
+        "project": ctx.project if ctx else None,
+        "member_ref": ctx.member_ref if ctx else None,
+        "provenance": "sdk",
+        # Streams pass the link snapshotted when the call was made (like ctx
+        # above); a non-streaming call builds its event in the caller's own
+        # context, where reading the active span now gives the same answer.
+        **(_otel_link_tags() if otel == "_unset" else otel),
+    }
+    # governed() scope: stamp the preflight decision so ingest's coverage
+    # report counts this call as preflight-protected (_preflight.py).
+    preflight_id = active_preflight_id()
+    if preflight_id:
+        event["tags"] = {**event.get("tags", {}), "preflight_id": preflight_id}
+    if usage_complete and error_code is None:
+        extra = optimization if optimization is not None else _snapshot_optimization({}, ctx)
+        if extra:
+            event["tags"] = enrich_tags(event.get("tags", {}), extra)
+    return event
+
+
+def _otel_link_tags() -> dict:
+    """The active OpenTelemetry span as tags.otel_trace_id / otel_span_id.
+
+    When an app runs both this SDK and an OTel GenAI instrumentation, the same
+    model call reaches ingest twice; ingest uses this link to keep only one.
+    opentelemetry-api is optional -- absent, or no valid span, adds nothing.
+    Read when the event is built: exact for non-streaming calls; for a stream
+    consumed in another context it may be missing (never wrong: ingest also
+    requires matching output tokens before treating two events as one call).
+    """
+    try:
+        from opentelemetry import trace  # optional dependency
+        sc = trace.get_current_span().get_span_context()
+        if not sc.is_valid:
+            return {}
+        return {"tags": {"otel_trace_id": format(sc.trace_id, "032x"), "otel_span_id": format(sc.span_id, "016x")}}
+    except Exception:
+        return {}
+
+
+_NORMALIZERS = {
+    "openai": normalize_openai_usage,
+    "azure_openai": normalize_azure_openai_usage,
+    "anthropic": normalize_anthropic_usage,
+    "google": normalize_gemini_usage,
+    "bedrock": normalize_bedrock_converse_usage,
+}
+
+
+def _normalize(provider: str, usage_dict: Optional[dict]) -> dict:
+    if usage_dict is None:
+        return dict(_NULL_USAGE)
+    return _NORMALIZERS[provider](usage_dict)
+
+
+# --------------------------------------------------------------------------
+# Non-streaming capture
+# --------------------------------------------------------------------------
+
+def _emit_success(provider: str, kwargs: dict, response: Any, t0: float, model_map: Optional[dict] = None,
+                  wire: Optional[str] = None, optimization: Optional[dict] = None, ctx: Any = "_unset") -> None:
+    latency_ms = int((time.monotonic() - t0) * 1000)
+    model_raw = getattr(response, "model", None)
+    usage_dict = _to_dict(getattr(response, "usage", None))
+    norm = _normalize(provider, usage_dict)
+    _emit(_build_event(provider=wire or provider, model_raw=model_raw, kwargs=kwargs, latency_ms=latency_ms,
+                        usage_complete=usage_dict is not None, norm=norm, model_map=model_map, optimization=optimization, ctx=ctx))
+
+
+def _emit_error(provider: str, kwargs: dict, t0: float, exc: BaseException, model_key: str = "model",
+                wire: Optional[str] = None) -> None:
+    latency_ms = int((time.monotonic() - t0) * 1000)
+    error_code = str(getattr(exc, "status_code", None) or type(exc).__name__)
+    _emit(_build_event(provider=wire or provider, model_raw=kwargs.get(model_key), kwargs=kwargs, latency_ms=latency_ms,
+                        usage_complete=True, norm=dict(_ZERO_USAGE), error_code=error_code))
+
+
+# --------------------------------------------------------------------------
+# Streaming capture: usage lives in a terminal chunk/event, not every chunk.
+# If the iterator raises or ends without ever seeing a usage frame, the event
+# is recorded usage_complete=False with NULL tokens -- never guessed from
+# partial deltas.
+# --------------------------------------------------------------------------
+
+_ExtractUsage = Callable[[Any], Optional[dict]]
+
+_STREAM_EXTRACTORS: dict = {
+    "openai_chat": lambda chunk: _to_dict(getattr(chunk, "usage", None)),
+    "openai_responses": lambda ev: (
+        _to_dict(getattr(ev, "usage", None))
+        or _to_dict(getattr(getattr(ev, "response", None), "usage", None))
+    ),
+    "azure_openai_chat": lambda chunk: _to_dict(getattr(chunk, "usage", None)),
+    "azure_openai_responses": lambda ev: (
+        _to_dict(getattr(ev, "usage", None))
+        or _to_dict(getattr(getattr(ev, "response", None), "usage", None))
+    ),
+    "anthropic_messages": lambda ev: (
+        _to_dict(getattr(ev, "usage", None))
+        or _to_dict(getattr(getattr(ev, "message", None), "usage", None))
+    ),
+}
+
+
+def _stream_model_of(chunk: Any) -> Optional[str]:
+    return getattr(chunk, "model", None)
+
+
+def _merge_usage(acc: Optional[dict], new: Optional[dict]) -> Optional[dict]:
+    """Per-field last-non-null-wins merge, not a wholesale overwrite.
+
+    Anthropic's real stream sends input/cache tokens on `message_start` and
+    the FINAL output_tokens on a later `message_delta` whose usage dict does
+    not repeat input/cache fields at all. Overwriting `acc` wholesale with
+    each new non-null usage dict (the earlier bug here) discards
+    message_start's input/cache counts the moment message_delta arrives,
+    which silently zero-costs every Anthropic streaming call's input tokens.
+    Merging field-by-field keeps every counter any frame reported.
+    """
+    if new is None:
+        return acc
+    if acc is None:
+        return dict(new)
+    merged = dict(acc)
+    for k, v in new.items():
+        if v is not None:
+            merged[k] = v
+    return merged
+
+
+def _wrap_sync_stream(orig_iter, provider: str, kind: str, kwargs: dict, t0: float, ctx: Any,
+                       model_map: Optional[dict] = None, wire: Optional[str] = None, otel: Any = "_unset", optimization: Optional[dict] = None):
+    # ctx is snapshotted by the caller at stream-creation time -- see the
+    # _build_event docstring note on why it can't be read lazily here.
+    extract = _STREAM_EXTRACTORS[f"{provider}_{kind}"]
+    final_usage: Optional[dict] = None
+    final_model: Optional[str] = None
+    errored = False
+    try:
+        for chunk in orig_iter:
+            m = _stream_model_of(chunk)
+            if m:
+                final_model = m
+            u = extract(chunk)
+            final_usage = _merge_usage(final_usage, u)
+            yield chunk
+    except BaseException:
+        errored = True
+        raise
+    finally:
+        usage_complete = (final_usage is not None) and not errored
+        norm = _normalize(provider, final_usage) if usage_complete else dict(_NULL_USAGE)
+        latency_ms = int((time.monotonic() - t0) * 1000)
+        _safe(_emit, _build_event(provider=wire or provider, model_raw=final_model, kwargs=kwargs,
+                                   latency_ms=latency_ms, usage_complete=usage_complete, norm=norm,
+                                   ctx=ctx, otel=otel, model_map=model_map, optimization=optimization))
+
+
+async def _wrap_async_stream(orig_aiter, provider: str, kind: str, kwargs: dict, t0: float, ctx: Any,
+                              model_map: Optional[dict] = None, wire: Optional[str] = None, otel: Any = "_unset", optimization: Optional[dict] = None):
+    # ctx is snapshotted by the caller at stream-creation time -- same reason
+    # as _wrap_sync_stream: an async generator does not snapshot contextvars
+    # at creation either, only asyncio.Task does.
+    extract = _STREAM_EXTRACTORS[f"{provider}_{kind}"]
+    final_usage: Optional[dict] = None
+    final_model: Optional[str] = None
+    errored = False
+    try:
+        async for chunk in orig_aiter:
+            m = _stream_model_of(chunk)
+            if m:
+                final_model = m
+            u = extract(chunk)
+            final_usage = _merge_usage(final_usage, u)
+            yield chunk
+    except BaseException:
+        errored = True
+        raise
+    finally:
+        usage_complete = (final_usage is not None) and not errored
+        norm = _normalize(provider, final_usage) if usage_complete else dict(_NULL_USAGE)
+        latency_ms = int((time.monotonic() - t0) * 1000)
+        _safe(_emit, _build_event(provider=wire or provider, model_raw=final_model, kwargs=kwargs,
+                                   latency_ms=latency_ms, usage_complete=usage_complete, norm=norm,
+                                   ctx=ctx, otel=otel, model_map=model_map, optimization=optimization))
+
+
+# --------------------------------------------------------------------------
+# Method wrapping
+# --------------------------------------------------------------------------
+
+def _wrap_method(obj: Any, attr_name: str, provider: str, kind: str, model_map: Optional[dict] = None,
+                 wire: Optional[str] = None) -> None:
+    # `provider` selects the usage shape (normalizer + stream extractor);
+    # `wire`, when set, is the provider_kind the event is recorded under.
+    # They differ for Anthropic's Bedrock/Vertex clients: Anthropic-shaped
+    # usage, but billed -- and reconciled -- as bedrock / google.
+    if obj is None or not hasattr(obj, attr_name):
+        return
+    original = getattr(obj, attr_name)
+    if getattr(original, "_tokenmeter_wrapped", False):
+        return
+
+    is_async = inspect.iscoroutinefunction(original)
+    # Azure's Chat Completions shape is the openai-python SDK verbatim, same
+    # stream_options requirement/restriction as plain OpenAI (spike §2.1).
+    injects_stream_options = provider in ("openai", "azure_openai") and kind == "chat"
+
+    if is_async:
+        async def wrapper(*args, **kwargs):
+            ctx = current()
+            optimization = _snapshot_optimization(kwargs, ctx)
+            is_stream = bool(kwargs.get("stream"))
+            # Chat Completions needs stream_options to get a terminal usage
+            # chunk; the Responses API always sends usage on its final event
+            # and has no such parameter -- injecting it there is a TypeError
+            # (openai-python declares explicit named params, not **kwargs
+            # passthrough), which would turn a working streaming call into a
+            # hard failure the instant wrap() is applied.
+            if is_stream and injects_stream_options:
+                kwargs.setdefault("stream_options", {"include_usage": True})
+            t0 = time.monotonic()
+            try:
+                result = await original(*args, **kwargs)
+            except Exception as exc:
+                _safe(_emit_error, provider, kwargs, t0, exc, wire=wire)
+                raise
+            if is_stream:
+                # Snapshot context NOW, at call time -- the stream object
+                # returned below may be iterated far later/elsewhere, after
+                # any action() scope that was open here has already exited.
+                try:
+                    return _wrap_async_stream(result, provider, kind, kwargs, t0, ctx, model_map, wire, _otel_link_tags(), optimization)
+                except Exception:
+                    logger.debug("tokenmeter: failed to wrap async stream", exc_info=True)
+                    return result
+            _safe(_emit_success, provider, kwargs, result, t0, model_map, wire, optimization, ctx)
+            return result
+
+        wrapper._tokenmeter_wrapped = True  # type: ignore[attr-defined]
+        setattr(obj, attr_name, wrapper)
+    else:
+        def wrapper(*args, **kwargs):
+            ctx = current()
+            optimization = _snapshot_optimization(kwargs, ctx)
+            is_stream = bool(kwargs.get("stream"))
+            # See the async wrapper above: Responses API has no stream_options
+            # parameter and openai-python rejects unknown kwargs with a
+            # TypeError, so this must stay scoped to chat/completions.
+            if is_stream and injects_stream_options:
+                kwargs.setdefault("stream_options", {"include_usage": True})
+            t0 = time.monotonic()
+            try:
+                result = original(*args, **kwargs)
+            except Exception as exc:
+                _safe(_emit_error, provider, kwargs, t0, exc, wire=wire)
+                raise
+            if is_stream:
+                # Snapshot context NOW, at call time -- the stream object
+                # returned below may be iterated far later/elsewhere, after
+                # any action() scope that was open here has already exited.
+                try:
+                    return _wrap_sync_stream(result, provider, kind, kwargs, t0, ctx, model_map, wire, _otel_link_tags(), optimization)
+                except Exception:
+                    logger.debug("tokenmeter: failed to wrap stream", exc_info=True)
+                    return result
+            _safe(_emit_success, provider, kwargs, result, t0, model_map, wire, optimization, ctx)
+            return result
+
+        wrapper._tokenmeter_wrapped = True  # type: ignore[attr-defined]
+        setattr(obj, attr_name, wrapper)
+
+
+def _wrap_openai_or_azure(client: Any, model_map: Optional[dict]) -> None:
+    # Azure's `AzureOpenAI`/`AsyncAzureOpenAI` clients ARE the openai-python
+    # SDK -- same `.chat.completions.create` / `.responses.create` shape, so
+    # detection is by client identity, not structure (spike §2.1). We check
+    # both the documented attribute and the class name since neither alone
+    # is guaranteed across openai-python versions, and this must stay a duck
+    # type (no `import openai` at module load).
+    is_azure = hasattr(client, "azure_endpoint") or type(client).__name__ in ("AzureOpenAI", "AsyncAzureOpenAI")
+    provider = "azure_openai" if is_azure else "openai"
+    azure_map = model_map if is_azure else None
+
+    chat = getattr(client, "chat", None)
+    completions = getattr(chat, "completions", None) if chat is not None else None
+    if completions is not None and hasattr(completions, "create"):
+        _wrap_method(completions, "create", provider, "chat", azure_map)
+
+    responses = getattr(client, "responses", None)
+    if responses is not None and hasattr(responses, "create"):
+        _wrap_method(responses, "create", provider, "responses", azure_map)
+
+
+_OPENAI_COMPATIBLE_PROVIDERS = frozenset({"xai", "mistral", "fireworks", "together", "deepinfra", "openrouter", "vercel_gateway"})
+
+
+def wrap_openai_compatible(client: Any, provider: str) -> Any:
+    """Wrap an OpenAI-compatible client under its actual billing provider.
+
+    Configure the upstream client with that provider's base URL, then call
+    this instead of :func:`wrap`. The OpenAI usage normalizer is deliberately
+    reused because these endpoints expose the Chat Completions / Responses
+    usage schema. Calls without provider usage remain incomplete, never
+    estimated from request text.
+    """
+    if provider not in _OPENAI_COMPATIBLE_PROVIDERS:
+        supported = ", ".join(sorted(_OPENAI_COMPATIBLE_PROVIDERS))
+        raise ValueError(f"provider must be one of: {supported}")
+    chat = getattr(client, "chat", None)
+    completions = getattr(chat, "completions", None) if chat is not None else None
+    if completions is not None and hasattr(completions, "create"):
+        _wrap_method(completions, "create", "openai", "chat", wire=provider)
+    responses = getattr(client, "responses", None)
+    if responses is not None and hasattr(responses, "create"):
+        _wrap_method(responses, "create", "openai", "responses", wire=provider)
+    return client
+
+
+def _wrap_anthropic(client: Any) -> None:
+    # Anthropic-shaped: has .messages.create, no .chat (distinguishes from
+    # both OpenAI and Azure OpenAI, which always have .chat).
+    chat = getattr(client, "chat", None)
+    messages = getattr(client, "messages", None)
+    if messages is not None and hasattr(messages, "create") and chat is None:
+        _wrap_method(messages, "create", "anthropic", "messages", wire=_anthropic_wire_provider(client))
+
+
+def _anthropic_wire_provider(client: Any) -> Optional[str]:
+    # anthropic-python's AnthropicBedrock / AnthropicVertex (and their Async
+    # variants) share the Messages API shape but are billed by AWS / Google,
+    # so their usage must land under bedrock / google or it can never
+    # reconcile against that provider's bill.
+    name = type(client).__name__
+    if "Bedrock" in name:
+        return "bedrock"
+    if "Vertex" in name:
+        return "google"
+    return None
+
+
+# --------------------------------------------------------------------------
+# Gemini (google-genai): client.models.generate_content / .generate_content_stream.
+# Streaming is a SEPARATE method, not a `stream=True` kwarg, so it gets its
+# own wrapper rather than reusing _wrap_method's is_stream branching.
+# ponytail: sync client only (google-genai's async client lives at
+# `client.aio.models`, same methods) -- add if a host needs it.
+# --------------------------------------------------------------------------
+
+def _emit_gemini_success(kwargs: dict, response: Any, t0: float) -> None:
+    latency_ms = int((time.monotonic() - t0) * 1000)
+    model_raw = getattr(response, "model_version", None) or kwargs.get("model")
+    usage_dict = _to_dict(getattr(response, "usage_metadata", None))
+    norm = normalize_gemini_usage(usage_dict)
+    _emit(_build_event(provider="google", model_raw=model_raw, kwargs=kwargs, latency_ms=latency_ms,
+                        usage_complete=usage_dict is not None, norm=norm))
+
+
+def _wrap_gemini_stream(orig_iter, kwargs: dict, t0: float, ctx: Any, otel: Any = "_unset"):
+    final_usage: Optional[dict] = None
+    final_model: Optional[str] = None
+    errored = False
+    try:
+        for chunk in orig_iter:
+            m = getattr(chunk, "model_version", None)
+            if m:
+                final_model = m
+            u = _to_dict(getattr(chunk, "usage_metadata", None))
+            # Community-observed (spike §3.2, unverified by live call): each
+            # chunk's usage_metadata is a complete cumulative snapshot, not a
+            # delta -- the LAST chunk is the correct one. Field-merge (rather
+            # than wholesale overwrite) is still the safe choice here since
+            # it degrades to "last non-null wins" for any chunk that omits a
+            # field, without a live probe to confirm cumulative-vs-partial.
+            final_usage = _merge_usage(final_usage, u)
+            yield chunk
+    except BaseException:
+        errored = True
+        raise
+    finally:
+        usage_complete = (final_usage is not None) and not errored
+        norm = normalize_gemini_usage(final_usage) if usage_complete else dict(_NULL_USAGE)
+        latency_ms = int((time.monotonic() - t0) * 1000)
+        _safe(_emit, _build_event(provider="google", model_raw=final_model, kwargs=kwargs,
+                                   latency_ms=latency_ms, usage_complete=usage_complete, norm=norm, ctx=ctx, otel=otel))
+
+
+def _wrap_gemini_method(obj: Any, attr_name: str, streaming: bool) -> None:
+    if obj is None or not hasattr(obj, attr_name):
+        return
+    original = getattr(obj, attr_name)
+    if getattr(original, "_tokenmeter_wrapped", False):
+        return
+
+    def wrapper(*args, **kwargs):
+        t0 = time.monotonic()
+        try:
+            result = original(*args, **kwargs)
+        except Exception as exc:
+            _safe(_emit_error, "google", kwargs, t0, exc)
+            raise
+        if streaming:
+            ctx = current()
+            try:
+                return _wrap_gemini_stream(result, kwargs, t0, ctx, _otel_link_tags())
+            except Exception:
+                logger.debug("tokenmeter: failed to wrap gemini stream", exc_info=True)
+                return result
+        _safe(_emit_gemini_success, kwargs, result, t0)
+        return result
+
+    wrapper._tokenmeter_wrapped = True  # type: ignore[attr-defined]
+    setattr(obj, attr_name, wrapper)
+
+
+def _wrap_gemini(client: Any) -> None:
+    models = getattr(client, "models", None)
+    if models is None or not (hasattr(models, "generate_content") or hasattr(models, "generate_content_stream")):
+        return
+    if hasattr(models, "generate_content"):
+        _wrap_gemini_method(models, "generate_content", streaming=False)
+    if hasattr(models, "generate_content_stream"):
+        _wrap_gemini_method(models, "generate_content_stream", streaming=True)
+
+
+# --------------------------------------------------------------------------
+# AWS Bedrock (boto3 `bedrock-runtime`). Responses are plain dicts already
+# (botocore, not a provider SDK object model), so no _to_dict() needed.
+# Converse/ConverseStream is the primary, normalized path (spike §1.1);
+# InvokeModel is per-family (Anthropic + Nova bodies, header fallback for
+# the rest -- spike §1.2). boto3 methods live directly on the client, not
+# nested, so these get their own thin wrapper rather than _wrap_method.
+# ponytail: sync boto3 client only -- aioboto3 is a separate package with a
+# different call shape; add if a host needs it.
+# --------------------------------------------------------------------------
+
+class _ReplayableBody:
+    """Minimal stand-in for botocore's StreamingBody so a host that calls
+    `response['body'].read()` after us still gets the real bytes -- reading
+    the body once (in-memory, no extra network hop) to extract usage would
+    otherwise leave the host with an already-exhausted stream.
+    ponytail: only `.read()`/`.close()` -- add `.iter_lines()` etc. if a host
+    needs them; nothing in this codebase's own tests does.
+    """
+
+    def __init__(self, data: bytes) -> None:
+        import io
+        self._buf = io.BytesIO(data)
+
+    def read(self, amt: Optional[int] = None) -> bytes:
+        return self._buf.read(amt) if amt is not None else self._buf.read()
+
+    def close(self) -> None:
+        pass
+
+
+def _emit_bedrock_converse_success(kwargs: dict, response: Any, t0: float) -> None:
+    latency_ms = int((time.monotonic() - t0) * 1000)
+    model_raw = kwargs.get("modelId")
+    usage_dict = response.get("usage") if isinstance(response, dict) else None
+    norm = normalize_bedrock_converse_usage(usage_dict)
+    _emit(_build_event(provider="bedrock", model_raw=model_raw, kwargs=kwargs, latency_ms=latency_ms,
+                        usage_complete=usage_dict is not None, norm=norm))
+
+
+def _wrap_bedrock_converse_event_stream(orig_iter, kwargs: dict, t0: float, ctx: Any, model_raw: Optional[str], otel: Any = "_unset"):
+    final_usage: Optional[dict] = None
+    errored = False
+    try:
+        for event in orig_iter:
+            u = (event.get("metadata") or {}).get("usage") if isinstance(event, dict) else None
+            final_usage = _merge_usage(final_usage, u)
+            yield event
+    except BaseException:
+        errored = True
+        raise
+    finally:
+        usage_complete = (final_usage is not None) and not errored
+        norm = normalize_bedrock_converse_usage(final_usage) if usage_complete else dict(_NULL_USAGE)
+        latency_ms = int((time.monotonic() - t0) * 1000)
+        _safe(_emit, _build_event(provider="bedrock", model_raw=model_raw, kwargs=kwargs,
+                                   latency_ms=latency_ms, usage_complete=usage_complete, norm=norm, ctx=ctx, otel=otel))
+
+
+def _emit_bedrock_invoke_model_success(kwargs: dict, response: Any, t0: float) -> None:
+    latency_ms = int((time.monotonic() - t0) * 1000)
+    model_raw = kwargs.get("modelId")
+    parsed: Optional[dict] = None
+    if isinstance(response, dict):
+        body_obj = response.get("body")
+        if body_obj is not None:
+            try:
+                raw_bytes = body_obj.read()
+            except Exception:
+                raw_bytes = b""
+            # Replace with a fresh readable stream over the SAME bytes so the
+            # host's own `response["body"].read()` still works -- reading it
+            # here does not add a network hop (already fully buffered locally
+            # by botocore for InvokeModel's single JSON body), only in-memory
+            # work the host would have done anyway.
+            response["body"] = _ReplayableBody(raw_bytes)
+            try:
+                parsed = _json.loads(raw_bytes) if raw_bytes else None
+            except Exception:
+                parsed = None
+        headers = ((response.get("ResponseMetadata") or {}).get("HTTPHeaders") or {})
+    else:
+        headers = {}
+    norm, usage_complete = normalize_bedrock_invoke_model(parsed, headers)
+    _emit(_build_event(provider="bedrock", model_raw=model_raw, kwargs=kwargs, latency_ms=latency_ms,
+                        usage_complete=usage_complete, norm=norm))
+
+
+def _wrap_bedrock_invoke_stream_events(orig_iter, kwargs: dict, t0: float, ctx: Any, model_raw: Optional[str], otel: Any = "_unset"):
+    final_metrics: Optional[dict] = None
+    errored = False
+    try:
+        for event in orig_iter:
+            try:
+                chunk_bytes = ((event or {}).get("chunk") or {}).get("bytes")
+                if chunk_bytes:
+                    payload = _json.loads(chunk_bytes)
+                    if isinstance(payload, dict) and "amazon-bedrock-invocationMetrics" in payload:
+                        final_metrics = payload["amazon-bedrock-invocationMetrics"]
+            except Exception:
+                pass
+            yield event
+    except BaseException:
+        errored = True
+        raise
+    finally:
+        usage_complete = final_metrics is not None and not errored
+        if usage_complete:
+            norm = {
+                "input_tokens": final_metrics.get("inputTokenCount"),
+                "output_tokens": final_metrics.get("outputTokenCount"),
+                "cache_read_tokens": None,
+                "cache_write_tokens": None,
+                "reasoning_tokens": None,
+            }
+        else:
+            norm = dict(_NULL_USAGE)
+        latency_ms = int((time.monotonic() - t0) * 1000)
+        _safe(_emit, _build_event(provider="bedrock", model_raw=model_raw, kwargs=kwargs,
+                                   latency_ms=latency_ms, usage_complete=usage_complete, norm=norm, ctx=ctx, otel=otel))
+
+
+def _wrap_bedrock_plain_method(client: Any, attr_name: str, on_success: Callable) -> None:
+    original = getattr(client, attr_name, None)
+    if original is None or getattr(original, "_tokenmeter_wrapped", False):
+        return
+
+    def wrapper(*args, **kwargs):
+        t0 = time.monotonic()
+        try:
+            result = original(*args, **kwargs)
+        except Exception as exc:
+            _safe(_emit_error, "bedrock", kwargs, t0, exc, model_key="modelId")
+            raise
+        _safe(on_success, kwargs, result, t0)
+        return result
+
+    wrapper._tokenmeter_wrapped = True  # type: ignore[attr-defined]
+    setattr(client, attr_name, wrapper)
+
+
+def _wrap_bedrock_converse_stream_method(client: Any) -> None:
+    attr_name = "converse_stream"
+    original = getattr(client, attr_name, None)
+    if original is None or getattr(original, "_tokenmeter_wrapped", False):
+        return
+
+    def wrapper(*args, **kwargs):
+        t0 = time.monotonic()
+        try:
+            result = original(*args, **kwargs)
+        except Exception as exc:
+            _safe(_emit_error, "bedrock", kwargs, t0, exc, model_key="modelId")
+            raise
+        if not isinstance(result, dict) or result.get("stream") is None:
+            return result
+        ctx = current()
+        try:
+            model_raw = kwargs.get("modelId")
+            result = dict(result)
+            result["stream"] = _wrap_bedrock_converse_event_stream(result["stream"], kwargs, t0, ctx, model_raw, _otel_link_tags())
+        except Exception:
+            logger.debug("tokenmeter: failed to wrap bedrock converse_stream", exc_info=True)
+        return result
+
+    wrapper._tokenmeter_wrapped = True  # type: ignore[attr-defined]
+    setattr(client, attr_name, wrapper)
+
+
+def _wrap_bedrock_invoke_stream_method(client: Any) -> None:
+    attr_name = "invoke_model_with_response_stream"
+    original = getattr(client, attr_name, None)
+    if original is None or getattr(original, "_tokenmeter_wrapped", False):
+        return
+
+    def wrapper(*args, **kwargs):
+        t0 = time.monotonic()
+        try:
+            result = original(*args, **kwargs)
+        except Exception as exc:
+            _safe(_emit_error, "bedrock", kwargs, t0, exc, model_key="modelId")
+            raise
+        if not isinstance(result, dict) or result.get("body") is None:
+            return result
+        ctx = current()
+        try:
+            model_raw = kwargs.get("modelId")
+            result = dict(result)
+            result["body"] = _wrap_bedrock_invoke_stream_events(result["body"], kwargs, t0, ctx, model_raw, _otel_link_tags())
+        except Exception:
+            logger.debug("tokenmeter: failed to wrap bedrock invoke_model_with_response_stream", exc_info=True)
+        return result
+
+    wrapper._tokenmeter_wrapped = True  # type: ignore[attr-defined]
+    setattr(client, attr_name, wrapper)
+
+
+def _wrap_bedrock(client: Any) -> None:
+    # boto3 bedrock-runtime client: a dynamically-generated botocore client,
+    # detected structurally by the methods it exposes (no `import boto3`).
+    if not any(
+        hasattr(client, m) for m in ("converse", "converse_stream", "invoke_model", "invoke_model_with_response_stream")
+    ):
+        return
+    _wrap_bedrock_plain_method(client, "converse", _emit_bedrock_converse_success)
+    _wrap_bedrock_converse_stream_method(client)
+    _wrap_bedrock_plain_method(client, "invoke_model", _emit_bedrock_invoke_model_success)
+    _wrap_bedrock_invoke_stream_method(client)
+
+
+def wrap(client: Any, *, model_map: Optional[dict] = None) -> Any:
+    """Wraps a client instance in place and returns it. Idempotent (re-wrapping
+    an already-wrapped client/method is a no-op, checked via a marker attribute).
+
+    `model_map`: optional `{deployment_name: model_id}` mapping, used only for
+    Azure OpenAI clients (spike §2.2 -- Azure's Responses API returns the
+    deployment name, not the canonical model id, which breaks §7.1's pricing
+    aliasing unless resolved here). Ignored for every other provider.
+
+    Each provider's detection+wrapping runs in its own try/except so one
+    provider's failure (e.g. an unexpected client shape) never prevents the
+    others from being wrapped.
+    """
+    for step in (
+        lambda: _wrap_openai_or_azure(client, model_map),
+        lambda: _wrap_anthropic(client),
+        lambda: _wrap_gemini(client),
+        lambda: _wrap_bedrock(client),
+    ):
+        try:
+            step()
+        except Exception:
+            logger.debug("tokenmeter: wrap() step failed, continuing", exc_info=True)
+    return client