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..1340e44 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,54 @@ 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"], + "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) # --------------------------------------------------------------------------- @@ -177,9 +376,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 +390,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 +523,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 +540,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 +617,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 +703,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..17d4668 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 == "" + # --------------------------------------------------------------------------- -# §1.3 Tool conversion +# 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 — 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 - 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" +# --------------------------------------------------------------------------- +# Tool span failure and abandonment +# --------------------------------------------------------------------------- + + +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 - 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. - """ + 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" - @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] = [] - result_msg = _make_result_message("done", input_tokens=3, output_tokens=2) +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" - # 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) @@ -1052,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"]