Rework telemetry to support durable execution case

src/ai/agents/agent.py

@@ -723,7 +723,8 @@ async def _real(
             # A tool exception is caught and converted to an error
             # result before it reaches this block, so it never hits the
             # span's own except path — thread it through explicitly.
-            sp.error = res.exception
+            if res.exception is not None:
+                sp.error = telemetry.SpanError.from_exception(res.exception)
             return res
 
 

src/ai/agents/hooks.py

@@ -229,7 +229,8 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
         )
 
         await rt.put_hook(hook_part)
-        await sp.add_event(telemetry.HOOK_DEFERRED)
+        sp.add_event(telemetry.HOOK_DEFERRED)
+        await sp.push()
 
         # Await resolution — may be resolved externally or cancelled.
         try:
@@ -240,15 +241,17 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel:
             attrs: dict[str, Any] = {}
             if exc.args and exc.args[0] is not None:
                 attrs["reason"] = exc.args[0]
-            await sp.add_event(telemetry.HOOK_CANCELLED, **attrs)
+            sp.add_event(telemetry.HOOK_CANCELLED, attrs)
+            await sp.push()
             raise
         finally:
             # Clean up live registry.
             registry._live_hooks.pop(label, None)
 
         sp.data.status = "resolved"
         sp.data.resolution = resolution
-        await sp.add_event(telemetry.HOOK_RESOLVED)
+        sp.add_event(telemetry.HOOK_RESOLVED)
+        await sp.push()
 
         # Emit resolved signal.
         await rt.put_hook(

src/ai/experimental_telemetry/__init__.py

@@ -1,4 +1,4 @@
-"""Telemetry: spans, adapters, and the ambient current span.
+"""Telemetry: spans, sinks, adapters, and the ambient current span.
 
 Experimental: not part of the stable API, may change or be removed.
 
@@ -14,21 +14,29 @@
     Adapter,
     AiGenerateSpanData,
     AiStreamSpanData,
+    Collector,
     CustomSpanData,
     HookSpanData,
     LoopTurnSpanData,
     RunSpanData,
+    Sink,
     Span,
     SpanData,
+    SpanError,
     SpanEvent,
-    SpanRef,
     ToolExecutionSpanData,
+    create_span,
     current,
-    current_ref,
+    enabled,
+    flush,
+    now_ns,
+    push_all,
     register,
     span,
     unregister,
     use_clock,
+    use_sink,
+    use_span,
     wrap_span,
 )
 
@@ -41,20 +49,28 @@
     "Adapter",
     "AiGenerateSpanData",
     "AiStreamSpanData",
+    "Collector",
     "CustomSpanData",
     "HookSpanData",
     "LoopTurnSpanData",
     "RunSpanData",
+    "Sink",
     "Span",
     "SpanData",
+    "SpanError",
     "SpanEvent",
-    "SpanRef",
     "ToolExecutionSpanData",
+    "create_span",
     "current",
-    "current_ref",
+    "enabled",
+    "flush",
+    "now_ns",
+    "push_all",
     "register",
     "span",
     "unregister",
     "use_clock",
+    "use_sink",
+    "use_span",
     "wrap_span",
 ]

src/ai/experimental_telemetry/otel.py

@@ -17,11 +17,28 @@
 nest under any raw otel span the caller already has open.  Spans
 opened with ``set_as_current=False`` are never attached: they don't
 parent concurrent work in our tree, so they must not do it in otel's.
