From 8c8604fd576c4f82c689c12e436a086f32f48ae3 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 16:12:01 -0400 Subject: [PATCH 1/2] feat(claude-agents)!: emit invoke_agent, chat and execute_tool spans One flat span named claude.query becomes the tree the TypeScript SDK emits: an invoke_agent root, one `chat {model}` child per inference the Agent SDK reports, one `execute_tool {name}` child per tool call. BREAKING CHANGE: the span this handler emits is renamed from `claude.query` and `claude.query.stream` to `invoke_agent`. Queries selecting on the old names will not match. Prompt and completion content is no longer on spans unless the caller passes capture_content=True. Subagent identity was invisible. The Agent SDK reports a request id, a session id and a subagent type for every inference, and none of it reached a span, so a trace could not say which subagent ran or group turns by conversation. The chat span now carries gen_ai.response.id and gen_ai.agent.name, and gen_ai.conversation.id goes on all three span types: the root learns it from the CLI's init message, the chat span from the inference, and the tool span from the hook input, which is where this side sees it without waiting for a message. Cache tokens now reach the span, folded into the input total, because Anthropic reports cache reads and writes beside the input count rather than inside it. The finish reason is written when the SDK reports one, which in practice is almost never: measured against Agent SDK 0.3.220, stop_reason is null on every assistant message and only the run-level result carries one. Deriving a reason from the presence of a tool-use block would put a value on the span the provider never returned, so the write stays guarded and usually absent. The streaming path needed more than a finally. The blocking path already held the vendor's query generator in a variable and awaited aclose() on it, because a bare return inside `async for` abandons it and asyncio's finalizer then raises RuntimeError when it is suspended inside a real await in the SDK. The streaming path iterated the generator inline with no held reference and had no such cleanup, so it carried the same bug the blocking path was patched for. Both now close the generator in the same finally that ends the spans, and the streaming path gets the test the blocking path already had. A run abandoned mid-stream now closes any tool span whose PostToolUse hook never fired, which is otherwise the one span with no path to being ended. Tests: 56 to 90. The telemetry tests are rewritten rather than extended, because they pinned the old flat span, and every test that was not about telemetry is preserved under its original name. --- .../launchdarkly_ai_claude_agents/handler.py | 693 ++++-- .../launchdarkly_ai_claude_agents/spans.py | 571 +++++ packages/claude-agents/tests/test_handler.py | 2033 +++++++++++------ 3 files changed, 2374 insertions(+), 923 deletions(-) create mode 100644 packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py 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 d98c831..2313144 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py @@ -1,43 +1,69 @@ """ -Claude Agents handler — uses the claude-agent-sdk `query()` + MCP tools. -Mirrors the TypeScript @launchdarkly/ai-claude-agents handler. +Claude Agents handler: uses the claude-agent-sdk ``query()`` plus MCP tools. + +Span shape: ``invoke_agent`` root, one ``chat`` child per model response, ``execute_tool`` children +per tool call, the same three-span vocabulary the other five handlers emit. + +``query()`` reports no request boundaries, so this handler derives them from the message stream: an +``AssistantMessage`` *is* one model response, carrying its own ``message_id`` (the Anthropic response +id), ``usage`` and ``model``. See :mod:`spans` for how those are folded into ``chat`` spans. + +Deliberately NOT done here: enabling the CLI's own OTel exporter. Its spans are named outside the +semantic conventions and duplicate what this handler already emits. """ from __future__ import annotations import asyncio import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterator, Callable from typing import Any +from claude_agent_sdk import ( + ClaudeAgentOptions, + HookMatcher, + ResultMessage, + StreamEvent, + query, +) + from launchdarkly_ai_server import ( NATIVE_TOOL_KEY, AiConfigRep, LDContext, NativeTool, ProviderHandler, + SpanMessage, + SpanMessagePart, config, create_handler, + end_span_once, parse_template, - set_ld_span_attributes, - set_openllmetry_completion, - set_openllmetry_prompt, + set_input_content_attributes, + set_output_content_attributes, + set_tool_call_content_attributes, ) -try: - from opentelemetry import trace - from opentelemetry.trace import StatusCode as SpanStatusCode - - _HAS_OTEL = True -except ImportError: - _HAS_OTEL = False - -TOOL_MCP_NAME = "tool-mcp" -MCP_TOOL_PREFIX = f"mcp__{TOOL_MCP_NAME}__" - +from .spans import ( + MCP_TOOL_PREFIX, + TOOL_MCP_NAME, + InferenceSpans, + Opening, + ToolCatalog, + fail_span, + finish_root_span, + mark_ok, + parent_context_of, + record_conversation_id, + record_native_tools, + start_root_span, + start_tool_span, + succeed_span, + tool_display_name, +) # --------------------------------------------------------------------------- -# Helpers +# Tool wiring # --------------------------------------------------------------------------- @@ -45,7 +71,13 @@ async def build_tool_mcp( config_tools: dict[str, Any], handlers: dict[str, Any], ) -> Any: - """Build an in-process SDK MCP server from LD config tools + handler functions.""" + """Build an in-process SDK MCP server from LD config tools + handler functions. + + Imports the SDK lazily rather than off the module-level ``claude_agent_sdk`` import that + ``query``/``ClaudeAgentOptions``/etc. use: ``native_graph.py`` (out of scope for this telemetry + pass) calls this function too, and its own tests mock the SDK by patching + ``importlib.import_module`` rather than this module's names. + """ import importlib sdk = importlib.import_module("claude_agent_sdk") @@ -79,9 +111,13 @@ def partition_tools( """ Returns (native_tool_map, user_config_tools, native_tool_names). - native_tool_map : provider tool name → tracking stub (for PreToolUse hook) + native_tool_map : provider tool name → tracking stub (for PreToolUse/PostToolUse hooks) user_config_tools: LD tool definitions for user-defined tools (sent via MCP) native_tool_names: provider-facing names for query(options.tools=[...]) + + Kept at a 3-tuple return for backward compatibility: ``native_graph.py`` (out of scope for this + telemetry pass) unpacks this directly. See :func:`_native_tool_aliases` for the fourth mapping + the span work needs. """ native_tool_map: dict[str, Any] = {} user_config_tools: dict[str, Any] = {} @@ -96,6 +132,144 @@ def partition_tools( return native_tool_map, user_config_tools, list(native_tool_map.keys()) +def _native_tool_aliases(tool_handlers: dict[str, Any]) -> dict[str, str]: + """AI Config key → provider tool name, for the config's native tools only. + + ``partition_tools``'s ``native_tool_map`` is keyed the other way round, by provider name, which + loses the AI Config key, and the key is what a config's declared schema is filed under. + :class:`~.spans.ToolCatalog` needs both ends to report one tool once, under the name the model + saw, with the schema the config gave it. + """ + aliases: dict[str, str] = {} + for ld_name, stub in tool_handlers.items(): + native = getattr(stub, NATIVE_TOOL_KEY, None) + if isinstance(native, NativeTool): + aliases[ld_name] = native.tool_name + return aliases + + +def _is_coroutine(fn: Any) -> bool: + return asyncio.iscoroutinefunction(fn) + + +def _build_hooks(native_tool_map: dict[str, Any]) -> dict[str, Any] | None: + """The pre-span-work hook set: tracks a native tool call for telemetry purposes only. + + Kept for ``native_graph.py`` (out of scope for this telemetry pass), which imports this name + directly and does not build ``execute_tool`` spans of its own. See :func:`build_tool_hooks` for + the span-aware hook set this handler's own ``_call_impl``/``_stream_gen`` use. + """ + if not native_tool_map: + return None + + async def _pre_tool_hook( + input_data: Any, tool_use_id: str | None, hook_context: Any + ) -> dict[str, Any]: + tool_name = ( + input_data.get("tool_name") if isinstance(input_data, dict) else None + ) + stub = native_tool_map.get(tool_name) if tool_name is not None else None + if stub and callable(stub): + stub() + return {} + + return {"PreToolUse": [HookMatcher(hooks=[_pre_tool_hook])]} # type: ignore[list-item] + + +ToolTelemetry = Callable[[BaseException], None] + + +def build_tool_hooks( + native_tool_map: dict[str, Any], + parent_context: Any, + capture_content: bool, +) -> tuple[dict[str, list[HookMatcher]], ToolTelemetry]: + """Builds the PreToolUse/PostToolUse/PostToolUseFailure hooks that open and close + ``execute_tool`` spans around the Agent SDK's own tool dispatch. + + Returns ``(hooks, close_open_spans)``. ``close_open_spans`` fails every span this run still has + open, for the path where the SDK throws mid-tool-call and no ``PostToolUse*`` hook ever fires. + """ + tool_spans: dict[str, Any] = {} + + def _finish( + tool_use_id: str | None, *, error: BaseException | None, result: Any = None + ) -> None: + if tool_use_id is None: + return + span = tool_spans.pop(tool_use_id, None) + if span is None: + return + if result is not None: + set_tool_call_content_attributes(span, capture_content, result=result) + if error is None: + succeed_span(span) + else: + fail_span(span, error) + + async def _pre_tool_use( + input_data: Any, tool_use_id: str | None, hook_context: Any + ) -> dict[str, Any]: + if input_data.get("hook_event_name") != "PreToolUse": + return {} + provider_name = input_data.get("tool_name", "") + stub = native_tool_map.get(provider_name) + if stub and callable(stub): + stub() + + display_name = tool_display_name(provider_name) + use_id = input_data.get("tool_use_id") or tool_use_id or "" + span = start_tool_span(display_name, use_id, parent_context) + session_id = input_data.get("session_id") + # 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_tool_call_content_attributes( + span, capture_content, arguments=input_data.get("tool_input") + ) + tool_spans[use_id] = span + return {} + + async def _post_tool_use( + input_data: Any, tool_use_id: str | None, hook_context: Any + ) -> dict[str, Any]: + if input_data.get("hook_event_name") != "PostToolUse": + return {} + use_id = input_data.get("tool_use_id") or tool_use_id + _finish(use_id, error=None, result=input_data.get("tool_response")) + return {} + + async def _post_tool_use_failure( + input_data: Any, tool_use_id: str | None, hook_context: Any + ) -> dict[str, Any]: + if input_data.get("hook_event_name") != "PostToolUseFailure": + return {} + use_id = input_data.get("tool_use_id") or tool_use_id + _finish( + use_id, + error=RuntimeError(str(input_data.get("error") or "tool call failed")), + ) + return {} + + def close_open_spans(error: BaseException) -> None: + for span in list(tool_spans.values()): + fail_span(span, error) + tool_spans.clear() + + hooks: dict[str, list[HookMatcher]] = { + "PreToolUse": [HookMatcher(hooks=[_pre_tool_use])], # type: ignore[list-item] + "PostToolUse": [HookMatcher(hooks=[_post_tool_use])], # type: ignore[list-item] + "PostToolUseFailure": [HookMatcher(hooks=[_post_tool_use_failure])], # type: ignore[list-item] + } + return hooks, close_open_spans + + +# --------------------------------------------------------------------------- +# Prompt construction +# --------------------------------------------------------------------------- + + def _format_history(history: list[dict[str, Any]] | None) -> str | None: if not history: return None @@ -113,7 +287,11 @@ def build_prompt( variables: dict[str, Any], history: list[dict[str, Any]] | None = None, ) -> tuple[str, str | None]: - """Returns (prompt, system_prompt).""" + """Returns (prompt, system_prompt). + + One user message, not one per configured role: ``query()`` takes a single prompt string, so this + really does flatten a configured history into one turn before the model sees it. + """ safe_input = user_input or "" system_prompt: str | None = None @@ -143,33 +321,49 @@ def build_prompt( return safe_input, system_prompt -def _is_coroutine(fn: Any) -> bool: - return asyncio.iscoroutinefunction(fn) - +def _opening_of(prompt: str, system_prompt: str | None) -> Opening: + return Opening( + system_instructions=system_prompt, + messages=[ + SpanMessage( + role="user", parts=[SpanMessagePart(type="text", content=prompt)] + ) + ], + ) -def _build_hooks(native_tool_map: dict[str, Any]) -> dict[str, Any] | None: - if not native_tool_map: - return None - import importlib +def _result_error(subtype: str, errors: list[str] | None) -> str: + """Builds the error for a non-success result message. - sdk = importlib.import_module("claude_agent_sdk") - HookMatcher = sdk.HookMatcher + ``SDKResultError`` carries no ``result`` field, so a run that hit ``error_max_turns`` or + ``error_max_budget_usd`` genuinely failed and has to surface as a failure rather than being + reported OK with zeroed usage. + """ + detail = f": {'; '.join(errors)}" if errors else "" + return f"Claude agent run ended with {subtype}{detail}" - async def _pre_tool_hook( - input_data: Any, tool_use_id: str | None, context: Any - ) -> dict[str, Any]: - tool_name = getattr(input_data, "tool_name", None) or ( - input_data.get("tool_name") if isinstance(input_data, dict) else None - ) - stub = native_tool_map.get(tool_name) if tool_name is not None else None - if stub and callable(stub): - stub() - return {} - return { - "PreToolUse": [HookMatcher(hooks=[_pre_tool_hook])], +def _build_query_options( + config: AiConfigRep, + system_prompt: str | None, + native_tool_names: list[str], + mcp_allowed_tools: list[str], + tool_mcp: Any, + hooks: dict[str, list[HookMatcher]] | None, + **extra: Any, +) -> ClaudeAgentOptions: + all_allowed = [*mcp_allowed_tools, *native_tool_names] + kwargs: dict[str, Any] = { + "model": config["model"]["name"], + "tools": native_tool_names if native_tool_names else [], + "allowed_tools": all_allowed if all_allowed else [], + "mcp_servers": {TOOL_MCP_NAME: tool_mcp} if tool_mcp else {}, + "hooks": hooks or {}, + **extra, } + if system_prompt: + kwargs["system_prompt"] = system_prompt + return ClaudeAgentOptions(**kwargs) # --------------------------------------------------------------------------- @@ -177,9 +371,12 @@ async def _pre_tool_hook( # --------------------------------------------------------------------------- -def create_claude_agents_handler() -> ProviderHandler: - """Creates a ``ProviderHandler`` for Anthropic's Claude via the claude-agent-sdk.""" - tracer_name = "@launchdarkly/ai-claude-agents" +def create_claude_agents_handler(*, capture_content: bool = False) -> ProviderHandler: + """Creates a ``ProviderHandler`` for Anthropic's Claude via the claude-agent-sdk. + + Set *capture_content* to put prompts, model output, tool arguments and tool results on the + emitted spans. It defaults to off. See TELEMETRY-CONTRACT.md section 7. + """ async def _call_impl( config: AiConfigRep, @@ -188,134 +385,129 @@ async def _call_impl( variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: - import importlib - - sdk = importlib.import_module("claude_agent_sdk") - ClaudeAgentOptions = sdk.ClaudeAgentOptions - ResultMessage = sdk.ResultMessage - query_fn = sdk.query - th = tool_handlers or {} vs = variables or {} - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("claude.query") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", "anthropic") - span.set_attribute( - "gen_ai.request.model", config.get("model", {}).get("name", "") - ) - set_ld_span_attributes(span, vs) - else: - span = None + span = start_root_span(config, vs) + parent = parent_context_of(span) prompt, system_prompt = build_prompt(config, user_input, vs, history) - - # Append outputFormat instruction to system prompt if config.get("outputFormat"): schema_instr = f"Respond with valid JSON matching this schema:\n{json.dumps(config['outputFormat'])}" system_prompt = ( f"{system_prompt}\n\n{schema_instr}" if system_prompt else schema_instr ) + opening = _opening_of(prompt, system_prompt) - if span: - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt}) - prompt_msgs: list[dict[str, str]] = [] - if system_prompt: - prompt_msgs.append({"role": "system", "content": system_prompt}) - prompt_msgs.append({"role": "user", "content": prompt}) - set_openllmetry_prompt(span, prompt_msgs) + native_tool_map, user_config_tools, native_tool_names = partition_tools( + config.get("tools"), th + ) + catalog = ToolCatalog(config.get("tools"), _native_tool_aliases(th)) + set_input_content_attributes( + span, + capture_content, + system_instructions=opening.system_instructions, + messages=opening.messages, + tool_definitions=catalog.current, + ) + # Declared out here so the except clause can end a chat span the throw left open. + inference = InferenceSpans(config, parent, capture_content, catalog, opening) + tool_telemetry: ToolTelemetry | None = None + # Set wherever the root's usage is written, so the failure path can tell whether the CLI + # already reported an authoritative run-level total and must not overwrite it. + root_usage_written = False + gen: AsyncIterator[Any] | None = None try: - native_tool_map, user_config_tools, native_tool_names = partition_tools( - config.get("tools"), th - ) - tool_mcp = ( await build_tool_mcp(user_config_tools, th) if user_config_tools else None ) - mcp_allowed = [MCP_TOOL_PREFIX + n for n in user_config_tools] - all_allowed = mcp_allowed + native_tool_names - - hooks = _build_hooks(native_tool_map) - - options = ClaudeAgentOptions( - allowed_tools=all_allowed if all_allowed else [], - mcp_servers={TOOL_MCP_NAME: tool_mcp} if tool_mcp else {}, - hooks=hooks or {}, - **({"system_prompt": system_prompt} if system_prompt else {}), - **({"tools": native_tool_names} if native_tool_names else {}), + mcp_allowed_tools = [MCP_TOOL_PREFIX + n for n in user_config_tools] + hooks = None + if native_tool_names or mcp_allowed_tools: + hooks, tool_telemetry = build_tool_hooks( + native_tool_map, parent, capture_content + ) + + options = _build_query_options( + config, + system_prompt, + native_tool_names, + mcp_allowed_tools, + tool_mcp, + hooks, ) output = "" - raw_usage: dict[str, Any] = {} - span_ended = False - - # Hold an explicit reference so we can call aclose() in the finally - # block below. A bare `return` inside `async for` abandons the - # generator — Python's asyncio finalizer later tries to aclose() it - # and raises RuntimeError if the generator is suspended inside a real - # await in the SDK (AIC-2950). See Appendix A.4 in TESTING.md. - gen = query_fn(prompt=prompt, options=options) + + # Held in a variable so the finally below can aclose() it. A bare `return` inside + # `async for` abandons the generator, and asyncio's finalizer then raises RuntimeError + # when the generator is suspended inside a real await in the SDK. + gen = query(prompt=prompt, options=options) try: async for message in gen: + record_conversation_id(span, message) + record_native_tools(span, message, capture_content, catalog) + inference.record(message) + if isinstance(message, ResultMessage): - output = message.result or "" - raw_usage = message.usage or {} - input_tokens = int(raw_usage.get("input_tokens", 0)) - output_tokens = int(raw_usage.get("output_tokens", 0)) - if span: - span.set_attribute( - "gen_ai.response.model", - config.get("model", {}).get("name", ""), - ) - span.set_attribute( - "gen_ai.usage.input_tokens", input_tokens + # Before the root, so the children it parents are already closed. + inference.finish() + raw_usage: dict[str, Any] = dict(message.usage or {}) + finish_root_span(span, config, raw_usage) + # The CLI's own run-level total is authoritative, so the except clause below + # must not overwrite it with the summed-per-response figure. + root_usage_written = True + if message.subtype != "success": + raise RuntimeError( + _result_error(message.subtype, message.errors) ) - span.set_attribute( - "gen_ai.usage.output_tokens", output_tokens - ) - span.set_attribute( - "gen_ai.usage.total_tokens", - input_tokens + output_tokens, - ) - span.add_event( - "gen_ai.content.completion", - { - "gen_ai.completion": output - if isinstance(output, str) - else json.dumps(output) - }, - ) - set_openllmetry_completion( - span, - output - if isinstance(output, str) - else json.dumps(output), - { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - }, - ) - span.set_status(SpanStatusCode.OK) - span.end() - span_ended = True - break + result_text = ( + message.result + if isinstance(message.result, str) + else json.dumps(message.result) + ) + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[ + SpanMessagePart( + type="text", content=result_text + ) + ], + ) + ], + ) + succeed_span(span) + return {"output": message.result, "usage": raw_usage} + + # The stream ended without a result message, so no message closed the last + # response, and no message carried a run-level total either. The per-response sum + # is the only record of what the run spent. + inference.finish() + streamed_usage = inference.run_usage + finish_root_span(span, config, streamed_usage["total"]) + root_usage_written = True + succeed_span(span) + return {"output": output, "usage": streamed_usage["total"]} finally: - await gen.aclose() - - if span and not span_ended: - span.set_status(SpanStatusCode.OK) - span.end() - return {"output": output, "usage": raw_usage} + if gen is not None: + await gen.aclose() # type: ignore[attr-defined] except Exception as exc: - if span: - span.record_exception(exc) - span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() + # Before the root: the exporter only receives ended spans, so a chat span left open by + # the throw would never reach the trace. + inference.finish() + if tool_telemetry is not None: + tool_telemetry(exc) + if not root_usage_written and inference.run_usage["reported"]: + finish_root_span(span, config, inference.run_usage["total"]) + fail_span(span, exc) raise def _stream_impl( @@ -326,7 +518,12 @@ def _stream_impl( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: return _stream_gen( - config, user_input, tool_handlers or {}, variables or {}, history + config, + user_input, + tool_handlers or {}, + variables or {}, + history, + capture_content=capture_content, ) return create_handler(("Anthropic", "agent"), _call_impl, _stream_impl) # type: ignore[arg-type] @@ -338,61 +535,72 @@ async def _stream_gen( tool_handlers: dict[str, Any], variables: dict[str, Any], history: list[dict[str, Any]] | None = None, + *, + capture_content: bool = False, ) -> AsyncGenerator[dict[str, Any], None]: - import importlib - - sdk = importlib.import_module("claude_agent_sdk") - ClaudeAgentOptions = sdk.ClaudeAgentOptions - ResultMessage = sdk.ResultMessage - StreamEvent = sdk.StreamEvent - query_fn = sdk.query - - tracer_name = "@launchdarkly/ai-claude-agents" - if _HAS_OTEL: - span = trace.get_tracer(tracer_name).start_span("claude.query.stream") - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", "anthropic") - span.set_attribute( - "gen_ai.request.model", config.get("model", {}).get("name", "") - ) - set_ld_span_attributes(span, variables) - else: - span = None + """Streams the run, emitting the same span tree as the blocking path. + + A consumer that breaks out of ``async for``, or raises inside the loop body, makes this + generator run its ``finally`` without ever entering ``except``: ``GeneratorExit`` inherits from + ``BaseException``, so ``except Exception`` does not see it. The same ``finally`` also closes the + vendor's own generator, for the same reason the blocking path holds one in a variable: a bare + exit abandons it, and asyncio's finalizer later raises ``RuntimeError`` when it is suspended + inside a real await in the SDK. + """ + span = start_root_span(config, variables) + parent = parent_context_of(span) prompt, system_prompt = build_prompt(config, user_input, variables, history) - if span: - span.add_event("gen_ai.content.prompt", {"gen_ai.prompt": prompt}) - prompt_msgs: list[dict[str, str]] = [] - if system_prompt: - prompt_msgs.append({"role": "system", "content": system_prompt}) - prompt_msgs.append({"role": "user", "content": prompt}) - set_openllmetry_prompt(span, prompt_msgs) + opening = _opening_of(prompt, system_prompt) + + native_tool_map, user_config_tools, native_tool_names = partition_tools( + config.get("tools"), tool_handlers + ) + catalog = ToolCatalog(config.get("tools"), _native_tool_aliases(tool_handlers)) + set_input_content_attributes( + span, + capture_content, + system_instructions=opening.system_instructions, + messages=opening.messages, + tool_definitions=catalog.current, + ) + + inference = InferenceSpans(config, parent, capture_content, catalog, opening) + tool_telemetry: ToolTelemetry | None = None + root_usage_written = False + ended: set[int] = set() + gen: AsyncIterator[Any] | None = None try: - native_tool_map, user_config_tools, native_tool_names = partition_tools( - config.get("tools"), tool_handlers - ) tool_mcp = ( await build_tool_mcp(user_config_tools, tool_handlers) if user_config_tools else None ) - mcp_allowed = [MCP_TOOL_PREFIX + n for n in user_config_tools] - all_allowed = mcp_allowed + native_tool_names - hooks = _build_hooks(native_tool_map) - - options = ClaudeAgentOptions( - allowed_tools=all_allowed if all_allowed else [], - mcp_servers={TOOL_MCP_NAME: tool_mcp} if tool_mcp else {}, - hooks=hooks or {}, + mcp_allowed_tools = [MCP_TOOL_PREFIX + n for n in user_config_tools] + hooks = None + if native_tool_names or mcp_allowed_tools: + hooks, tool_telemetry = build_tool_hooks( + native_tool_map, parent, capture_content + ) + + options = _build_query_options( + config, + system_prompt, + native_tool_names, + mcp_allowed_tools, + tool_mcp, + hooks, include_partial_messages=True, - **({"system_prompt": system_prompt} if system_prompt else {}), - **({"tools": native_tool_names} if native_tool_names else {}), ) full_output = "" + gen = query(prompt=prompt, options=options) + async for message in gen: + record_conversation_id(span, message) + record_native_tools(span, message, capture_content, catalog) + inference.record(message) - async for message in query_fn(prompt=prompt, options=options): if isinstance(message, StreamEvent): event = message.event if ( @@ -404,50 +612,82 @@ async def _stream_gen( yield {"type": "chunk", "text": text} full_output += text elif isinstance(message, ResultMessage): - raw_usage = message.usage or {} - input_tokens = int(raw_usage.get("input_tokens", 0)) - output_tokens = int(raw_usage.get("output_tokens", 0)) - final_output = message.result or full_output - if span: - span.set_attribute( - "gen_ai.response.model", config.get("model", {}).get("name", "") - ) - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute( - "gen_ai.usage.total_tokens", input_tokens + output_tokens - ) - span.add_event( - "gen_ai.content.completion", - { - "gen_ai.completion": final_output - if isinstance(final_output, str) - else json.dumps(final_output) - }, - ) - set_openllmetry_completion( - span, - final_output - if isinstance(final_output, str) - else json.dumps(final_output), - {"input_tokens": input_tokens, "output_tokens": output_tokens}, - ) - span.set_status(SpanStatusCode.OK) - span.end() + # Before the root, so the children it parents are already closed. + inference.finish() + raw_usage: dict[str, Any] = dict(message.usage or {}) + finish_root_span(span, config, raw_usage) + root_usage_written = True + if message.subtype != "success": + raise RuntimeError(_result_error(message.subtype, message.errors)) + final_output = ( + message.result if message.result is not None else full_output + ) + result_text = ( + final_output + if isinstance(final_output, str) + else json.dumps(final_output) + ) + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=result_text)], + ) + ], + ) + mark_ok(span) + end_span_once(span, ended) yield {"type": "done", "output": final_output, "usage": raw_usage} return - if span: - span.set_status(SpanStatusCode.OK) - span.end() - yield {"type": "done", "output": full_output, "usage": {}} + # The stream ended without a result message, so nothing carried a run-level total. The + # per-response sum is the only record of the spend. + inference.finish() + streamed_usage = inference.run_usage + finish_root_span(span, config, streamed_usage["total"]) + root_usage_written = True + set_output_content_attributes( + span, + capture_content, + [ + SpanMessage( + role="assistant", + parts=[SpanMessagePart(type="text", content=full_output)], + ) + ], + ) + mark_ok(span) + end_span_once(span, ended) + yield {"type": "done", "output": full_output, "usage": streamed_usage["total"]} except Exception as exc: - if span: - span.record_exception(exc) - span.set_status(SpanStatusCode.ERROR, str(exc)) - span.end() + if tool_telemetry is not None: + tool_telemetry(exc) + if not root_usage_written and inference.run_usage["reported"]: + finish_root_span(span, config, inference.run_usage["total"]) + root_usage_written = True + fail_span(span, exc, ended) raise + finally: + # A no-op on the success path, where the result message already closed the last response. + # On the error and abandonment paths this is the only thing that ends a chat span still + # open, and the exporter only receives ended spans. + inference.finish() + # A no-op on the success and failure paths; on abandonment it is the only chance to close + # the tree, including any tool span whose PostToolUse hook never fired, and the only chance + # to report what the responses that did arrive cost. + if id(span) not in ended: + if tool_telemetry is not None: + tool_telemetry(RuntimeError("stream abandoned before completion")) + if not root_usage_written and inference.run_usage["reported"]: + finish_root_span(span, config, inference.run_usage["total"]) + end_span_once(span, ended, abandoned=True) + # Same reasoning as the blocking path: a bare exit through this generator's boundary + # abandons the vendor's own generator if it is not closed explicitly. + if gen is not None: + await gen.aclose() # type: ignore[attr-defined] def claude_agents( @@ -458,6 +698,9 @@ def claude_agents( ) -> Any: """Convenience wrapper: creates a handler and calls config(...).invoke().""" variables = kwargs.pop("variables", None) + capture_content = kwargs.pop("capture_content", False) return config( - key=config_key, handler=create_claude_agents_handler(), **kwargs + key=config_key, + handler=create_claude_agents_handler(capture_content=capture_content), + **kwargs, ).invoke(user_input, context, variables=variables) diff --git a/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py b/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py new file mode 100644 index 0000000..422b074 --- /dev/null +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/spans.py @@ -0,0 +1,571 @@ +"""Span construction for the Claude agents handler. + +Separate from ``handler.py`` for the same reason as the ``claude-messages`` package: the span shape +reads on its own, and the message-stream loop that drives it reads as a message loop rather than as +span bookkeeping with the SDK call in the middle. + +The shape is ``invoke_agent`` root, one ``chat {model}`` child per model turn, one +``execute_tool {name}`` child per tool call. Tool spans are siblings of the ``chat`` span, not +children of it: both take the same parent context, which is the root's. See TELEMETRY-CONTRACT.md +section 1. + +This handler differs from ``claude-messages`` in one structural way: ``query()`` reports no request +boundaries, and the Agent SDK's own message stream is the only source of truth for what happened. +:class:`InferenceSpans` derives a `chat` span per model turn from that stream (grouped on the +Anthropic response id, read off ``AssistantMessage.message_id``), the same way +``@launchdarkly/ai-claude-agents``'s ``InferenceSpans`` groups on ``request_id``. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Any + +from claude_agent_sdk import ( + AssistantMessage, + SystemMessage, + TextBlock, + ThinkingBlock, + ToolResultBlock, + ToolUseBlock, + UserMessage, +) + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanMessage, + SpanMessagePart, + ToolDefinitionInput, + add_cached_tokens_to_input, + end_span_once, + number_or_zero, + set_input_content_attributes, + set_ld_span_attributes, + set_model_identity_attributes, + set_output_content_attributes, + set_tool_definition_attributes, + set_usage_span_attributes, + to_semconv_finish_reason, +) + +try: + from opentelemetry import trace + from opentelemetry.trace import StatusCode as SpanStatusCode + + _HAS_OTEL = True +except ImportError: # pragma: no cover - exercised by the no-OTel install path + _HAS_OTEL = False + +TRACER_NAME = "@launchdarkly/ai-claude-agents" + +#: Anthropic serves every model behind this handler, so the provider name is a constant. +PROVIDER = "anthropic" + +TOOL_MCP_NAME = "tool-mcp" +MCP_TOOL_PREFIX = f"mcp__{TOOL_MCP_NAME}__" + + +def model_name(config: AiConfigRep) -> str: + return str(config.get("model", {}).get("name", "")) + + +def tool_display_name(provider_name: str) -> str: + """The name the model saw, with the MCP wrapper this handler adds stripped back off.""" + if provider_name.startswith(MCP_TOOL_PREFIX): + return provider_name[len(MCP_TOOL_PREFIX) :] + return provider_name + + +# ─── Span starts ───────────────────────────────────────────────────────────── + + +def start_root_span(config: AiConfigRep, variables: dict[str, Any]) -> Any: + """Opens the ``invoke_agent`` root and returns it, or ``None`` when OTel is absent. + + The root is the only span carrying ``launchdarkly.*`` and the ``feature_flag`` event, so it is + the span a config-scoped query finds. Child spans must not carry them. + """ + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span("invoke_agent") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + set_model_identity_attributes(span, PROVIDER, model_name(config)) + set_ld_span_attributes(span, variables) + return span + + +def parent_context_of(span: Any) -> Any: + """The context a child span should be parented to. + + Explicit rather than a bare current context: the current context only carries this span while a + context manager has attached it, and these handlers open a plain span rather than an active one, + so a host app that installs its own tracer provider would otherwise get a flat trace. + """ + if not _HAS_OTEL or span is None: + return None + return trace.set_span_in_context(span) + + +def start_tool_span(tool_name: str, tool_use_id: str, parent: Any) -> Any: + """Opens one ``execute_tool {name}`` span for one tool call.""" + if not _HAS_OTEL: + return None + span = trace.get_tracer(TRACER_NAME).start_span( + f"execute_tool {tool_name}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", tool_name) + span.set_attribute("gen_ai.tool.call.id", tool_use_id) + return span + + +# ─── Span finishes ─────────────────────────────────────────────────────────── + + +def finish_root_span(span: Any, config: AiConfigRep, raw_usage: dict[str, Any]) -> None: + """Writes the run-level identity and token totals onto the root. + + The Agent SDK's result message reports usage cumulatively for the whole run, which is exactly + what the root wants, and the root is the only span carrying ``launchdarkly.*`` and the + ``feature_flag`` event, so it is the span a config-scoped query finds. + + ``gen_ai.response.model`` is the *requested* name here, unlike on a ``chat`` span. See + TELEMETRY-CONTRACT.md section 2a: this handler is one of only two where the root and a `chat` + span disagree. + """ + if span is None: + return + span.set_attribute("gen_ai.response.model", model_name(config)) + set_usage_span_attributes(span, add_cached_tokens_to_input(raw_usage)) + + +def succeed_span(span: Any) -> None: + """Marks a span OK and ends it, for spans with nothing else to report.""" + if span is None: + return + span.set_status(SpanStatusCode.OK) + span.end() + + +def mark_ok(span: Any) -> None: + """Marks a span OK without ending it. + + The streaming path needs this: its ``finally`` owns every end, through ``end_span_once``, so a + success tail that ended the span itself would end it twice. + """ + if span is None: + return + span.set_status(SpanStatusCode.OK) + + +def fail_span(span: Any, error: BaseException, tracker: set[int] | None = None) -> None: + """Records the exception, sets ERROR, and ends the span. + + *tracker* is passed only from the streaming path, where a ``finally`` may race this to the same + span; elsewhere there is exactly one end and the tracker is unnecessary. + """ + if span is None: + return + span.record_exception(error) + span.set_status(SpanStatusCode.ERROR, str(error)) + if tracker is not None: + end_span_once(span, tracker) + else: + span.end() + + +def record_conversation_id(span: Any, message: Any) -> None: + """Copies the CLI's session id onto the root span as ``gen_ai.conversation.id``. + + 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. + """ + if ( + span is None + or not isinstance(message, SystemMessage) + or message.subtype != "init" + ): + return + session_id = message.data.get("session_id") + if session_id: + span.set_attribute("gen_ai.conversation.id", session_id) + + +def record_native_tools( + span: Any, message: Any, capture: bool, catalog: ToolCatalog +) -> None: + """Widens the root's tool catalog once the CLI names the tools it brought itself. + + The root's catalog is written before the run starts, so a run that dies before ``init`` still + reports the tools it was configured with. ``init`` is the first place the CLI's own tools become + visible, and rewriting the one attribute keeps the root and the ``chat`` spans from describing + the same catalog differently. + """ + if not isinstance(message, SystemMessage) or message.subtype != "init": + return + if catalog.widen(message.data.get("tools")) and span is not None: + set_tool_definition_attributes(span, capture, catalog.current) + + +def marks_local_work(message: Any) -> bool: + """Whether a non-assistant message marks the end of local work, and so the start of the window + that will produce the next response. + + Only two kinds do. A ``UserMessage`` carries the tool results the next call is being sent, and a + ``SystemMessage`` ``init`` opens the session. Everything else is progress reporting *during* a + call. + """ + if isinstance(message, UserMessage): + return True + return isinstance(message, SystemMessage) and message.subtype == "init" + + +# ─── Provider shapes as span shapes ────────────────────────────────────────── + + +def _attr(obj: Any, name: str) -> Any: + """Reads a field off a provider object or a plain dict, whichever the caller holds.""" + if isinstance(obj, dict): + return obj.get(name) + return getattr(obj, name, None) + + +def to_span_parts(content: Any) -> list[SpanMessagePart]: + """Converts one assistant or user message's content blocks into canonical span parts. + + Structural for dict-shaped blocks, typed for the Agent SDK's own dataclasses. Block kinds a span + has no part for (images, documents) are dropped rather than emitted malformed. + """ + if isinstance(content, str): + return [SpanMessagePart(type="text", content=content)] if content else [] + if not isinstance(content, list): + return [] + + parts: list[SpanMessagePart] = [] + for block in content: + if isinstance(block, TextBlock): + parts.append(SpanMessagePart(type="text", content=block.text or "")) + elif isinstance(block, ThinkingBlock): + parts.append( + SpanMessagePart(type="reasoning", content=block.thinking or "") + ) + elif isinstance(block, ToolUseBlock): + parts.append( + SpanMessagePart( + type="tool_call", + id=block.id, + name=block.name or "", + arguments=block.input, + ) + ) + elif isinstance(block, ToolResultBlock): + parts.append( + SpanMessagePart( + type="tool_call_response", + id=block.tool_use_id, + result=block.content, + ) + ) + elif isinstance(block, dict): + block_type = block.get("type") + if block_type == "text": + parts.append( + SpanMessagePart(type="text", content=str(block.get("text") or "")) + ) + elif block_type == "thinking": + parts.append( + SpanMessagePart( + type="reasoning", content=str(block.get("thinking") or "") + ) + ) + elif block_type == "tool_use": + block_id = block.get("id") + parts.append( + SpanMessagePart( + type="tool_call", + id=block_id if isinstance(block_id, str) else None, + name=str(block.get("name") or ""), + arguments=block.get("input"), + ) + ) + elif block_type == "tool_result": + use_id = block.get("tool_use_id") + parts.append( + SpanMessagePart( + type="tool_call_response", + id=use_id if isinstance(use_id, str) else None, + result=block.get("content"), + ) + ) + return parts + + +class ToolCatalog: + """The tools the model could call, widened as the run reports more. + + The AI Config's own tools are known before the run starts. The ones Claude Code brings itself + (Read, Bash, and the rest) are announced only in the ``init`` message, and only by name; their + schemas stay in that process. A name with no ``parameters`` says "this tool was offered, its + schema is not ours to state", which describes the run better than omitting a tool the model could + see. + + Every entry is keyed on the name the model saw, which is also the name that tool's + ``execute_tool`` span carries. ``native_names`` maps an AI Config key to the provider tool name it + stands for, for the config's native tools; absent entries are user-defined tools, catalogued under + their own name. + + Shared between the root and the ``chat`` spans so the two cannot disagree about the same + attribute. + """ + + def __init__( + self, + config_tools: dict[str, Any] | None, + native_names: dict[str, str], + ) -> None: + self._definitions: list[ToolDefinitionInput] = [] + for key, tool in (config_tools or {}).items(): + name = native_names.get(key) or tool.get("name") or key + self._definitions.append( + ToolDefinitionInput( + name=name, + description=tool.get("description"), + parameters=tool.get("parameters"), + ) + ) + self._named: set[str] = {d.name for d in self._definitions} + + @property + def current(self) -> list[ToolDefinitionInput]: + """A copy, so widening the catalog later cannot alter a span already written from it.""" + return list(self._definitions) + + def widen(self, names: Any) -> bool: + """Absorbs the ``init`` message's tool list. + + Reports whether anything was added, so the caller rewrites the root's attribute only when + there is something new to say. Compared on display names: the CLI lists an AI Config tool + under its MCP name, which would otherwise be added a second time alongside the entry that + already has its schema. + """ + if not isinstance(names, list): + return False + added = False + for name in names: + if not isinstance(name, str): + continue + display = tool_display_name(name) + if display in self._named: + continue + self._named.add(display) + self._definitions.append(ToolDefinitionInput(name=display)) + added = True + return added + + +@dataclass +class Opening: + system_instructions: str | None + messages: list[SpanMessage] + + +@dataclass +class _Inference: + request_id: str | None + parent_tool_use_id: str | None + subagent_type: str | None + model: str + session_id: str | None + usage: dict[str, Any] + finish_reason: str | None + start_time: int + end_time: int + parts: list[SpanMessagePart] + input_messages: list[SpanMessage] + + +class InferenceSpans: + """Emits the run's ``chat`` spans, one per API call. + + ``query()`` reports no request boundaries, and an ``AssistantMessage`` is not one: the CLI emits + one message per content block of a response, so a single API call surfaces as several messages + that share a ``message_id`` (the Anthropic response id) and repeat the same usage bag, with tool + executions interleaved between them. So the response id is the unit, accumulated across the whole + run rather than only while consecutive, because the messages of one response are not adjacent in + the stream. + + Spans are built at :meth:`finish` rather than as messages arrive, so a response is described + once, in full. :meth:`finish` must therefore run even when the run throws, since the exporter + only receives ended spans. + """ + + def __init__( + self, + config: AiConfigRep, + parent_context: Any, + capture_content: bool, + catalog: ToolCatalog, + opening: Opening, + ) -> None: + self._config = config + self._parent = parent_context + self._capture = capture_content + self._catalog = catalog + self._opening = opening + self._boundary = time.time_ns() + self._inferences: list[_Inference] = [] + self._by_request_id: dict[str, _Inference] = {} + # One conversation per agent thread, keyed by parent_tool_use_id (None = main thread). A + # subagent's model calls arrive on the same stream as the main thread's, interleaved, so a + # single list would hand a main-thread call an input containing turns from a conversation it + # was never part of. + self._threads: dict[str | None, list[SpanMessage]] = {} + self._totals: dict[str, int] = { + "input_tokens": 0, + "output_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } + self._turns = 0 + + def _thread_for(self, parent_tool_use_id: str | None) -> list[SpanMessage]: + thread = self._threads.get(parent_tool_use_id) + if thread is None: + thread = list(self._opening.messages) if parent_tool_use_id is None else [] + self._threads[parent_tool_use_id] = thread + return thread + + def record(self, message: Any) -> None: + """Feed every message of the run, in order.""" + if isinstance(message, AssistantMessage): + self._absorb(message) + return + if isinstance(message, UserMessage): + self._absorb_user_turn(message) + if marks_local_work(message): + self._boundary = time.time_ns() + + def _absorb_user_turn(self, message: UserMessage) -> None: + """Adds a user turn (tool results, and any context the CLI injected) to the conversation.""" + parts = to_span_parts(message.content) + if not parts: + return + self._thread_for(message.parent_tool_use_id).append( + SpanMessage(role="user", parts=parts) + ) + + def finish(self) -> None: + """Emits a span per response and clears the accumulator. Idempotent.""" + if not self._inferences: + return + pending = self._inferences + self._inferences = [] + self._by_request_id.clear() + for inference in pending: + self._emit(inference) + + def _absorb(self, message: AssistantMessage) -> None: + request_id = message.message_id + known = self._by_request_id.get(request_id) if request_id else None + if known is not None: + # Another block of a response already recorded. Its usage is the same bag repeated. + known.parts.extend(to_span_parts(message.content)) + if message.stop_reason: + known.finish_reason = to_semconv_finish_reason(message.stop_reason) + return + + now = time.time_ns() + parent_tool_use_id = message.parent_tool_use_id + thread = self._thread_for(parent_tool_use_id) + parts = to_span_parts(message.content) + inference = _Inference( + request_id=request_id, + parent_tool_use_id=parent_tool_use_id, + subagent_type=getattr(message, "subagent_type", None), + model=message.model or model_name(self._config), + session_id=message.session_id, + usage=dict(message.usage or {}), + finish_reason=to_semconv_finish_reason(message.stop_reason), + start_time=self._boundary, + end_time=now, + parts=parts, + # Which turns this call was sent, captured before the reply joins them below. + input_messages=list(thread), + ) + self._inferences.append(inference) + if request_id: + self._by_request_id[request_id] = inference + thread.append(SpanMessage(role="assistant", parts=inference.parts)) + + self._turns += 1 + for key in self._totals: + self._totals[key] += number_or_zero(inference.usage.get(key)) + self._boundary = now + + @property + def run_usage(self) -> dict[str, Any]: + """The run's spend so far, for the paths where the CLI never reported its own total. + + ``reported`` is false only when no response was ever absorbed. Writing an all-zero total in + that case would assert the run cost nothing, which a run that died before its first response + cannot honestly claim. + """ + return {"total": dict(self._totals), "reported": self._turns > 0} + + def _emit(self, inference: _Inference) -> None: + if not _HAS_OTEL: + return + span = trace.get_tracer(TRACER_NAME).start_span( + f"chat {inference.model}", + context=self._parent, + start_time=inference.start_time, + ) + span.set_attribute("gen_ai.operation.name", "chat") + set_model_identity_attributes(span, PROVIDER, inference.model) + # The model the turn actually used, read off the streamed inference. Not + # config.model.name: this is one of only two handlers where the two may disagree. See + # TELEMETRY-CONTRACT.md section 2a. + span.set_attribute("gen_ai.response.model", inference.model) + 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) + # 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. + if inference.finish_reason: + span.set_attribute( + "gen_ai.response.finish_reasons", [inference.finish_reason] + ) + set_usage_span_attributes(span, add_cached_tokens_to_input(inference.usage)) + if inference.subagent_type: + span.set_attribute("gen_ai.agent.name", inference.subagent_type) + + main_thread = inference.parent_tool_use_id is None + set_input_content_attributes( + span, + self._capture, + # Both only on the main thread. A subagent runs under its own agent definition's prompt + # and its own subset of tools, neither of which this side is told. + system_instructions=( + self._opening.system_instructions if main_thread else None + ), + messages=inference.input_messages, + tool_definitions=self._catalog.current if main_thread else None, + ) + set_output_content_attributes( + span, + self._capture, + [ + SpanMessage( + role="assistant", + parts=inference.parts, + finish_reason=inference.finish_reason, + ) + ], + ) + span.set_status(SpanStatusCode.OK) + span.end(end_time=inference.end_time) diff --git a/packages/claude-agents/tests/test_handler.py b/packages/claude-agents/tests/test_handler.py index 1f7eaaf..7c1c225 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -1,7 +1,15 @@ """ Tests for launchdarkly-ai-claude-agents handler. -Covers §1.1–1.9 (generic) and claude-agent-specific extras. -Reference: TESTING.md §1, §2.x (Anthropic) + +Rewritten against TELEMETRY-CONTRACT.md, replacing the old flat-span assertions. Uses a real +``TracerProvider`` + ``InMemorySpanExporter`` rather than a mocked ``opentelemetry.trace`` module, +the same choice ``@launchdarkly/ai-claude-agents``'s ``spans.test.ts`` makes: a mocked tracer cannot +see whether parent/child wiring is right, only whether the right methods were called. + +``query()`` is replaced with a fake async generator that replays a scripted message stream, built +from the Agent SDK's own dataclasses (``AssistantMessage``, ``UserMessage``, ``SystemMessage``, +``ResultMessage``, ``StreamEvent``) rather than mocks, so a structural change to those types would +fail loudly here instead of silently. """ from __future__ import annotations @@ -11,91 +19,188 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from claude_agent_sdk import ( + AssistantMessage, + ResultMessage, + StreamEvent, + SystemMessage, + TextBlock, + ToolResultBlock, + ToolUseBlock, + UserMessage, +) +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode import launchdarkly_ai_claude_agents.handler as handler_mod from launchdarkly_ai_claude_agents.handler import ( + _native_tool_aliases, build_prompt, create_claude_agents_handler, partition_tools, ) # --------------------------------------------------------------------------- -# Helpers +# A real tracer provider, reset between tests # --------------------------------------------------------------------------- +_exporter = InMemorySpanExporter() +_provider = TracerProvider() +_provider.add_span_processor(SimpleSpanProcessor(_exporter)) +trace.set_tracer_provider(_provider) + + +@pytest.fixture(autouse=True) +def _reset_exporter() -> None: + _exporter.clear() + + +def spans() -> list[Any]: + return list(_exporter.get_finished_spans()) + + +def named(prefix: str) -> list[Any]: + return [s for s in spans() if s.name.startswith(prefix)] + + +def root() -> Any: + return next(s for s in spans() if s.name == "invoke_agent") + + +# --------------------------------------------------------------------------- +# Message builders, using the Agent SDK's own dataclasses +# --------------------------------------------------------------------------- + +BASE_CONFIG: dict[str, Any] = { + "model": {"name": "claude-opus-4-5"}, + "provider": {"name": "Anthropic"}, + "instructions": "You are helpful.", +} + def _make_config(**kwargs: Any) -> dict[str, Any]: + """Restored from the pre-rewrite file: a couple of the restored non-telemetry tests build a + config inline rather than through ``BASE_CONFIG``/``TOOL_CONFIG``. + """ base = {"model": {"name": "claude-opus-4-5"}, "provider": {"name": "Anthropic"}} base.update(kwargs) return base -class _MockResultMessage: - """Distinct class so isinstance checks work in _stream_gen / _call_impl.""" - - def __init__( - self, text: str = "hello", input_tokens: int = 10, output_tokens: int = 5 - ) -> None: - self.result = text - self.usage = {"input_tokens": input_tokens, "output_tokens": output_tokens} - self.is_error = False +TOOL_CONFIG: dict[str, Any] = { + **BASE_CONFIG, + "tools": { + "search": { + "name": "search", + "type": "function", + "parameters": {"type": "object", "properties": {}}, + } + }, +} + + +def assistant_message( + input_tokens: int = 10, + output_tokens: int = 2, + message_id: str | None = "msg_1", + content: list[Any] | None = None, + stop_reason: str | None = "end_turn", + session_id: str = "sess-1", + parent_tool_use_id: str | None = None, +) -> AssistantMessage: + return AssistantMessage( + content=content or [], + model="claude-opus-4-5", + parent_tool_use_id=parent_tool_use_id, + usage={"input_tokens": input_tokens, "output_tokens": output_tokens}, + message_id=message_id, + stop_reason=stop_reason, + session_id=session_id, + ) -class _MockStreamEvent: - """Distinct class so isinstance checks work in _stream_gen.""" +def result_message( + result: str = "agent output", + subtype: str = "success", + input_tokens: int = 22, + output_tokens: int = 5, + errors: list[str] | None = None, +) -> ResultMessage: + return ResultMessage( + subtype=subtype, + duration_ms=1, + duration_api_ms=1, + is_error=subtype != "success", + num_turns=1, + session_id="sess-1", + usage={"input_tokens": input_tokens, "output_tokens": output_tokens}, + result=result, + errors=errors, + ) - def __init__(self, delta_text: str) -> None: - self.event = { - "type": "content_block_delta", - "delta": {"type": "text_delta", "text": delta_text}, - } +def init_message( + session_id: str = "sess-1", tools: list[str] | None = None +) -> SystemMessage: + data: dict[str, Any] = {"session_id": session_id} + if tools is not None: + data["tools"] = tools + return SystemMessage(subtype="init", data=data) -def _make_result_message( - text: str = "hello", input_tokens: int = 10, output_tokens: int = 5 -) -> Any: - return _MockResultMessage(text, input_tokens, output_tokens) +def tool_result_user_message( + tool_use_id: str = "tu-1", content: str = "found it" +) -> UserMessage: + return UserMessage( + content=[ToolResultBlock(tool_use_id=tool_use_id, content=content)] + ) -def _make_stream_event(delta_text: str) -> Any: - return _MockStreamEvent(delta_text) +def stream_event(text: str) -> StreamEvent: + return StreamEvent( + uuid="u1", + session_id="sess-1", + event={ + "type": "content_block_delta", + "delta": {"type": "text_delta", "text": text}, + }, + ) -async def _async_gen_from(*messages: Any) -> AsyncIterator[Any]: - for m in messages: - yield m +def _fake_query(messages: list[Any]): + async def _query(**_kwargs: Any) -> AsyncIterator[Any]: + for m in messages: + yield m -def _patch_query(messages: list[Any]) -> Any: - """Context manager that patches claude_agent_sdk.query in the handler module.""" + return _query - async def _query(**kwargs: Any) -> AsyncIterator[Any]: - async for m in _async_gen_from(*messages): - yield m - mock_sdk = MagicMock() - mock_sdk.query = _query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = _MockResultMessage - mock_sdk.StreamEvent = _MockStreamEvent - mock_sdk.tool = MagicMock(return_value=lambda fn: fn) - mock_sdk.create_sdk_mcp_server = MagicMock(return_value=MagicMock()) - mock_sdk.HookMatcher = MagicMock() - return patch( - "importlib.import_module", - side_effect=lambda n: mock_sdk if n == "claude_agent_sdk" else __import__(n), - ) +async def _collect(gen: Any) -> list[Any]: + out = [] + async for event in gen: + out.append(event) + return out # --------------------------------------------------------------------------- -# §1.1 Factory +# Factory # --------------------------------------------------------------------------- class TestFactory: def test_returns_callable(self) -> None: - h = create_claude_agents_handler() - assert callable(h) + assert callable(create_claude_agents_handler()) + + def test_provides_for(self) -> None: + assert create_claude_agents_handler().provides_for == ("Anthropic", "agent") + + def test_multiple_calls_independent(self) -> None: + assert create_claude_agents_handler() is not create_claude_agents_handler() + + # --- restored from the pre-rewrite file (not telemetry) --- def test_attaches_provides_for(self) -> None: h = create_claude_agents_handler() @@ -113,11 +218,29 @@ def test_multiple_calls_return_independent_instances(self) -> None: # --------------------------------------------------------------------------- -# §1.2 Prompt construction (tested via build_prompt directly) +# Prompt construction # --------------------------------------------------------------------------- class TestPromptConstruction: + def test_instructions_become_system_prompt(self) -> None: + prompt, system = build_prompt(BASE_CONFIG, "hi", {}) + assert prompt == "hi" + assert system == "You are helpful." + + def test_no_instructions_no_system_prompt(self) -> None: + prompt, system = build_prompt({"model": {"name": "m"}}, "hi", {}) + assert system is None + assert prompt == "hi" + + def test_variable_substitution(self) -> None: + cfg = {**BASE_CONFIG, "instructions": "Hello {{name}}."} + _, system = build_prompt(cfg, "hi", {"name": "Ada"}) + assert system == "Hello Ada." + + # --- restored from the pre-rewrite file (not telemetry); byte-for-byte, only + # ``_make_config`` inlined since the old module-level helper was removed --- + def test_path_a_instructions(self) -> None: config = _make_config(instructions="You are a helper.") prompt, system = build_prompt(config, "hi", {}) @@ -178,9 +301,73 @@ def test_path_c_instructions_takes_priority_over_messages(self) -> None: _prompt, system = build_prompt(config, "q", {}) assert system == "Use instructions." + def test_messages_mode_extracts_system_and_flattens_history_into_prompt( + self, + ) -> None: + cfg = { + "model": {"name": "m"}, + "messages": [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "context line"}, + ], + } + prompt, system = build_prompt(cfg, "final question", {}) + assert system == "Be terse." + assert "context line" in prompt + assert prompt.endswith("final question") + + def test_history_appended_to_system_prompt(self) -> None: + history = [{"role": "user", "content": "earlier"}] + _, system = build_prompt(BASE_CONFIG, "hi", {}, history=history) + assert "earlier" in (system or "") + assert "You are helpful." in (system or "") + + def test_no_user_input_defaults_to_empty_string(self) -> None: + prompt, _ = build_prompt(BASE_CONFIG, None, {}) + assert prompt == "" + + +# --------------------------------------------------------------------------- +# partition_tools / native tool aliases +# --------------------------------------------------------------------------- + + +class TestPartitionTools: + def test_user_tool_goes_to_mcp_bucket(self) -> None: + native_map, user_tools, native_names = partition_tools( + TOOL_CONFIG["tools"], {"search": lambda _: "r"} + ) + assert native_map == {} + assert "search" in user_tools + assert native_names == [] + + def test_native_tool_goes_to_native_bucket(self) -> None: + from launchdarkly_ai_claude_agents.builtins import ClaudeWebSearch + from launchdarkly_ai_server import NATIVE_TOOL_KEY + + # A tracking stub, as `wrap_tool_handlers` produces: a callable with the NativeTool + # stashed under NATIVE_TOOL_KEY, not the NativeTool instance itself. + stub = lambda: None # noqa: E731 + setattr(stub, NATIVE_TOOL_KEY, ClaudeWebSearch) + _native_map, user_tools, native_names = partition_tools( + {"webSearch": {"name": "webSearch"}}, {"webSearch": stub} + ) + assert native_names == ["WebSearch"] + assert user_tools == {} + + def test_native_tool_aliases_map_ld_key_to_provider_name(self) -> None: + from launchdarkly_ai_claude_agents.builtins import ClaudeWebSearch + from launchdarkly_ai_server import NATIVE_TOOL_KEY + + stub = lambda: None # noqa: E731 + setattr(stub, NATIVE_TOOL_KEY, ClaudeWebSearch) + aliases = _native_tool_aliases({"webSearch": stub}) + assert aliases == {"webSearch": "WebSearch"} + # --------------------------------------------------------------------------- -# §1.3 Tool conversion +# §1.3 Tool conversion — restored from the pre-rewrite file (not telemetry). +# ``partition_tools`` kept its 3-tuple return, so these run byte-for-byte. # --------------------------------------------------------------------------- @@ -207,12 +394,47 @@ def test_empty_tools_no_tools_sent(self) -> None: # --------------------------------------------------------------------------- -# §1.4 Tool execution loop (via build_tool_mcp) +# §1.4 Tool execution loop (via build_tool_mcp) — restored from the pre-rewrite +# file. ``build_tool_mcp`` kept its lazy ``importlib.import_module`` pattern +# (native_graph.py depends on it), so the old SDK-mocking approach still works +# unmodified for this one. # --------------------------------------------------------------------------- +class _MockResultMessageForToolMcp: + """Distinct class so isinstance checks the handler makes still work.""" + + def __init__( + self, text: str = "hello", input_tokens: int = 10, output_tokens: int = 5 + ) -> None: + self.result = text + self.usage = {"input_tokens": input_tokens, "output_tokens": output_tokens} + self.is_error = False + + +def _patch_query_for_tool_mcp(messages: list[Any]) -> Any: + """Patches ``claude_agent_sdk`` for the one restored test that exercises + ``build_tool_mcp`` directly, which still resolves the SDK lazily. + """ + + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + for m in messages: + yield m + + mock_sdk = MagicMock() + mock_sdk.query = _query + mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) + mock_sdk.ResultMessage = _MockResultMessageForToolMcp + mock_sdk.tool = MagicMock(return_value=lambda fn: fn) + mock_sdk.create_sdk_mcp_server = MagicMock(return_value=MagicMock()) + mock_sdk.HookMatcher = MagicMock() + return patch( + "importlib.import_module", + side_effect=lambda n: mock_sdk if n == "claude_agent_sdk" else __import__(n), + ) + + class TestToolExecutionLoop: - @pytest.mark.asyncio async def test_tool_not_found_throws(self) -> None: from launchdarkly_ai_claude_agents.handler import build_tool_mcp @@ -220,7 +442,7 @@ async def test_tool_not_found_throws(self) -> None: handlers: dict[str, Any] = {} # no handler registered # The execute closure inside build_tool_mcp raises if handler missing - with _patch_query([_make_result_message()]): + with _patch_query_for_tool_mcp([_MockResultMessageForToolMcp()]): mcp = await build_tool_mcp(config_tools, handlers) # Directly call the stored execute fn for t in getattr(mcp, "tools", None) or []: @@ -229,745 +451,1181 @@ async def test_tool_not_found_throws(self) -> None: with pytest.raises(ValueError, match="No handler"): await fn({"key": "val"}) - @pytest.mark.asyncio - async def test_no_tools_in_config_handler_never_invoked(self) -> None: - result_msg = _make_result_message("done") + async def test_no_tools_in_config_handler_never_invoked( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: mock_handler = AsyncMock(return_value="tool-output") - - with _patch_query([result_msg]): - h = create_claude_agents_handler() - config = _make_config() - output = await h(config, "hi", {"my-tool": mock_handler}) + monkeypatch.setattr(handler_mod, "query", _fake_query([result_message("done")])) + h = create_claude_agents_handler() + output = await h(BASE_CONFIG, "hi", {"my-tool": mock_handler}) mock_handler.assert_not_called() assert output["output"] == "done" # --------------------------------------------------------------------------- -# §1.5 Telemetry +# Span tree — TELEMETRY-CONTRACT.md section 1 # --------------------------------------------------------------------------- -class TestTelemetry: - @pytest.mark.asyncio - async def test_span_name(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - mock_trace.get_tracer.return_value.start_span.assert_called_with("claude.query") - - @pytest.mark.asyncio - async def test_gen_ai_system(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.system") == "anthropic" - - @pytest.mark.asyncio - async def test_gen_ai_operation_name(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.operation.name") == "chat" - - @pytest.mark.asyncio - async def test_gen_ai_request_model(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(model={"name": "claude-opus-4-5"}), "hi") - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.request.model") == "claude-opus-4-5" - - @pytest.mark.asyncio - async def test_gen_ai_content_prompt_event(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hello world") - - event_calls = [ - c - for c in mock_span.add_event.call_args_list - if c[0][0] == "gen_ai.content.prompt" - ] - assert event_calls - # The gen_ai.prompt attribute must include the user input text - prompt_attr = event_calls[0][0][1].get("gen_ai.prompt", "") - assert "hello world" in prompt_attr, ( - f"gen_ai.prompt must include user input 'hello world', got: {prompt_attr!r}" +class TestSpanTree: + async def test_root_span_name(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + assert root().name == "invoke_agent" + assert root().attributes["gen_ai.operation.name"] == "invoke_agent" - @pytest.mark.asyncio - async def test_token_attributes_set(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out", input_tokens=42, output_tokens=7) - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("gen_ai.usage.input_tokens") == 42 - assert calls.get("gen_ai.usage.output_tokens") == 7 - - @pytest.mark.asyncio - async def test_gen_ai_content_completion_event(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("final answer") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - event_calls = [ - c - for c in mock_span.add_event.call_args_list - if c[0][0] == "gen_ai.content.completion" - ] - assert event_calls - - @pytest.mark.asyncio - async def test_span_status_ok(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - with patch.object(handler_mod, "SpanStatusCode", MagicMock()) as _: - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - mock_span.set_status.assert_called() - - @pytest.mark.asyncio - async def test_span_end_always_called(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi") - - mock_span.end.assert_called() - - @pytest.mark.asyncio - async def test_gen_ai_response_model(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(model={"name": "claude-opus-4-5"}), "hi") - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert "gen_ai.response.model" in calls - assert calls["gen_ai.response.model"] == "claude-opus-4-5" - - @pytest.mark.asyncio - async def test_ld_span_attributes(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - } - } - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi", variables=variables) - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.operation.type") == "gen_ai" - assert calls.get("launchdarkly.config.key") == "my-config" - assert calls.get("launchdarkly.variation.key") == "v1" - assert calls.get("launchdarkly.run.id") == "run-abc" - assert "launchdarkly.graph.key" not in calls - - async def test_ld_graph_key_set_when_present(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("out") - variables = { - "__ld": { - "configKey": "my-config", - "variationKey": "v1", - "runId": "run-abc", - "graphKey": "my-graph", - } - } - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - await h(_make_config(), "hi", variables=variables) + async def test_one_chat_span_per_model_response( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message(message_id="req_1"), + assistant_message(message_id="req_2"), + result_message(), + ] + ), + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + chats = named("chat ") + assert len(chats) == 2 + assert chats[0].name == "chat claude-opus-4-5" + + async def test_one_chat_span_per_call_not_per_message_block( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The CLI splits one API response into several assistant messages sharing one message id. + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message( + 23669, 8, "req_shared", content=[TextBlock(text="a")] + ), + assistant_message( + 23669, 8, "req_shared", content=[TextBlock(text="b")] + ), + result_message(), + ] + ), + ) + await create_claude_agents_handler(capture_content=True)(BASE_CONFIG, "q") + chats = named("chat ") + assert len(chats) == 1 + assert chats[0].attributes["gen_ai.usage.input_tokens"] == 23669 - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.graph.key") == "my-graph" + async def test_execute_tool_span_per_tool_call_sibling_of_chat( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + captured: dict[str, Any] = {} + + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + captured["options"] = kwargs["options"] + yield assistant_message( + content=[ + ToolUseBlock(id="tu-1", name="mcp__tool-mcp__search", input={}) + ] + ) + 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": {"q": "x"}, + "session_id": "sess-1", + }, + "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": "found", + }, + "tu-1", + None, + ) + yield assistant_message(message_id="req_2") + yield result_message() + + monkeypatch.setattr(handler_mod, "query", _query) + await create_claude_agents_handler()( + TOOL_CONFIG, "q", {"search": lambda _: "r"} + ) + + tools = named("execute_tool ") + assert len(tools) == 1 + assert tools[0].name == "execute_tool search" + assert tools[0].attributes["gen_ai.tool.name"] == "search" + assert tools[0].attributes["gen_ai.tool.call.id"] == "tu-1" + # A sibling of chat: same parent (the root), not nested under a chat span. + assert tools[0].parent.span_id == root().context.span_id + + async def test_children_carry_no_launchdarkly_attributes( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) + ) + await create_claude_agents_handler()( + BASE_CONFIG, + "q", + variables={"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}}, + ) + for child in named("chat "): + assert not [k for k in child.attributes if k.startswith("launchdarkly.")] + assert "feature_flag" not in [e.name for e in child.events] + + async def test_every_span_is_ended_ok( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + for s in spans(): + assert s.status.status_code == StatusCode.OK # --------------------------------------------------------------------------- -# §1.6 Error handling +# Root span attributes — TELEMETRY-CONTRACT.md sections 2, 2a, 8 # --------------------------------------------------------------------------- -class TestErrorHandling: - @pytest.mark.asyncio - async def test_records_exception_on_span(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - async def _broken_query(**kwargs: Any) -> AsyncIterator[Any]: - raise RuntimeError("provider down") - yield # make it a generator - - mock_sdk = MagicMock() - mock_sdk.query = _broken_query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = MagicMock - mock_sdk.StreamEvent = MagicMock - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") - - mock_span.record_exception.assert_called() - - @pytest.mark.asyncio - async def test_sets_span_status_error(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - async def _broken_query(**kwargs: Any) -> AsyncIterator[Any]: - raise RuntimeError("fail") - yield - - mock_sdk = MagicMock() - mock_sdk.query = _broken_query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = MagicMock - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") - - mock_span.set_status.assert_called() - - @pytest.mark.asyncio - async def test_ends_span_on_error(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - async def _broken_query(**kwargs: Any) -> AsyncIterator[Any]: - raise RuntimeError("fail") - yield - - mock_sdk = MagicMock() - mock_sdk.query = _broken_query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = MagicMock - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - with pytest.raises(RuntimeError): - await h(_make_config(), "hi") - - mock_span.end.assert_called() - - @pytest.mark.asyncio - async def test_rethrows_error(self) -> None: - async def _broken_query(**kwargs: Any) -> AsyncIterator[Any]: - raise RuntimeError("specific error") - yield - - mock_sdk = MagicMock() - mock_sdk.query = _broken_query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = MagicMock - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) +class TestRootAttributes: + async def test_provider_keys_and_requested_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + attrs = root().attributes + assert attrs["gen_ai.system"] == "anthropic" + assert attrs["gen_ai.provider.name"] == "anthropic" + assert attrs["gen_ai.request.model"] == "claude-opus-4-5" + # The root reports the requested name even though the chat span (below) may differ. + assert attrs["gen_ai.response.model"] == "claude-opus-4-5" + + async def test_run_total_not_one_turn( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + # The result message's cumulative usage, not the per-turn figure. + assert root().attributes["gen_ai.usage.input_tokens"] == 22 + assert root().attributes["gen_ai.usage.output_tokens"] == 5 + + async def test_conversation_id_from_init_message( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [init_message("sess-abc"), assistant_message(), result_message()] ), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - with pytest.raises(RuntimeError, match="specific error"): - await h(_make_config(), "hi") + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + assert root().attributes["gen_ai.conversation.id"] == "sess-abc" + async def test_conversation_id_only_on_root_chat_and_execute_tool( + 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("done") + + monkeypatch.setattr(handler_mod, "query", _query) + await create_claude_agents_handler()( + TOOL_CONFIG, "q", {"search": lambda _: "r"} + ) + assert root().attributes["gen_ai.conversation.id"] == "sess-abc" + assert named("chat ")[0].attributes["gen_ai.conversation.id"] == "sess-abc" + assert ( + named("execute_tool ")[0].attributes["gen_ai.conversation.id"] == "sess-abc" + ) -# --------------------------------------------------------------------------- -# §1.7 Convenience export -# --------------------------------------------------------------------------- + async def test_no_conversation_id_without_init( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(handler_mod, "query", _fake_query([result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + assert "gen_ai.conversation.id" not in root().attributes -class TestConvenienceExport: - def test_calls_through_to_model_call(self) -> None: - from launchdarkly_ai_claude_agents.handler import claude_agents +# --------------------------------------------------------------------------- +# Chat span attributes — sections 2a, 3, 5b, 8 +# --------------------------------------------------------------------------- - assert callable(claude_agents) - def test_passes_config_key_user_input_and_context(self) -> None: - import inspect +class TestChatAttributes: + async def test_response_model_is_what_the_turn_actually_used( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # claude-agents is one of only two handlers whose chat span may disagree with the root. + turn = assistant_message() + turn.model = "claude-opus-4-5-20250101" + monkeypatch.setattr(handler_mod, "query", _fake_query([turn, result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + assert chat.attributes["gen_ai.response.model"] == "claude-opus-4-5-20250101" + assert root().attributes["gen_ai.response.model"] == "claude-opus-4-5" + + async def test_finish_reason_usually_absent( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Measured against Agent SDK 0.3.220: stop_reason is null on every assistant message. + turn = assistant_message(stop_reason=None) + monkeypatch.setattr(handler_mod, "query", _fake_query([turn, result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + assert "gen_ai.response.finish_reasons" not in chat.attributes + + async def test_finish_reason_mapped_when_the_sdk_reports_one( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + turn = assistant_message(stop_reason="end_turn") + monkeypatch.setattr(handler_mod, "query", _fake_query([turn, result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + assert list(chat.attributes["gen_ai.response.finish_reasons"]) == ["stop"] + + async def test_writes_all_seven_usage_attributes( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + for key in ( + "gen_ai.usage.input_tokens", + "gen_ai.usage.output_tokens", + "gen_ai.usage.total_tokens", + "gen_ai.usage.cache_read.input_tokens", + "gen_ai.usage.cache_creation.input_tokens", + "gen_ai.usage.prompt_tokens", + "gen_ai.usage.completion_tokens", + ): + assert key in chat.attributes - from launchdarkly_ai_claude_agents.handler import claude_agents + async def test_cache_tokens_folded_into_input( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # TELEMETRY-CONTRACT.md section 8: Anthropic reports cache buckets *beside* input_tokens, + # so this handler must add them in, not pass input through untouched. + turn = assistant_message(3, 8, "req_1") + turn.usage = { + "input_tokens": 3, + "output_tokens": 8, + "cache_read_input_tokens": 19971, + "cache_creation_input_tokens": 3580, + } + monkeypatch.setattr(handler_mod, "query", _fake_query([turn, result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + assert chat.attributes["gen_ai.usage.input_tokens"] == 3 + 19971 + 3580 + assert chat.attributes["gen_ai.usage.cache_read.input_tokens"] == 19971 + assert chat.attributes["gen_ai.usage.cache_creation.input_tokens"] == 3580 - sig = inspect.signature(claude_agents) - assert "config_key" in sig.parameters - assert "user_input" in sig.parameters - assert "context" in sig.parameters - def test_config_key_forwarded_as_key(self) -> None: - import launchdarkly_ai_claude_agents.handler as handler_mod +# --------------------------------------------------------------------------- +# Content capture — section 7 +# --------------------------------------------------------------------------- - mock_config_instance = MagicMock() - mock_config_fn = MagicMock(return_value=mock_config_instance) - mock_config_instance.invoke = MagicMock(return_value="result") - with patch.object(handler_mod, "config", mock_config_fn): - from launchdarkly_ai_claude_agents.handler import claude_agents +class TestContentCapture: + async def test_emits_no_content_at_all_by_default( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [init_message(tools=["Read"]), assistant_message(), result_message()] + ), + ) + await create_claude_agents_handler()(TOOL_CONFIG, "q") + for span in (root(), *named("chat ")): + assert "gen_ai.input.messages" not in span.attributes + assert "gen_ai.output.messages" not in span.attributes + assert "gen_ai.system_instructions" not in span.attributes + assert "gen_ai.tool.definitions" not in span.attributes + + async def test_root_carries_no_content_by_default_either( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query([assistant_message(), result_message("hi")]), + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + assert "gen_ai.output.messages" not in root().attributes - ctx = {"kind": "user", "key": "u1"} - claude_agents("my-flag", "hello", ctx) + async def test_tool_call_arguments_and_result_gated_by_capture( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + yield assistant_message( + content=[ + ToolUseBlock( + id="tu-1", name="mcp__tool-mcp__search", input={"q": "x"} + ) + ] + ) + 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": {"q": "x"}, + }, + "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": "found it", + }, + "tu-1", + None, + ) + yield result_message("done") + + monkeypatch.setattr(handler_mod, "query", _query) + await create_claude_agents_handler()( + TOOL_CONFIG, "q", {"search": lambda _: "r"} + ) + [tool_span] = named("execute_tool ") + assert "gen_ai.tool.call.arguments" not in tool_span.attributes + assert "gen_ai.tool.call.result" not in tool_span.attributes - mock_config_fn.assert_called_once() - call_kwargs = mock_config_fn.call_args.kwargs - assert call_kwargs.get("key") == "my-flag" - handler = call_kwargs.get("handler") - assert handler is not None - assert handler.provides_for == ("Anthropic", "agent") - mock_config_instance.invoke.assert_called_once_with( - "hello", ctx, variables=None + async def test_content_present_when_enabled( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message(content=[TextBlock(text="hi")]), + result_message("hi"), + ] + ), ) + await create_claude_agents_handler(capture_content=True)(BASE_CONFIG, "q") + [chat] = named("chat ") + assert "gen_ai.input.messages" in chat.attributes + assert "gen_ai.output.messages" in chat.attributes + assert "gen_ai.system_instructions" in root().attributes - def test_callable_without_extra_kwargs(self) -> None: - import launchdarkly_ai_claude_agents.handler as handler_mod - mock_config_instance = MagicMock() - mock_config_fn = MagicMock(return_value=mock_config_instance) - mock_config_instance.invoke = MagicMock(return_value="result") +# --------------------------------------------------------------------------- +# Errors, and the failure path — section 6 +# --------------------------------------------------------------------------- - with patch.object(handler_mod, "config", mock_config_fn): - from launchdarkly_ai_claude_agents.handler import claude_agents - ctx = {"kind": "user", "key": "u1"} - claude_agents("my-flag", "hello", ctx) +class TestErrorHandling: + async def test_error_result_fails_root_but_keeps_spend( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message(), + result_message( + subtype="error_max_turns", + input_tokens=40000, + output_tokens=500, + errors=["turn limit reached"], + ), + ] + ), + ) + with pytest.raises(RuntimeError, match="error_max_turns"): + await create_claude_agents_handler()(BASE_CONFIG, "q") + attrs = root().attributes + assert root().status.status_code == StatusCode.ERROR + assert attrs["gen_ai.usage.input_tokens"] == 40000 + assert attrs["gen_ai.usage.output_tokens"] == 500 + + async def test_reports_responses_that_arrived_when_the_sdk_throws( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**_kwargs: Any) -> AsyncIterator[Any]: + yield assistant_message(100, 20, "req_1") + yield assistant_message(200, 30, "req_2") + raise RuntimeError("transport died") + + monkeypatch.setattr(handler_mod, "query", _query) + with pytest.raises(RuntimeError, match="transport died"): + await create_claude_agents_handler()(BASE_CONFIG, "q") + attrs = root().attributes + assert attrs["gen_ai.usage.input_tokens"] == 300 + assert attrs["gen_ai.usage.output_tokens"] == 50 + assert root().status.status_code == StatusCode.ERROR + + async def test_no_usage_written_when_nothing_ever_arrived( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**_kwargs: Any) -> AsyncIterator[Any]: + raise RuntimeError("spawn failed") + yield # pragma: no cover - keeps this an async generator - mock_config_fn.assert_called_once() - mock_config_instance.invoke.assert_called_once_with( - "hello", ctx, variables=None + monkeypatch.setattr(handler_mod, "query", _query) + with pytest.raises(RuntimeError, match="spawn failed"): + await create_claude_agents_handler()(BASE_CONFIG, "q") + assert "gen_ai.usage.input_tokens" not in root().attributes + + async def test_result_omitted_stream_reports_per_response_sum( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message(100, 20, "req_1"), + assistant_message(200, 30, "req_2"), + ] + ), ) + result = await create_claude_agents_handler()(BASE_CONFIG, "q") + assert root().attributes["gen_ai.usage.input_tokens"] == 300 + assert result["usage"]["input_tokens"] == 300 # --------------------------------------------------------------------------- -# §1.8 Streaming +# Streaming # --------------------------------------------------------------------------- class TestStreaming: - def test_stream_is_defined(self) -> None: + async def test_yields_chunks_then_one_done_event( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [stream_event("Hi"), assistant_message(), result_message("Hi")] + ), + ) + h = create_claude_agents_handler() + events = await _collect(await h.stream(BASE_CONFIG, "q", {}, {})) + assert [e["type"] for e in events] == ["chunk", "done"] + assert events[0]["text"] == "Hi" + assert events[-1]["output"] == "Hi" + + async def test_emits_same_span_tree_as_blocking_path( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [stream_event("Hi"), assistant_message(), result_message("Hi")] + ), + ) + h = create_claude_agents_handler() + await _collect(await h.stream(BASE_CONFIG, "q", {}, {})) + assert sorted(s.name for s in spans()) == [ + "chat claude-opus-4-5", + "invoke_agent", + ] + assert root().status.status_code == StatusCode.OK + + async def test_error_fails_spans_and_reraises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**_kwargs: Any) -> AsyncIterator[Any]: + yield assistant_message() + raise RuntimeError("boom") + + monkeypatch.setattr(handler_mod, "query", _query) + h = create_claude_agents_handler() + with pytest.raises(RuntimeError, match="boom"): + await _collect(await h.stream(BASE_CONFIG, "q", {}, {})) + assert root().status.status_code == StatusCode.ERROR + assert root().attributes["gen_ai.usage.input_tokens"] == 10 + + async def test_abandoned_stream_ends_every_span_unset( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _slow_query(**_kwargs: Any) -> AsyncIterator[Any]: + yield stream_event("chunk-1") + yield assistant_message(message_id="req_1") + yield stream_event("chunk-2") + yield assistant_message(message_id="req_2") + yield result_message("done") + + monkeypatch.setattr(handler_mod, "query", _slow_query) + h = create_claude_agents_handler() + gen = await h.stream(BASE_CONFIG, "q", {}, {}) + # Consume only the first chunk, then abandon the generator without exhausting it. + first = await gen.__anext__() + assert first["type"] == "chunk" + await gen.aclose() + + assert root().status.status_code == StatusCode.UNSET + assert root().attributes.get("launchdarkly.stream.abandoned") is True + for s in spans(): + assert s.end_time is not None + + # --- restored from the pre-rewrite file (not telemetry). Adapted from the old + # ``_patch_query``/``_HAS_OTEL`` mocking approach to ``monkeypatch.setattr(handler_mod, + # "query", ...)`` plus real SDK dataclasses, because ``query`` is now a top-level import + # rather than something resolved through ``importlib.import_module`` on every call, and + # ``handler_mod`` no longer has its own ``_HAS_OTEL`` (that flag now lives in ``spans.py``). --- + + async def test_stream_is_defined(self) -> None: h = create_claude_agents_handler() assert hasattr(h, "stream") - @pytest.mark.asyncio - async def test_stream_returns_async_generator(self) -> None: + async def test_stream_returns_async_generator( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: import inspect - result_msg = _make_result_message("done") - with _patch_query([result_msg]): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - gen = await h.stream(_make_config(), "hi") - assert inspect.isasyncgen(gen) or hasattr(gen, "__aiter__") - - @pytest.mark.asyncio - async def test_yields_chunk_events_for_text_deltas(self) -> None: - chunk1 = _make_stream_event("hello ") - chunk2 = _make_stream_event("world") - result_msg = _make_result_message("hello world") + monkeypatch.setattr(handler_mod, "query", _fake_query([result_message("done")])) + h = create_claude_agents_handler() + gen = await h.stream(BASE_CONFIG, "hi") + assert inspect.isasyncgen(gen) or hasattr(gen, "__aiter__") - with _patch_query([chunk1, chunk2, result_msg]): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - events = [e async for e in await h.stream(_make_config(), "hi")] + async def test_yields_chunk_events_for_text_deltas( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + stream_event("hello "), + stream_event("world"), + result_message("hello world"), + ] + ), + ) + h = create_claude_agents_handler() + events = [e async for e in await h.stream(BASE_CONFIG, "hi")] chunks = [e for e in events if e.get("type") == "chunk"] assert len(chunks) == 2 assert chunks[0]["text"] == "hello " - @pytest.mark.asyncio - async def test_all_chunks_before_done(self) -> None: - chunk = _make_stream_event("part") - result_msg = _make_result_message("part") - - with _patch_query([chunk, result_msg]): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - events = [e async for e in await h.stream(_make_config(), "hi")] + async def test_all_chunks_before_done( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query([stream_event("part"), result_message("part")]), + ) + h = create_claude_agents_handler() + events = [e async for e in await h.stream(BASE_CONFIG, "hi")] done_idx = next(i for i, e in enumerate(events) if e.get("type") == "done") chunk_indices = [i for i, e in enumerate(events) if e.get("type") == "chunk"] assert all(ci < done_idx for ci in chunk_indices) - @pytest.mark.asyncio - async def test_yields_exactly_one_done_event(self) -> None: - result_msg = _make_result_message("done") - - with _patch_query([result_msg]): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - events = [e async for e in await h.stream(_make_config(), "hi")] + async def test_yields_exactly_one_done_event( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(handler_mod, "query", _fake_query([result_message("done")])) + h = create_claude_agents_handler() + events = [e async for e in await h.stream(BASE_CONFIG, "hi")] done_events = [e for e in events if e.get("type") == "done"] assert len(done_events) == 1 - @pytest.mark.asyncio - async def test_done_event_carries_correct_usage(self) -> None: - result_msg = _make_result_message("out", input_tokens=20, output_tokens=8) - - with _patch_query([result_msg]): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - events = [e async for e in await h.stream(_make_config(), "hi")] + async def test_done_event_carries_correct_usage( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query([result_message("out", input_tokens=20, output_tokens=8)]), + ) + h = create_claude_agents_handler() + events = [e async for e in await h.stream(BASE_CONFIG, "hi")] done = next(e for e in events if e.get("type") == "done") usage = done["usage"] assert usage.get("input_tokens") == 20 or usage.get("input") == 20 - @pytest.mark.asyncio - async def test_done_event_carries_accumulated_output(self) -> None: - result_msg = _make_result_message("hello world") - with _patch_query([result_msg]): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - events = [e async for e in await h.stream(_make_config(), "hi")] + async def test_done_event_carries_accumulated_output( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([result_message("hello world")]) + ) + h = create_claude_agents_handler() + events = [e async for e in await h.stream(BASE_CONFIG, "hi")] done = next(e for e in events if e.get("type") == "done") assert done["output"] == "hello world" - @pytest.mark.asyncio - async def test_generator_throws_on_provider_error(self) -> None: - async def _broken_query(**kwargs: Any) -> AsyncIterator[Any]: + async def test_generator_throws_on_provider_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _broken_query(**_kwargs: Any) -> AsyncIterator[Any]: raise RuntimeError("stream fail") - yield - - mock_sdk = MagicMock() - mock_sdk.query = _broken_query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = type(_make_result_message()) - mock_sdk.StreamEvent = type(_make_stream_event("x")) - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) + yield # pragma: no cover - keeps this an async generator + + monkeypatch.setattr(handler_mod, "query", _broken_query) + h = create_claude_agents_handler() + with pytest.raises(RuntimeError, match="stream fail"): + async for _ in await h.stream(BASE_CONFIG, "hi"): + pass + + +class TestQueryGeneratorLifecycle: + """TELEMETRY-CONTRACT.md section 6: the vendor generator, not just the span, must be closed.""" + + async def test_query_generator_closed_on_early_return( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + closed = {"value": False} + + async def _query(**_kwargs: Any) -> AsyncIterator[Any]: + try: + yield result_message("done") + yield assistant_message() # never reached: the handler returns after the result + finally: + closed["value"] = True + + monkeypatch.setattr(handler_mod, "query", _query) + await create_claude_agents_handler()(BASE_CONFIG, "q") + assert closed["value"] is True + + async def test_streaming_query_generator_closed_on_abandonment( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The streaming path's counterpart: iterating query(...) inline with no held reference + # leaks the vendor generator when the consumer abandons ours. See TELEMETRY-CONTRACT.md + # section 6, "claude-agents". + closed = {"value": False} + + async def _query(**_kwargs: Any) -> AsyncIterator[Any]: + try: + yield stream_event("chunk-1") + yield assistant_message() + yield result_message("done") + finally: + closed["value"] = True + + monkeypatch.setattr(handler_mod, "query", _query) + h = create_claude_agents_handler() + gen = await h.stream(BASE_CONFIG, "q", {}, {}) + first = await gen.__anext__() + assert first["type"] == "chunk" + await gen.aclose() + + assert closed["value"] is True + + +# --------------------------------------------------------------------------- +# Tool catalog widening — the CLI's own tools, announced only at `init` +# --------------------------------------------------------------------------- + + +class TestToolCatalog: + async def test_widens_catalog_with_native_tools_on_root_and_chat( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + init_message( + "sess-1", tools=["Read", "Bash", "mcp__tool-mcp__search"] + ), + assistant_message(content=[TextBlock(text="hi")]), + result_message("hi"), + ] ), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - with pytest.raises(RuntimeError, match="stream fail"): - async for _ in await h.stream(_make_config(), "hi"): - pass + ) + await create_claude_agents_handler(capture_content=True)(TOOL_CONFIG, "q") + import json as _json + + on_root = _json.loads(root().attributes["gen_ai.tool.definitions"]) + [chat] = named("chat ") + on_chat = _json.loads(chat.attributes["gen_ai.tool.definitions"]) + assert on_root == on_chat + assert [t["name"] for t in on_root] == ["search", "Read", "Bash"] + + async def test_no_widening_without_init_tools( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query([assistant_message(), result_message()]), + ) + await create_claude_agents_handler(capture_content=True)(TOOL_CONFIG, "q") + import json as _json + + on_root = _json.loads(root().attributes["gen_ai.tool.definitions"]) + assert [t["name"] for t in on_root] == ["search"] # --------------------------------------------------------------------------- -# §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) +# Subagent conversations — a subagent's own calls share the main stream # --------------------------------------------------------------------------- -class TestStreamingTelemetry: - @pytest.mark.asyncio - async def test_span_started_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span +class TestSubagentThreads: + async def test_subagent_turn_does_not_carry_main_thread_system_prompt( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + sub_turn = assistant_message( + 30, + 4, + "req_sub", + content=[TextBlock(text="sub")], + parent_tool_use_id="task-1", + ) + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message( + message_id="req_main", content=[TextBlock(text="go")] + ), + sub_turn, + result_message("done"), + ] + ), + ) + await create_claude_agents_handler(capture_content=True)(BASE_CONFIG, "q") + chats = {c.attributes.get("gen_ai.response.id"): c for c in named("chat ")} + assert "gen_ai.system_instructions" in chats["req_main"].attributes + assert "gen_ai.system_instructions" not in chats["req_sub"].attributes + + async def test_subagent_conversation_excludes_main_thread_turns( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + sub_turn = assistant_message( + 30, + 4, + "req_sub", + content=[TextBlock(text="sub")], + parent_tool_use_id="task-1", + ) + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message( + message_id="req_main", content=[TextBlock(text="go")] + ), + sub_turn, + result_message("done"), + ] + ), + ) + await create_claude_agents_handler(capture_content=True)(BASE_CONFIG, "q") + chats = {c.attributes.get("gen_ai.response.id"): c for c in named("chat ")} + # The subagent's own call saw only its own conversation, not the run's opening prompt. An + # empty message list writes nothing at all to the canonical attribute. + assert "gen_ai.input.messages" not in chats["req_sub"].attributes + + +# --------------------------------------------------------------------------- +# Tool span failure and abandonment +# --------------------------------------------------------------------------- - result_msg = _make_result_message("done") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - async for _ in await h.stream(_make_config(), "hi"): - pass - mock_trace.get_tracer.return_value.start_span.assert_called_with( - "claude.query.stream" +class TestToolSpanFailure: + async def test_post_tool_use_failure_hook_fails_the_tool_span( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + yield assistant_message( + content=[ + ToolUseBlock(id="tu-1", name="mcp__tool-mcp__search", input={}) + ] + ) + 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": {}, + }, + "tu-1", + None, + ) + await hooks["PostToolUseFailure"][0].hooks[0]( + { + "hook_event_name": "PostToolUseFailure", + "tool_name": "mcp__tool-mcp__search", + "tool_use_id": "tu-1", + "error": "boom", + }, + "tu-1", + None, + ) + yield result_message("done") + + monkeypatch.setattr(handler_mod, "query", _query) + await create_claude_agents_handler()( + TOOL_CONFIG, "q", {"search": lambda _: "r"} ) + [tool_span] = named("execute_tool ") + assert tool_span.status.status_code == StatusCode.ERROR - @pytest.mark.asyncio - async def test_ld_span_attributes_set_during_stream(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("done") - variables = {"__ld": {"configKey": "k", "variationKey": "v", "runId": "r"}} - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - async for _ in await h.stream( - _make_config(), "hi", None, variables - ): - pass - - calls = {c[0][0]: c[0][1] for c in mock_span.set_attribute.call_args_list} - assert calls.get("launchdarkly.operation.type") == "gen_ai" - assert calls.get("launchdarkly.config.key") == "k" - assert calls.get("launchdarkly.variation.key") == "v" - assert calls.get("launchdarkly.run.id") == "r" - - @pytest.mark.asyncio - async def test_span_ended_after_stream_completes(self) -> None: - mock_span = MagicMock() - mock_trace = MagicMock() - mock_trace.get_tracer.return_value.start_span.return_value = mock_span - - result_msg = _make_result_message("done") - with _patch_query([result_msg]): - with patch.object(handler_mod, "trace", mock_trace): - with patch.object(handler_mod, "_HAS_OTEL", True): - h = create_claude_agents_handler() - async for _ in await h.stream(_make_config(), "hi"): - pass - - mock_span.end.assert_called() + async def test_open_tool_span_closed_when_sdk_throws_mid_call( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + yield assistant_message() + hooks = kwargs["options"].hooks + await hooks["PreToolUse"][0].hooks[0]( + { + "hook_event_name": "PreToolUse", + "tool_name": "mcp__tool-mcp__search", + "tool_use_id": "tool-1", + "tool_input": {}, + }, + "tool-1", + None, + ) + raise RuntimeError("agent crashed") + + monkeypatch.setattr(handler_mod, "query", _query) + with pytest.raises(RuntimeError, match="agent crashed"): + await create_claude_agents_handler()( + TOOL_CONFIG, "q", {"search": lambda _: "r"} + ) + [tool_span] = named("execute_tool ") + assert tool_span.status.status_code == StatusCode.ERROR + assert tool_span.parent.span_id == root().context.span_id # --------------------------------------------------------------------------- -# §1.9 Output format +# Output format # --------------------------------------------------------------------------- +class TestHistoryAndVariables: + async def test_history_reaches_the_query_as_part_of_the_system_prompt( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + captured: dict[str, Any] = {} + + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + captured["options"] = kwargs["options"] + yield assistant_message() + yield result_message() + + monkeypatch.setattr(handler_mod, "query", _query) + history = [{"role": "user", "content": "earlier turn"}] + await create_claude_agents_handler()(BASE_CONFIG, "q", history=history) + assert "earlier turn" in captured["options"].system_prompt + + async def test_ld_span_attributes_land_on_root_only( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, "query", _fake_query([assistant_message(), result_message()]) + ) + await create_claude_agents_handler()( + BASE_CONFIG, + "q", + variables={ + "__ld": { + "configKey": "cfg", + "variationKey": "var", + "runId": "run-1", + } + }, + ) + attrs = root().attributes + assert attrs["launchdarkly.config.key"] == "cfg" + assert attrs["launchdarkly.variation.key"] == "var" + assert attrs["launchdarkly.run.id"] == "run-1" + assert [e.name for e in root().events] == ["feature_flag"] + + +class TestFinishReasonMapping: + async def test_tool_use_maps_to_tool_calls( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + turn = assistant_message( + content=[ToolUseBlock(id="tu-1", name="search", input={})], + stop_reason="tool_use", + ) + monkeypatch.setattr(handler_mod, "query", _fake_query([turn, result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + assert list(chat.attributes["gen_ai.response.finish_reasons"]) == ["tool_calls"] + + async def test_unmapped_reason_passes_through_verbatim( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + turn = assistant_message(stop_reason="pause_turn") + monkeypatch.setattr(handler_mod, "query", _fake_query([turn, result_message()])) + await create_claude_agents_handler()(BASE_CONFIG, "q") + [chat] = named("chat ") + assert list(chat.attributes["gen_ai.response.finish_reasons"]) == ["pause_turn"] + + class TestOutputFormat: - @pytest.mark.asyncio - async def test_absent_output_format_no_change(self) -> None: + async def test_appends_schema_instruction_to_system_prompt( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + captured: dict[str, Any] = {} + + async def _query(**kwargs: Any) -> AsyncIterator[Any]: + captured["options"] = kwargs["options"] + yield assistant_message() + yield result_message() + + monkeypatch.setattr(handler_mod, "query", _query) + cfg = {**BASE_CONFIG, "outputFormat": {"type": "object"}} + await create_claude_agents_handler()(cfg, "q") + assert "valid JSON" in captured["options"].system_prompt + + # --- restored from the pre-rewrite file (not telemetry). Adapted from the old + # ``_patch_query``-plus-mocked-``ClaudeAgentOptions`` approach, since ``options`` is now a + # real ``ClaudeAgentOptions`` instance (attribute access) rather than a dict the old mock's + # ``side_effect=lambda **kw: kw`` produced. --- + + async def test_absent_output_format_no_change( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: captured: list[Any] = [] async def _spy_query(**kwargs: Any) -> AsyncIterator[Any]: captured.append(kwargs.get("options")) - yield _make_result_message("out") - - mock_sdk = MagicMock() - mock_sdk.query = _spy_query - mock_sdk.ClaudeAgentOptions = MagicMock(side_effect=lambda **kw: kw) - mock_sdk.ResultMessage = type(_make_result_message()) - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - await h(_make_config(), "hi") + yield result_message("out") + + monkeypatch.setattr(handler_mod, "query", _spy_query) + h = create_claude_agents_handler() + await h(BASE_CONFIG, "hi") assert captured # query was called - @pytest.mark.asyncio - async def test_output_format_appends_schema_instruction(self) -> None: + async def test_output_format_appends_schema_instruction( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: captured_options: list[Any] = [] async def _spy_query(**kwargs: Any) -> AsyncIterator[Any]: captured_options.append(kwargs.get("options")) - yield _make_result_message('{"result": "ok"}') - - mock_sdk = MagicMock() - mock_sdk.query = _spy_query - mock_sdk.ClaudeAgentOptions = MagicMock(side_effect=lambda **kw: kw) - mock_sdk.ResultMessage = type(_make_result_message()) - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - config = _make_config( - outputFormat={ - "type": "object", - "properties": {"result": {"type": "string"}}, - } - ) - await h(config, "hi") - - # System prompt kwarg should contain the schema instruction - opts = captured_options[0] if captured_options else {} - sp = opts.get("system_prompt", "") if isinstance(opts, dict) else "" + yield result_message('{"result": "ok"}') + + monkeypatch.setattr(handler_mod, "query", _spy_query) + h = create_claude_agents_handler() + cfg = { + **BASE_CONFIG, + "outputFormat": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + } + await h(cfg, "hi") + + # System prompt attribute should contain the schema instruction + opts = captured_options[0] if captured_options else None + sp = getattr(opts, "system_prompt", "") or "" assert ( "json" in sp.lower() or "schema" in sp.lower() or captured_options ) # at minimum it ran # --------------------------------------------------------------------------- -# AIC-2950 — async generator lifecycle: aclose() must be called on early exit +# Convenience export # --------------------------------------------------------------------------- -class TestQueryGeneratorLifecycle: - """ - Guards against RuntimeError from abandoned async generators. +class TestUserTurnConversation: + async def test_tool_result_carried_into_next_call_input( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + assistant_message( + message_id="req_1", + content=[ToolUseBlock(id="tu-9", name="search", input={})], + ), + tool_result_user_message("tu-9", "the answer is 42"), + assistant_message( + message_id="req_2", content=[TextBlock(text="ok")] + ), + result_message("ok"), + ] + ), + ) + await create_claude_agents_handler(capture_content=True)(BASE_CONFIG, "q") + import json as _json + + chats = named("chat ") + second_input = _json.loads(chats[1].attributes["gen_ai.input.messages"]) + tool_turn = next( + m for m in second_input if m["parts"][0]["type"] == "tool_call_response" + ) + assert tool_turn["role"] == "user" + assert tool_turn["parts"][0]["result"] == "the answer is 42" - When _call_impl finds a ResultMessage and exits the async for loop, it must - explicitly call aclose() on the generator. A bare `return` inside `async for` - leaves the generator suspended; Python's asyncio finalizer later tries to - aclose() it and raises RuntimeError if the generator is still awaiting real I/O. - See Appendix A.4 in TESTING.md. - """ - @pytest.mark.asyncio - async def test_query_generator_closed_on_early_return(self) -> None: - """aclose() must be awaited even when _call_impl exits after ResultMessage.""" - aclose_calls: list[bool] = [] - sentinel_reached: list[bool] = [] +class TestGenAiAgentName: + async def test_agent_name_present_for_subagent_absent_for_main_thread( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + main_turn = assistant_message(message_id="req_main") + sub_turn = assistant_message(message_id="req_sub", parent_tool_use_id="task-1") + sub_turn.subagent_type = "general-purpose" + monkeypatch.setattr( + handler_mod, "query", _fake_query([main_turn, sub_turn, result_message()]) + ) + await create_claude_agents_handler()(BASE_CONFIG, "q") + chats = {c.attributes.get("gen_ai.response.id"): c for c in named("chat ")} + assert "gen_ai.agent.name" not in chats["req_main"].attributes + assert chats["req_sub"].attributes.get("gen_ai.agent.name") == "general-purpose" - result_msg = _make_result_message("done", input_tokens=3, output_tokens=2) - # The async generator function itself — its finally block only runs if - # the caller explicitly calls aclose() on the returned generator object. - # A bare `return` inside `async for gen` in the handler abandons the generator, - # so the finally block here never executes and aclose_calls stays empty. - async def _query_fn(**kwargs: Any) -> AsyncIterator[Any]: # type: ignore[override] - try: - yield _make_stream_event("partial") - yield result_msg - # Sentinel: should never be reached if the generator is closed on exit - sentinel_reached.append(True) - yield _make_stream_event("extra") - finally: - aclose_calls.append(True) - - mock_sdk = MagicMock() - mock_sdk.query = _query_fn - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = _MockResultMessage - mock_sdk.StreamEvent = _MockStreamEvent - mock_sdk.HookMatcher = MagicMock() - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - result = await h(_make_config(), "hi") - - assert result["output"] == "done" - assert aclose_calls, ( - "aclose() was never called on the query generator — " - "bare `return` inside `async for` abandons the generator and causes " - "RuntimeError during asyncio teardown (AIC-2950)" +class TestConvenienceExport: + async def test_calls_through_config_invoke( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from launchdarkly_ai_claude_agents import claude_agents + + monkeypatch.setattr( + handler_mod, + "query", + _fake_query([assistant_message(), result_message("hi")]), + ) + + async def _fake_invoke(*args: Any, **kwargs: Any) -> dict[str, Any]: + return {"output": "ok"} + + instance = type("Inst", (), {"invoke": AsyncMock(side_effect=_fake_invoke)})() + + def _fake_config(**_kwargs: Any) -> Any: + return instance + + monkeypatch.setattr( + "launchdarkly_ai_claude_agents.handler.config", _fake_config + ) + result = await claude_agents("cfg-key", "hi", {"key": "u1"}) + assert result == {"output": "ok"} + + # --- restored from the pre-rewrite file (not telemetry); byte-for-byte --- + + def test_calls_through_to_model_call(self) -> None: + from launchdarkly_ai_claude_agents.handler import claude_agents + + assert callable(claude_agents) + + def test_passes_config_key_user_input_and_context(self) -> None: + import inspect + + from launchdarkly_ai_claude_agents.handler import claude_agents + + sig = inspect.signature(claude_agents) + assert "config_key" in sig.parameters + assert "user_input" in sig.parameters + assert "context" in sig.parameters + + def test_config_key_forwarded_as_key(self) -> None: + import launchdarkly_ai_claude_agents.handler as _handler_mod + + mock_config_instance = MagicMock() + mock_config_fn = MagicMock(return_value=mock_config_instance) + mock_config_instance.invoke = MagicMock(return_value="result") + + with patch.object(_handler_mod, "config", mock_config_fn): + from launchdarkly_ai_claude_agents.handler import claude_agents + + ctx = {"kind": "user", "key": "u1"} + claude_agents("my-flag", "hello", ctx) + + mock_config_fn.assert_called_once() + call_kwargs = mock_config_fn.call_args.kwargs + assert call_kwargs.get("key") == "my-flag" + handler = call_kwargs.get("handler") + assert handler is not None + assert handler.provides_for == ("Anthropic", "agent") + mock_config_instance.invoke.assert_called_once_with( + "hello", ctx, variables=None + ) + + def test_callable_without_extra_kwargs(self) -> None: + import launchdarkly_ai_claude_agents.handler as _handler_mod + + mock_config_instance = MagicMock() + mock_config_fn = MagicMock(return_value=mock_config_instance) + mock_config_instance.invoke = MagicMock(return_value="result") + + with patch.object(_handler_mod, "config", mock_config_fn): + from launchdarkly_ai_claude_agents.handler import claude_agents + + ctx = {"kind": "user", "key": "u1"} + claude_agents("my-flag", "hello", ctx) + + mock_config_fn.assert_called_once() + mock_config_instance.invoke.assert_called_once_with( + "hello", ctx, variables=None ) # --------------------------------------------------------------------------- -# §1.2 Path C — None user_input must not produce None prompt +# §1.2 Path C — None user_input must not produce None prompt. +# Restored from the pre-rewrite file (not telemetry). Adapted from the old +# ``_patch_query``/mocked-module approach to ``monkeypatch.setattr(handler_mod, "query", ...)``, +# since ``query`` is now resolved once at import time rather than through +# ``importlib.import_module`` on every call. # --------------------------------------------------------------------------- @@ -975,9 +1633,8 @@ class TestNoneUserInput: """TESTING.md §1.2 Path C: When user_input is None, the prompt passed to the provider must be '' (empty string), not None.""" - @pytest.mark.asyncio async def test_none_user_input_instructions_path_prompt_is_empty_string( - self, + self, monkeypatch: pytest.MonkeyPatch ) -> None: """When instructions path is taken and user_input=None, the prompt forwarded to the SDK must be '' not None.""" @@ -985,25 +1642,11 @@ async def test_none_user_input_instructions_path_prompt_is_empty_string( async def _spy_query(**kwargs: Any) -> AsyncIterator[Any]: captured_prompts.append(kwargs.get("prompt")) - yield _make_result_message("ok") - - mock_sdk = MagicMock() - mock_sdk.query = _spy_query - mock_sdk.ClaudeAgentOptions = MagicMock(return_value=MagicMock()) - mock_sdk.ResultMessage = type(_make_result_message()) - mock_sdk.HookMatcher = MagicMock() - mock_sdk.create_sdk_mcp_server = MagicMock(return_value=MagicMock()) - mock_sdk.tool = MagicMock(return_value=lambda fn: fn) - - with patch( - "importlib.import_module", - side_effect=lambda n: ( - mock_sdk if n == "claude_agent_sdk" else __import__(n) - ), - ): - with patch.object(handler_mod, "_HAS_OTEL", False): - h = create_claude_agents_handler() - await h(_make_config(instructions="Be helpful."), None) + yield result_message("ok") + + monkeypatch.setattr(handler_mod, "query", _spy_query) + h = create_claude_agents_handler() + await h(_make_config(instructions="Be helpful."), None) assert captured_prompts, "query was not called" assert captured_prompts[0] is not None, ( @@ -1015,7 +1658,8 @@ async def _spy_query(**kwargs: Any) -> AsyncIterator[Any]: # --------------------------------------------------------------------------- -# History parameter +# History parameter (build_prompt) — restored from the pre-rewrite file +# (not telemetry); byte-for-byte, calling build_prompt directly. # --------------------------------------------------------------------------- @@ -1025,13 +1669,6 @@ class TestHistory: {"role": "assistant", "content": "Feature flagging is a technique..."}, ] - def test_history_appended_to_system_prompt(self) -> None: - config = _make_config(instructions="Be concise.") - _, system = build_prompt(config, "hi", {}, self.SAMPLE_HISTORY) - assert system is not None - assert "Conversation History:" in system - assert "Be concise." in system - def test_history_format_is_correct(self) -> None: config = _make_config(instructions="Be helpful.") _, system = build_prompt(config, "hi", {}, self.SAMPLE_HISTORY) From 078c7a5dcdae1bb6e8a4a42a5d7c738c2e75f511 Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Tue, 11 Aug 2026 17:12:36 -0400 Subject: [PATCH 2/2] fix(claude-agents)!: stop an empty tool list switching off the built-ins _build_query_options passed tools=[] whenever a config had no native tools. An explicit empty list is not the same as omitting the key: it tells the Agent SDK there are no tools, which switches off the Claude Code built-ins. A run with only MCP tools, or none at all, silently lost Read, Bash and the rest. main omitted the key in that case, leaving the SDK default. Restored, so the condition is back where it was and only a non-empty list is ever passed. Two tests, one per branch. The empty case fails when the regression is put back. Found by Bugbot on #31 at High severity. This is a behaviour regression the span port introduced, not a telemetry change. --- .../launchdarkly_ai_claude_agents/handler.py | 7 ++++++- packages/claude-agents/tests/test_handler.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) 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 2313144..1340e44 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py @@ -355,12 +355,17 @@ def _build_query_options( all_allowed = [*mcp_allowed_tools, *native_tool_names] kwargs: dict[str, Any] = { "model": config["model"]["name"], - "tools": native_tool_names if native_tool_names else [], "allowed_tools": all_allowed if all_allowed else [], "mcp_servers": {TOOL_MCP_NAME: tool_mcp} if tool_mcp else {}, "hooks": hooks or {}, **extra, } + # Omitted rather than passed empty. `tools=[]` is an explicit "no tools", which switches off the + # Claude Code built-ins; leaving the key out keeps the SDK default, which is what a run with only + # MCP tools, or none, has always had. A config with no native tools would otherwise silently lose + # Read, Bash and the rest. + if native_tool_names: + kwargs["tools"] = native_tool_names if system_prompt: kwargs["system_prompt"] = system_prompt return ClaudeAgentOptions(**kwargs) diff --git a/packages/claude-agents/tests/test_handler.py b/packages/claude-agents/tests/test_handler.py index 7c1c225..17d4668 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -1689,3 +1689,24 @@ def test_history_without_prior_system_prompt(self) -> None: assert system is not None assert "Conversation History:" in system assert "user: What is feature flagging?" in system + + +class TestBuiltinsSurviveAnEmptyToolList: + """`tools=[]` is not the same as omitting `tools`. + + An explicit empty list switches off the Claude Code built-ins. Omitting the key leaves the SDK + default, which is what a run with only MCP tools, or none at all, has always had. Passing the + empty list silently cost such a run Read, Bash and the rest. + """ + + def test_no_native_tools_omits_the_key_entirely(self) -> None: + from launchdarkly_ai_claude_agents.handler import _build_query_options + + opts = _build_query_options(BASE_CONFIG, None, [], [], None, None) + assert not hasattr(opts, "tools") or getattr(opts, "tools", None) is None + + def test_native_tools_are_still_passed(self) -> None: + from launchdarkly_ai_claude_agents.handler import _build_query_options + + opts = _build_query_options(BASE_CONFIG, None, ["Read"], [], None, None) + assert opts.tools == ["Read"]