diff --git a/TELEMETRY-CONTRACT.md b/TELEMETRY-CONTRACT.md index 50c619c..513b8fd 100644 --- a/TELEMETRY-CONTRACT.md +++ b/TELEMETRY-CONTRACT.md @@ -119,7 +119,13 @@ finds. Without totals on it, that query returns nothing, because summing the chi having already found them. `claude-agents` also writes `gen_ai.conversation.id` on the root, from the session id on the -CLI's `system` / `init` message. See section 4 for the third place it appears. +CLI's `system` / `init` message, write-if-absent. A caller-supplied id from `conversation_id(...)` +(TypeScript: `withConversationId`) wins. When the caller supplies none, the session id is used. +An app that opens a fresh CLI session per turn and re-feeds history must pass its own id, or each +turn becomes its own conversation. See section 4 for the third place it appears. + +When a conversation id is bound, every handler stamps it on root, `chat`, and `execute_tool` (and +on `ld.ai.graph`) via the shared span processor. No id is invented when the caller supplies none. --- @@ -182,7 +188,7 @@ requested name. See section 2a. | Key | Value | |---|---| | `gen_ai.response.id` | provider request id | -| `gen_ai.conversation.id` | session id | +| `gen_ai.conversation.id` | session id, write-if-absent against a caller-supplied id | | `gen_ai.agent.name` | subagent type | --- @@ -200,8 +206,9 @@ requested name. See section 2a. No usage attributes. A tool call spends no tokens. `claude-agents` also writes `gen_ai.conversation.id` here, from the session id on the tool-use -hook input, when present. That attribute therefore appears on all three span types for this one -handler: root, `chat`, and `execute_tool`. +hook input, when present, write-if-absent. That attribute therefore appears on all three span types +for this one handler: root, `chat`, and `execute_tool`. A caller-supplied id is already on the span +when `conversation_id(...)` is bound. --- diff --git a/examples/history.py b/examples/history.py index e2ec188..8087504 100644 --- a/examples/history.py +++ b/examples/history.py @@ -17,8 +17,8 @@ import sys import examples.register # noqa: F401 – side-effect: populate global_registry -from examples.utils import new_context, write_output -from launchdarkly_ai_server import config, global_registry +from examples.utils import new_context, new_conversation_id, write_output +from launchdarkly_ai_server import config, conversation_id, global_registry HISTORY = [ { @@ -48,10 +48,11 @@ async def run(key: str, user_input: str) -> None: prompt = user_input or HISTORY_PROMPT - response = await config( - key=key, - registry=global_registry, - ).invoke(prompt, new_context(), variables=None, history=HISTORY) + with conversation_id(new_conversation_id("history-example")): + response = await config( + key=key, + registry=global_registry, + ).invoke(prompt, new_context(), variables=None, history=HISTORY) text = str( response.get("response", "") diff --git a/examples/streaming.py b/examples/streaming.py index 6cedc33..2dec1c7 100644 --- a/examples/streaming.py +++ b/examples/streaming.py @@ -2,6 +2,12 @@ Example: config().stream() — tokens are printed as they arrive, then the final usage + judge results are logged when the stream ends. +Also the end-to-end check for call-time conversation binding: the generator is built inside the +``conversation_id`` block and iterated *outside* it, which is what a chat app does when it hands +the stream to a transport. An async generator body does not run until the first ``__anext__``, so +before ``stream()`` bound at call time this produced spans with no ``gen_ai.conversation.id`` at +all — silently. Every span of this run should carry the id printed below. + Usage (via main.py): python main.py streaming "" """ @@ -11,15 +17,19 @@ import sys import examples.register # noqa: F401 – side-effect: populate global_registry -from examples.utils import new_context -from launchdarkly_ai_server import config, global_registry +from examples.utils import new_context, new_conversation_id +from launchdarkly_ai_server import config, conversation_id, global_registry async def run(key: str, user_input: str) -> None: - stream = config( - key=key, - registry=global_registry, - ).stream(user_input, new_context()) + conversation = new_conversation_id("streaming-example") + print(f"[conversation] {conversation}", file=sys.stderr) + + with conversation_id(conversation): + stream = config( + key=key, + registry=global_registry, + ).stream(user_input, new_context()) async for event in stream: if event["type"] == "chunk": diff --git a/examples/utils.py b/examples/utils.py index 6bb4f5a..191a94c 100644 --- a/examples/utils.py +++ b/examples/utils.py @@ -9,6 +9,7 @@ from datetime import UTC, datetime from pathlib import Path from typing import Any +from uuid import uuid4 def new_context() -> dict[str, Any]: @@ -17,6 +18,17 @@ def new_context() -> dict[str, Any]: return {"kind": "user", "key": key} +def new_conversation_id(label: str) -> str: + """A fresh conversation id per run. + + A constant would collapse every run — by every developer, and every CI pass — into one + ever-growing conversation in LaunchDarkly's view: a misleading demo of the very feature it is + demonstrating. An id should be stable across the turns of one conversation and distinct across + conversations. + """ + return f"{label}-{uuid4().hex[:8]}" + + def _default_encoder(obj: Any) -> Any: if dataclasses.is_dataclass(obj) and not isinstance(obj, type): return dataclasses.asdict(obj) diff --git a/packages/claude-agents/agents.md b/packages/claude-agents/agents.md index 4028234..8c4ff71 100644 --- a/packages/claude-agents/agents.md +++ b/packages/claude-agents/agents.md @@ -136,6 +136,11 @@ Span event after the call: On error: `span.record_exception(exc)`, status set to ERROR, span ended, error re-raised. +`gen_ai.conversation.id` is a caller-supplied id from `conversation_id(...)`, or the CLI +`session_id` from the `init` message when the caller supplied none. Write-if-absent: the caller +id wins. An app that opens a fresh CLI session per turn and re-feeds history must pass its own +conversation id, or each turn becomes its own conversation. + --- ## OTel Setup diff --git a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py index 4e7a54a..cc8bf5e 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py @@ -40,6 +40,7 @@ end_span_once, end_unfinished_spans, parse_template, + set_conversation_id_if_absent, set_input_content_attributes, set_output_content_attributes, set_tool_call_content_attributes, @@ -247,7 +248,7 @@ async def _pre_tool_use( # Same grouping key as the root and as the CLI's own spans; the hook input is where this # side sees it without waiting for a message. See TELEMETRY-CONTRACT.md section 4. if span is not None and session_id: - span.set_attribute("gen_ai.conversation.id", session_id) + set_conversation_id_if_absent(span, session_id) # Filed before the content write, not after. Serialising the arguments can raise, and a raise # out of an unfiled span leaves it open with nothing tracking it: close_open_spans and the # teardown both walk this dict, so a span missing from it is a span that never exports. diff --git a/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py b/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py index 422b074..2921b29 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py @@ -40,6 +40,7 @@ add_cached_tokens_to_input, end_span_once, number_or_zero, + set_conversation_id_if_absent, set_input_content_attributes, set_ld_span_attributes, set_model_identity_attributes, @@ -181,7 +182,9 @@ def record_conversation_id(span: Any, message: Any) -> None: It is the only key LaunchDarkly's trace view groups a conversation on, and the ``init`` system message is where this side first learns it. The ``chat`` and ``execute_tool`` children read the same id off their own message and hook input, so one run does not split into several - conversations. Set once: the id does not change within a run. + conversations. Write-if-absent: a caller-supplied id from ``conversation_id`` is already on the + span and must not be overwritten. Apps that open a fresh CLI session per turn and re-feed history + must pass their own conversation id, or each turn becomes its own conversation. """ if ( span is None @@ -191,7 +194,7 @@ def record_conversation_id(span: Any, message: Any) -> None: return session_id = message.data.get("session_id") if session_id: - span.set_attribute("gen_ai.conversation.id", session_id) + set_conversation_id_if_absent(span, session_id) def record_native_tools( @@ -532,7 +535,7 @@ def _emit(self, inference: _Inference) -> None: if inference.request_id: span.set_attribute("gen_ai.response.id", inference.request_id) if inference.session_id: - span.set_attribute("gen_ai.conversation.id", inference.session_id) + set_conversation_id_if_absent(span, inference.session_id) # Absent in practice: measured against Agent SDK 0.3.220, stop_reason and stop_details are # both null on every assistant message. Written only when the SDK populates it — never # synthesised from the presence of a tool-use block. diff --git a/packages/claude-agents/tests/test_handler.py b/packages/claude-agents/tests/test_handler.py index 23535a8..dc138c3 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -42,6 +42,7 @@ create_claude_agents_handler, partition_tools, ) +from launchdarkly_ai_server import ConversationIdSpanProcessor, conversation_id # --------------------------------------------------------------------------- # A real tracer provider, reset between tests @@ -49,6 +50,7 @@ _exporter = InMemorySpanExporter() _provider = TracerProvider() +_provider.add_span_processor(ConversationIdSpanProcessor()) _provider.add_span_processor(SimpleSpanProcessor(_exporter)) trace.set_tracer_provider(_provider) @@ -687,6 +689,48 @@ async def test_no_conversation_id_without_init( await create_claude_agents_handler()(BASE_CONFIG, "q") assert "gen_ai.conversation.id" not in root().attributes + async def test_caller_conversation_id_wins_over_session_id( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + yield init_message("sess-abc") + yield assistant_message(session_id="sess-abc") + hooks = kwargs["options"].hooks + await hooks["PreToolUse"][0].hooks[0]( + { + "hook_event_name": "PreToolUse", + "tool_name": "mcp__tool-mcp__search", + "tool_use_id": "tu-1", + "tool_input": {}, + "session_id": "sess-abc", + }, + "tu-1", + None, + ) + await hooks["PostToolUse"][0].hooks[0]( + { + "hook_event_name": "PostToolUse", + "tool_name": "mcp__tool-mcp__search", + "tool_use_id": "tu-1", + "tool_response": "r", + }, + "tu-1", + None, + ) + yield result_message() + + monkeypatch.setattr(handler_mod, "query", _query) + with conversation_id("thread-stable"): + await create_claude_agents_handler()( + TOOL_CONFIG, "q", {"search": lambda _: "r"} + ) + assert root().attributes["gen_ai.conversation.id"] == "thread-stable" + assert named("chat ")[0].attributes["gen_ai.conversation.id"] == "thread-stable" + assert ( + named("execute_tool ")[0].attributes["gen_ai.conversation.id"] + == "thread-stable" + ) + # --------------------------------------------------------------------------- # Chat span attributes — sections 2a, 3, 5b, 8 diff --git a/packages/claude-agents/tests/test_native_graph.py b/packages/claude-agents/tests/test_native_graph.py index 9f204db..bf301ff 100644 --- a/packages/claude-agents/tests/test_native_graph.py +++ b/packages/claude-agents/tests/test_native_graph.py @@ -652,3 +652,62 @@ async def test_build_tool_mcp_throws_when_tool_not_in_handlers(self) -> None: handlers: dict[str, Any] = {} await build_tool_mcp(config_tools, handlers) # The actual error is raised when the fn is *called*, not when mcp is built + + +class TestNativeGraphConversationId: + """The telemetry contract claims the conversation id reaches ``ld.ai.graph`` spans. + + The span is opened after an ``await`` on the graph definition, so this pins that the binding + survives the await chain rather than only covering spans started synchronously in the block. + """ + + @pytest.mark.asyncio + async def test_stamps_conversation_id_on_graph_span(self) -> None: + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + import launchdarkly_ai_claude_agents.native_graph as ng_mod + from launchdarkly_ai_server.conversation import ( + GEN_AI_CONVERSATION_ID, + ConversationIdSpanProcessor, + conversation_id, + ) + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(ConversationIdSpanProcessor()) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + # A real tracer behind the module's `trace` lookup: the span and the processor are real, + # only the global-provider registration is bypassed (it is process-wide and set once). + real_trace = MagicMock() + real_trace.get_tracer.return_value = provider.get_tracer( + "@launchdarkly/ai-claude-agents" + ) + + mock_sdk = _make_sdk_mock("done") + graph_def = _make_graph_def() + + with patch( + "importlib.import_module", + side_effect=lambda n: ( + mock_sdk if n == "claude_agent_sdk" else __import__(n) + ), + ): + with patch.object(ng_mod, "trace", real_trace): + with patch.object(ng_mod, "_HAS_OTEL", True): + with conversation_id("thread-graph"): + await to_claude_agents(_make_def_promise(graph_def)).invoke( + "hi" + ) + + graph_spans = [ + s for s in exporter.get_finished_spans() if s.name == "ld.ai.graph" + ] + assert len(graph_spans) == 1 + assert (graph_spans[0].attributes or {}).get( + GEN_AI_CONVERSATION_ID + ) == "thread-graph" diff --git a/packages/client/agents.md b/packages/client/agents.md index 9780296..61d99cf 100644 --- a/packages/client/agents.md +++ b/packages/client/agents.md @@ -21,6 +21,7 @@ No other `launchdarkly-ai-*` package may define or duplicate these. They import | File | Responsibility | |---|---| +| `src/launchdarkly_ai_server/conversation.py` | `conversation_id`, `ConversationIdSpanProcessor` — stamps `gen_ai.conversation.id` | | `src/launchdarkly_ai_server/lifecycle.py` | `init_client`, `get_client`, `shutdown`, `extract_variation` | | `src/launchdarkly_ai_server/client.py` | `config()`, `ConfigInstance` | | `src/launchdarkly_ai_server/tracking.py` | `execute_and_track`, `execute_and_stream`, `wrap_tool_handlers`, `parse_usage` | @@ -41,6 +42,7 @@ Key symbols exported from `launchdarkly_ai_server`: ```python # Lifecycle from launchdarkly_ai_server import init_client, get_client, shutdown, extract_variation +from launchdarkly_ai_server import conversation_id, set_conversation_id_if_absent, ConversationIdSpanProcessor # Types from launchdarkly_ai_server import ( @@ -123,9 +125,39 @@ Handlers may return any of these — the client normalizes them before emitting --- +## Conversation grouping + +LaunchDarkly's conversation view groups spans on `gen_ai.conversation.id`. Bind a caller-supplied id around any `invoke()` / `stream()` / `graph().invoke()` call: + +```python +from launchdarkly_ai_server import conversation_id, config + +with conversation_id("thread-123"): + await config(key=key, handler=handler).invoke(user_input, ctx) +``` + +`stream()` binds at call time rather than on first `__anext__`, so building the generator inside +the block and iterating it later — the normal shape for a chat app — keeps the id: + +```python +with conversation_id("thread-123"): + gen = config(key=key, handler=handler).stream(user_input, ctx) +async for event in gen: # spans opened here still carry thread-123 + ... +``` + +Only the id is re-applied per step; the ambient context at iteration time is otherwise untouched, +so streaming span parenting is the same as it is with no id bound. + +`init_client()` registers a span processor that stamps the id write-if-absent on every SDK span (root, chat, execute_tool, graph). The processor is registered on the *global* tracer provider, so it is scoped to spans from `@launchdarkly/ai-*` tracers only — a caller-supplied id must not land on third-party instrumentation spans (HTTP, Postgres, the outbound provider call). No id is invented when the caller supplies none — a UUID, a trace id, or a content hash would violate the semantic conventions. + +This is an OTel context value, not W3C baggage, so the id does not leak onto outbound provider HTTP calls. A multi-tenant process must bind a different id per request; do not put it on the tracer resource. + +--- + ## OTel Setup -The core client owns all OTel initialization. `init_client()` configures a `TracerProvider` with a `BatchSpanProcessor` and an OTLP HTTP exporter when the optional OTel packages are installed. +The core client owns all OTel initialization. `init_client()` configures a `TracerProvider` with `ConversationIdSpanProcessor` and a `BatchSpanProcessor` plus an OTLP HTTP exporter when the optional OTel packages are installed. **Required packages:** diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 73bfc3d..80b959b 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -16,6 +16,11 @@ text_message, to_semconv_finish_reason, ) +from .conversation import ( + ConversationIdSpanProcessor, + conversation_id, + set_conversation_id_if_absent, +) from .graph import GraphInstance, graph, resolve_graph from .judges import build_judge_tasks, run_judge, run_judges from .lifecycle import ( @@ -186,6 +191,10 @@ # client "config", "ConfigInstance", + # conversation + "ConversationIdSpanProcessor", + "conversation_id", + "set_conversation_id_if_absent", # graph "graph", "resolve_graph", diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 570ac0d..f64422d 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator, Callable from typing import Any +from .conversation import bind_conversation_id from .judges import build_judge_tasks, run_judges from .lifecycle import extract_variation from .registry import resolve_handlers, resolve_tools @@ -147,7 +148,25 @@ async def invoke( track_data=track_data, ) - async def stream( + def stream( + self, + user_input: str | None, + context: LDContext, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> AsyncGenerator[StreamEvent, None]: + """Stream events for this config. + + Deliberately not an ``async def`` with ``yield``: a generator body does not run until the + first ``__anext__``, by which point a :func:`conversation_id` block wrapped around this + call has already exited. Binding here — at call time — is what lets a caller hand the + generator off and iterate it later. + """ + return bind_conversation_id( + self._stream_events(user_input, context, variables, history) + ) + + async def _stream_events( self, user_input: str | None, context: LDContext, diff --git a/packages/client/src/launchdarkly_ai_server/conversation.py b/packages/client/src/launchdarkly_ai_server/conversation.py new file mode 100644 index 0000000..72dc969 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/conversation.py @@ -0,0 +1,164 @@ +"""Caller-supplied ``gen_ai.conversation.id``. + +A dedicated OTel context key, not W3C baggage: the id must not leak onto outbound +provider HTTP calls. A multi-tenant process binds a different id per request. +""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator, Iterator +from contextlib import aclosing, contextmanager +from typing import TYPE_CHECKING, Any + +from opentelemetry import context as otel_context +from opentelemetry.trace import Span + +if TYPE_CHECKING: + # ``opentelemetry-sdk`` is an optional extra (``[otel]``), so it must not be imported at + # runtime — this module has to work on an api-only install. Subclassing under + # ``TYPE_CHECKING`` still gets the interface checked, which a ``cast`` at the registration + # site would not. + from opentelemetry.sdk.trace import SpanProcessor as _SpanProcessorBase +else: + _SpanProcessorBase = object + +GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" + +_CONV_KEY = otel_context.create_key("launchdarkly.gen_ai.conversation.id") + +# Every tracer this SDK creates is named "@launchdarkly/ai-". The processor is registered +# on the *global* provider, so without this gate it stamps a caller-supplied id onto every span in +# the process — Postgres queries, inbound HTTP server spans, and the outbound provider call itself. +_LD_TRACER_PREFIX = "@launchdarkly/" + + +def _read_attribute(span: Any, key: str) -> Any: + attrs = getattr(span, "attributes", None) + if attrs is None: + return None + return attrs.get(key) + + +def set_conversation_id_if_absent(span: Any, conversation: str) -> None: + """Write ``gen_ai.conversation.id`` only when the span does not already carry one.""" + if not conversation or span is None: + return + existing = _read_attribute(span, GEN_AI_CONVERSATION_ID) + if isinstance(existing, str) and existing: + return + span.set_attribute(GEN_AI_CONVERSATION_ID, conversation) + + +def _is_launchdarkly_span(span: Any) -> bool: + """True only when the span came from one of this SDK's own tracers. + + Deliberately conservative: an unrecognisable scope means "not ours", so an id is never sprayed + across unrelated telemetry. The companion test asserts LD spans *are* stamped, so a rename of + the scope attribute fails the suite loudly rather than silently disabling the feature. + """ + scope_name = getattr(getattr(span, "instrumentation_scope", None), "name", None) + return isinstance(scope_name, str) and scope_name.startswith(_LD_TRACER_PREFIX) + + +def _conversation_id_from(ctx: otel_context.Context | None) -> str | None: + if ctx is None: + return None + value = otel_context.get_value(_CONV_KEY, ctx) + return value if isinstance(value, str) and value else None + + +@contextmanager +def conversation_id(conversation: str | None) -> Iterator[None]: + """Bind a caller-supplied conversation id for the duration of the ``with`` block. + + Every span the SDK creates while this is bound receives ``gen_ai.conversation.id``, + provided :class:`ConversationIdSpanProcessor` is registered — which ``init_client()`` + does when OTel is installed. An empty, whitespace, or ``None`` id is treated as unbound — + the natural call site is an optional value such as ``conversation_id(request.thread_id)``, + which must degrade to an unstamped trace rather than raise. + """ + trimmed = conversation.strip() if isinstance(conversation, str) else "" + if not trimmed: + yield + return + token = otel_context.attach(otel_context.set_value(_CONV_KEY, trimmed)) + try: + yield + finally: + otel_context.detach(token) + + +async def _stream_with_bound_id( + generator: AsyncGenerator[Any, None], conversation: str +) -> AsyncGenerator[Any, None]: + """Bind ``conversation`` for the whole iteration of ``generator``. + + Attached once and detached once, deliberately. Re-attaching per step would have to detach per + step too, and a streaming handler normally holds a span current *across* a ``yield`` — so that + detach would reset the contextvar past the handler's own live ``start_as_current_span`` and + reparent every span it opens after resuming. Binding for the whole iteration leaves the + generator's own context untouched, so span parenting matches an unbound stream exactly. + + The generator is closed before the detach so its ``__exit__`` unwinds inside our scope, keeping + contextvar tokens in LIFO order. + + Trade-off: the id is also bound while the consumer's body runs between chunks, so spans the + consumer opens mid-stream are stamped too. That is the conversation the consumer is streaming. + """ + token = otel_context.attach(otel_context.set_value(_CONV_KEY, conversation)) + try: + async with aclosing(generator): + async for item in generator: + yield item + finally: + otel_context.detach(token) + + +def bind_conversation_id( + generator: AsyncGenerator[Any, None], +) -> AsyncGenerator[Any, None]: + """Pin the currently-bound conversation id onto ``generator`` at call time. + + An ``async def`` with ``yield`` does not start its body until the first ``__anext__``, which + for a streaming caller is normally after the :func:`conversation_id` block has exited. Reading + the id here — eagerly, before any iteration — is what lets a caller bind, hand the generator + off, and iterate it later. + """ + conversation = _conversation_id_from(otel_context.get_current()) + if not conversation: + return generator + return _stream_with_bound_id(generator, conversation) + + +class ConversationIdSpanProcessor(_SpanProcessorBase): + """Stamps ``gen_ai.conversation.id`` write-if-absent on every span. + + Structurally a ``SpanProcessor``; the base is only real to a type checker so that an + api-only install (no ``[otel]`` extra) still imports this module. + ``init_client()`` registers it ahead of ``BatchSpanProcessor``. + """ + + def on_start( + self, span: Span, parent_context: otel_context.Context | None = None + ) -> None: + ctx = ( + parent_context if parent_context is not None else otel_context.get_current() + ) + conv = _conversation_id_from(ctx) or _conversation_id_from( + otel_context.get_current() + ) + if conv and _is_launchdarkly_span(span): + set_conversation_id_if_absent(span, conv) + + def on_end(self, span: Any) -> None: + return None + + def _on_ending(self, span: Any) -> None: + # Python SDK calls ``_on_ending`` from ``span.end()``; JS uses ``onEnd``. + return None + + def shutdown(self) -> None: + return None + + def force_flush(self, timeout_millis: int | None = None) -> bool: + return True diff --git a/packages/client/src/launchdarkly_ai_server/lifecycle.py b/packages/client/src/launchdarkly_ai_server/lifecycle.py index acad334..8969f2e 100644 --- a/packages/client/src/launchdarkly_ai_server/lifecycle.py +++ b/packages/client/src/launchdarkly_ai_server/lifecycle.py @@ -87,6 +87,9 @@ def _setup_telemetry(sdk_key: str, options: InitClientOptions | None = None) -> resource = Resource.create(resource_attrs) provider = TracerProvider(resource=resource) + from .conversation import ConversationIdSpanProcessor + + provider.add_span_processor(ConversationIdSpanProcessor()) if exporter: provider.add_span_processor(BatchSpanProcessor(exporter)) diff --git a/packages/client/tests/test_conversation.py b/packages/client/tests/test_conversation.py new file mode 100644 index 0000000..5bd8a57 --- /dev/null +++ b/packages/client/tests/test_conversation.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from launchdarkly_ai_server.conversation import ( + GEN_AI_CONVERSATION_ID, + ConversationIdSpanProcessor, + conversation_id, + set_conversation_id_if_absent, +) + +_exporter = InMemorySpanExporter() +_provider = TracerProvider() +_provider.add_span_processor(ConversationIdSpanProcessor()) +_provider.add_span_processor(SimpleSpanProcessor(_exporter)) +_tracer = _provider.get_tracer("@launchdarkly/ai-server") + + +@pytest.fixture(autouse=True) +def _reset_exporter() -> Iterator[None]: + _exporter.clear() + yield + _exporter.clear() + + +def finished() -> list[ReadableSpan]: + return list(_exporter.get_finished_spans()) + + +class TestConversationId: + def test_stamps_id_on_every_span_in_scope(self) -> None: + with conversation_id("thread-123"): + root = _tracer.start_span("invoke_agent") + child = _tracer.start_span( + "chat gpt-4o", context=trace.set_span_in_context(root) + ) + child.end() + root.end() + spans = finished() + assert len(spans) == 2 + for span in spans: + assert span.attributes[GEN_AI_CONVERSATION_ID] == "thread-123" + + def test_writes_nothing_when_unbound(self) -> None: + span = _tracer.start_span("invoke_agent") + span.end() + assert GEN_AI_CONVERSATION_ID not in finished()[0].attributes + + def test_whitespace_id_is_unbound(self) -> None: + with conversation_id(" "): + span = _tracer.start_span("invoke_agent") + span.end() + assert GEN_AI_CONVERSATION_ID not in finished()[0].attributes + + def test_does_not_invent_an_id_from_the_trace_id(self) -> None: + span = _tracer.start_span("invoke_agent") + span.end() + recorded = finished()[0] + assert GEN_AI_CONVERSATION_ID not in recorded.attributes + assert recorded.context.trace_id != 0 + + +class TestSetConversationIdIfAbsent: + def test_leaves_caller_id_in_place(self) -> None: + with conversation_id("caller-id"): + span = _tracer.start_span("invoke_agent") + set_conversation_id_if_absent(span, "sess-abc") + span.end() + assert finished()[0].attributes[GEN_AI_CONVERSATION_ID] == "caller-id" + + def test_writes_session_id_when_unbound(self) -> None: + span = _tracer.start_span("invoke_agent") + set_conversation_id_if_absent(span, "sess-abc") + span.end() + assert finished()[0].attributes[GEN_AI_CONVERSATION_ID] == "sess-abc" + + +class TestProcessorScope: + """The processor is registered on the *global* provider, so it sees every span in the process. + + Stamping a caller-supplied conversation id onto unrelated telemetry — Postgres queries, inbound + HTTP server spans, and the outbound provider call itself — is both semantically wrong and the + leak the module docstring says this design avoids. + """ + + def test_stamps_spans_from_launchdarkly_tracers(self) -> None: + tracer = _provider.get_tracer("@launchdarkly/ai-claude-messages") + with conversation_id("thread-123"): + span = tracer.start_span("invoke_agent") + span.end() + assert finished()[0].attributes[GEN_AI_CONVERSATION_ID] == "thread-123" + + def test_does_not_stamp_third_party_instrumentation_spans(self) -> None: + http = _provider.get_tracer("opentelemetry.instrumentation.httpx") + db = _provider.get_tracer("opentelemetry.instrumentation.psycopg") + with conversation_id("thread-123"): + for tracer, name in ((http, "POST /v1/messages"), (db, "SELECT users")): + span = tracer.start_span(name) + span.end() + for span in finished(): + assert (span.attributes or {}).get(GEN_AI_CONVERSATION_ID) is None, ( + span.name + ) diff --git a/packages/client/tests/test_stream_conversation.py b/packages/client/tests/test_stream_conversation.py new file mode 100644 index 0000000..198f993 --- /dev/null +++ b/packages/client/tests/test_stream_conversation.py @@ -0,0 +1,266 @@ +"""Conversation id across ``config().stream()``. + +An ``async def`` with ``yield`` does not run its body until the first ``__anext__``, which for a +streaming caller is normally after the ``conversation_id`` block has already exited. These tests +pin the id onto spans a handler opens while streaming, and pin that binding no id leaves span +parenting alone. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, Iterator +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +import launchdarkly_ai_server.lifecycle as lifecycle_module +from launchdarkly_ai_server import ProviderHandler, config +from launchdarkly_ai_server.conversation import ( + GEN_AI_CONVERSATION_ID, + ConversationIdSpanProcessor, + bind_conversation_id, + conversation_id, +) + +CONTEXT = {"kind": "user", "key": "u1"} + +_exporter = InMemorySpanExporter() +_provider = TracerProvider() +_provider.add_span_processor(ConversationIdSpanProcessor()) +_provider.add_span_processor(SimpleSpanProcessor(_exporter)) +_tracer = _provider.get_tracer("@launchdarkly/ai-server") + + +@pytest.fixture(autouse=True) +def _reset_exporter() -> Iterator[None]: + _exporter.clear() + yield + _exporter.clear() + + +@pytest.fixture +def mock_ld_client() -> Iterator[MagicMock]: + c = MagicMock() + c.track = MagicMock() + c.flush = AsyncMock() + c.close = AsyncMock() + c.variation = AsyncMock( + return_value={ + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": "Be helpful.", + "_ldMeta": { + "enabled": True, + "variationKey": "v1", + "version": 1, + "mode": "messages", + }, + } + ) + lifecycle_module._set_client_for_testing(c) + yield c + lifecycle_module._reset_for_testing() + + +def _span_creating_handler(chunks: list[str] | None = None) -> ProviderHandler: + """Mirrors a real streaming handler: spans are opened as the generator runs, not at call time.""" + _chunks = chunks or ["Hello", " World"] + _usage = {"input_tokens": 10, "output_tokens": 5} + + async def fn(cfg, user_input, tool_handlers, variables, history=None) -> dict: # type: ignore[override] + return {"output": "".join(_chunks), "usage": _usage} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator: # type: ignore[override] + root = _tracer.start_span("invoke_agent") + for c in _chunks: + with trace.use_span(root, end_on_exit=False): + chat = _tracer.start_span("chat gpt-4") + chat.end() + yield {"type": "chunk", "text": c} + root.end() + yield {"type": "done", "output": "".join(_chunks), "usage": _usage} + + return ProviderHandler( + fn=fn, provides_for=("TestProvider", "messages"), stream_fn=stream_fn + ) + + +def finished() -> list[ReadableSpan]: + return list(_exporter.get_finished_spans()) + + +def _ids_on(prefixes: tuple[str, ...]) -> list[Any]: + return [ + s.attributes.get(GEN_AI_CONVERSATION_ID) if s.attributes else None + for s in finished() + if any(s.name.startswith(p) for p in prefixes) + ] + + +class TestStreamConversationId: + async def test_stamps_id_when_iterated_outside_the_block( + self, mock_ld_client: MagicMock + ) -> None: + m = config(key="flag", handler=_span_creating_handler()) + + # The natural shape: bind, build the generator, hand it off, iterate later. + with conversation_id("thread-123"): + gen = m.stream("q", CONTEXT) + async for _ in gen: + pass + + ids = _ids_on(("invoke_agent", "chat")) + assert len(ids) > 0 + assert all(i == "thread-123" for i in ids) + + async def test_stamps_id_when_iterated_inside_the_block( + self, mock_ld_client: MagicMock + ) -> None: + m = config(key="flag", handler=_span_creating_handler()) + + with conversation_id("thread-inside"): + async for _ in m.stream("q", CONTEXT): + pass + + ids = _ids_on(("invoke_agent", "chat")) + assert len(ids) > 0 + assert all(i == "thread-inside" for i in ids) + + async def test_leaves_parenting_unchanged_when_unbound( + self, mock_ld_client: MagicMock + ) -> None: + m = config(key="flag", handler=_span_creating_handler(["only"])) + + caller = _tracer.start_span("caller") + with trace.use_span(caller, end_on_exit=False): + async for _ in m.stream("q", CONTEXT): + pass + caller.end() + + root = next(s for s in finished() if s.name == "invoke_agent") + caller_span = next(s for s in finished() if s.name == "caller") + assert root.parent is not None + assert root.parent.span_id == caller_span.context.span_id + assert (root.attributes or {}).get(GEN_AI_CONVERSATION_ID) is None + + async def test_overlapping_streams_stay_isolated( + self, mock_ld_client: MagicMock + ) -> None: + async def run(tag: str) -> None: + m = config(key="flag", handler=_span_creating_handler([tag])) + with conversation_id(tag): + gen = m.stream("q", CONTEXT) + async for _ in gen: + pass + + await asyncio.gather(run("tenant-a"), run("tenant-b")) + + roots = [s for s in finished() if s.name == "invoke_agent"] + assert len(roots) == 2 + assert sorted( + (s.attributes or {}).get(GEN_AI_CONVERSATION_ID) for s in roots + ) == ["tenant-a", "tenant-b"] + + for root in roots: + expected = (root.attributes or {}).get(GEN_AI_CONVERSATION_ID) + same_trace = [ + s for s in finished() if s.context.trace_id == root.context.trace_id + ] + assert all( + (s.attributes or {}).get(GEN_AI_CONVERSATION_ID) == expected + for s in same_trace + ) + + +class TestConcurrentConversationIsolation: + async def test_two_overlapping_scopes_stay_isolated(self) -> None: + async def one(tag: str) -> None: + with conversation_id(tag): + span = _tracer.start_span(f"{tag}:invoke_agent") + await asyncio.sleep(0.01) + span.end() + + await asyncio.gather(one("tenant-a"), one("tenant-b")) + + by_name = { + s.name: (s.attributes or {}).get(GEN_AI_CONVERSATION_ID) for s in finished() + } + assert by_name["tenant-a:invoke_agent"] == "tenant-a" + assert by_name["tenant-b:invoke_agent"] == "tenant-b" + + +class TestStreamParentingWithIdBound: + """The unbound parenting test cannot catch a regression in the wrapper — it never builds one. + + A real streaming handler holds a span current across a ``yield``. The wrapper must not disturb + that: a span the handler opens after resuming belongs to its own ``chat`` span, exactly as it + would with no id bound. + """ + + async def _run(self, bind: bool) -> tuple[Any, Any, Any]: + _exporter.clear() + + async def handler() -> AsyncGenerator: + with trace.use_span(_tracer.start_span("chat"), end_on_exit=True): + yield "c1" + child = _tracer.start_span("child-after-resume") + child.end() + yield "c2" + + caller = _tracer.start_span("caller") + with trace.use_span(caller, end_on_exit=False): + if bind: + with conversation_id("thread-x"): + gen = bind_conversation_id(handler()) + else: + gen = handler() + async for _ in gen: + pass + caller.end() + + spans = {s.name: s for s in finished()} + return caller, spans.get("chat"), spans.get("child-after-resume") + + async def test_child_after_resume_keeps_its_parent_when_id_is_bound(self) -> None: + _, chat, child = await self._run(bind=True) + assert chat is not None and child is not None + assert child.parent is not None + assert child.parent.span_id == chat.context.span_id + + async def test_parenting_matches_the_unbound_control(self) -> None: + caller_u, chat_u, child_u = await self._run(bind=False) + unbound = ( + chat_u.parent.span_id == caller_u.context.span_id, + child_u.parent.span_id == chat_u.context.span_id, + ) + caller_b, chat_b, child_b = await self._run(bind=True) + bound = ( + chat_b.parent.span_id == caller_b.context.span_id, + child_b.parent.span_id == chat_b.context.span_id, + ) + assert bound == unbound + + async def test_binding_does_not_leak_into_the_caller_context(self) -> None: + from opentelemetry import context as otel_context + + from launchdarkly_ai_server.conversation import _CONV_KEY + + await self._run(bind=True) + assert otel_context.get_value(_CONV_KEY) is None + + +class TestNullishConversationId: + async def test_none_id_is_treated_as_unbound_rather_than_raising(self) -> None: + # The natural call site is an optional header: conversation_id(request.thread_id) + with conversation_id(None): # type: ignore[arg-type] + span = _tracer.start_span("invoke_agent") + span.end() + assert (finished()[0].attributes or {}).get(GEN_AI_CONVERSATION_ID) is None