+
+Identity: otel ids are derived deterministically from the framework's
+span/trace ids (a truncated hash), via an id generator installed on the
+SDK tracer provider.  That is what makes spans that cross processes
+line up: a span finished and pushed in a different process than it
+started in exports under the same otel identity its children already
+parented to, and a re-emitted span comes out under the same ids instead
+of duplicating the tree.  A span whose parent is not live in this
+process is parented through those derived ids directly.
+
+With a non-SDK tracer provider (no id generator to install) the
+adapter still works for in-process traces, but ids fall back to
+whatever the tracer mints, so cross-process parenting degrades; a
+warning is logged once if that comes up.
 """
 
 from __future__ import annotations
 
+import contextvars
+import hashlib
 import json
+import logging
 from typing import TYPE_CHECKING, Any
 
 from .. import errors
@@ -30,6 +47,8 @@
 try:
     from opentelemetry import context as otel_context
     from opentelemetry import trace as otel_trace
+    from opentelemetry.sdk import trace as sdk_trace
+    from opentelemetry.sdk.trace import id_generator as sdk_id_generator
 except ModuleNotFoundError as exc:  # pragma: no cover
     raise errors.InstallationError(
         "could not import `opentelemetry`, which is required for the otel "
@@ -42,6 +61,8 @@
 
     from ..types import messages as messages_
 
+logger = logging.getLogger(__name__)
+
 
 def _messages_json(messages: list[messages_.Message]) -> str:
     return "[" + ",".join(m.model_dump_json() for m in messages) + "]"
@@ -112,21 +133,139 @@ def _attributes(sp: telemetry.Span) -> dict[str, Any]:
     return attrs
 
 
-def install(
-    *, tracer_provider: otel_trace.TracerProvider | None = None
-) -> telemetry.Adapter:
-    """Create the otel adapter, register it, and return it.
+def _derive_trace_id(id_: str) -> int:
+    """Derive a stable, nonzero otel trace id (128 bits) from one of ours."""
+    return int.from_bytes(hashlib.sha256(id_.encode()).digest()[:16]) or 1
 
-    Uses the global tracer provider unless one is passed.
-    """
-    tracer = otel_trace.get_tracer("ai", tracer_provider=tracer_provider)
 
-    @telemetry.wrap_span
-    async def otel_spans(span: telemetry.Span) -> AsyncGenerator[None, Any]:
-        otel_span = tracer.start_span(_name(span), start_time=span.started_at)
+def _derive_span_id(id_: str) -> int:
+    """Derive a stable, nonzero otel span id (64 bits) from one of ours."""
+    return int.from_bytes(hashlib.sha256(id_.encode()).digest()[:8]) or 1
+
+
+# The otel Tracer offers no way to dictate a span's ids, but derived
+# ids must come out exactly (see module docstring); the generator
+# installed on the SDK provider honors a preset while it is set.
+_preset_ids: contextvars.ContextVar[tuple[int, int] | None] = (
+    contextvars.ContextVar("otel_preset_ids", default=None)
+)
+
+
+class _PresetIdGenerator(sdk_id_generator.RandomIdGenerator):
+    """Uses the preset ids when set; defers to ``inner`` otherwise."""
+
+    def __init__(self, inner: sdk_id_generator.IdGenerator) -> None:
+        self._inner = inner
+
+    def generate_trace_id(self) -> int:
+        preset = _preset_ids.get()
+        return (
+            preset[0] if preset is not None else self._inner.generate_trace_id()
+        )
+
+    def generate_span_id(self) -> int:
+        preset = _preset_ids.get()
+        return (
+            preset[1] if preset is not None else self._inner.generate_span_id()
+        )
+
+
+class OtelAdapter(telemetry.Adapter):
+    """Maps framework spans onto otel spans; see the module docstring."""
+
+    def __init__(
+        self, *, tracer_provider: otel_trace.TracerProvider | None = None
+    ) -> None:
+        provider = tracer_provider or otel_trace.get_tracer_provider()
+        self._provider = provider
+        self._live: dict[str, otel_trace.Span] = {}
+        self._deterministic_ids = False
+        self._warned_random_ids = False
+        if isinstance(provider, sdk_trace.TracerProvider):
+            provider.id_generator = _PresetIdGenerator(provider.id_generator)
+            self._deterministic_ids = True
+        self._tracer = otel_trace.get_tracer("ai", tracer_provider=provider)
+
+    def span_name(self, span_: telemetry.Span, /) -> str:
+        """Return the exported otel span name.  Override to customize."""
+        return _name(span_)
+
+    def span_attributes(self, span_: telemetry.Span, /) -> dict[str, Any]:
+        """Return the attributes set at span end.  Override to enrich.
+
+        ::
+
+            class MyAdapter(otel.OtelAdapter):
+                def span_attributes(self, span_):
+                    return super().span_attributes(span_) | {"k": "v"}
+        """
+        return _attributes(span_)
+
+    def flush(self) -> None:
+        """Flush the provider's exporters, when it has any (SDK provider)."""
+        force_flush = getattr(self._provider, "force_flush", None)
+        if force_flush is not None:
+            force_flush()
+
+    def shutdown(self) -> None:
+        """Flush and stop the provider; spans pushed after this are lost."""
+        shutdown = getattr(self._provider, "shutdown", None)
+        if shutdown is not None:
+            shutdown()
+
+    def _parent_context(
+        self, span_: telemetry.Span
+    ) -> otel_context.Context | None:
+        # ``None`` means the ambient otel context: the parent's otel
+        # span is attached there when it is live in this process, and
+        # for our roots it lets any raw otel span the caller holds
+        # adopt the trace.
+        if span_.parent_id is None or span_.parent_id in self._live:
+            return None
+        # The parent is not live here — it started in another process,
+        # or this span arrived as a finished record.  Parent on the
+        # derived identity in an empty context so the pieces line up
+        # when the parent itself exports.
+        if not self._deterministic_ids and not self._warned_random_ids:
+            self._warned_random_ids = True
+            logger.warning(
+                "otel adapter: continuing a trace from another process "
+                "needs an SDK tracer provider to control span ids; ids "
+                "will not line up across processes"
+            )
+        return otel_trace.set_span_in_context(
+            otel_trace.NonRecordingSpan(
+                otel_trace.SpanContext(
+                    trace_id=_derive_trace_id(span_.trace_id),
+                    span_id=_derive_span_id(span_.parent_id),
+                    is_remote=True,
+                    trace_flags=otel_trace.TraceFlags(
+                        otel_trace.TraceFlags.SAMPLED
+                    ),
+                )
+            ),
+            otel_context.Context(),
+        )
+
+    async def wrap_span(
+        self, span_: telemetry.Span, /
+    ) -> AsyncGenerator[None, Any]:
+        parent_context = self._parent_context(span_)
+        preset_token = _preset_ids.set(
+            (_derive_trace_id(span_.trace_id), _derive_span_id(span_.id))
+        )
+        try:
+            otel_span = self._tracer.start_span(
+                self.span_name(span_),
+                context=parent_context,
+                start_time=span_.started_at,
+            )
+        finally:
+            _preset_ids.reset(preset_token)
+        self._live[span_.id] = otel_span
         token = (
             otel_context.attach(otel_trace.set_span_in_context(otel_span))
-            if span.set_as_current
+            if span_.set_as_current
             else None
         )
         try:
@@ -146,15 +285,24 @@ async def otel_spans(span: telemetry.Span) -> AsyncGenerator[None, Any]:
         finally:
             if token is not None:
                 otel_context.detach(token)
-            for key, value in _attributes(span).items():
+            self._live.pop(span_.id, None)
+            for key, value in self.span_attributes(span_).items():
                 otel_span.set_attribute(key, value)
-            if span.error is not None:
-                if isinstance(span.error, Exception):
-                    otel_span.record_exception(span.error)
+            if span_.error is not None:
                 otel_span.set_status(
-                    otel_trace.StatusCode.ERROR, str(span.error)
+                    otel_trace.StatusCode.ERROR,
+                    f"{span_.error.type}: {span_.error.message}",
                 )
-            otel_span.end(end_time=span.ended_at)
+            otel_span.end(end_time=span_.ended_at)
+
+
+def install(
+    *, tracer_provider: otel_trace.TracerProvider | None = None
+) -> OtelAdapter:
+    """Create the otel adapter, register it, and return it.
 
-    telemetry.register(otel_spans)
-    return otel_spans
+    Uses the global tracer provider unless one is passed.
+    """
+    adapter = OtelAdapter(tracer_provider=tracer_provider)
+    telemetry.register(adapter)
+    return adapter