"""Client instrumentation for Claude Agent SDK.""" import logging import time import weakref from collections.abc import AsyncGenerator, AsyncIterable from datetime import datetime, timezone from typing import Any, Optional from langsmith._internal import _context from langsmith._internal._package_version import get_package_version from langsmith.run_helpers import get_current_run_tree, trace from ._config import get_tracing_config from ._hooks import ( SessionState, _current_session, _register_session, _set_session_root, _unregister_session, clear_active_tool_runs, get_subagent_run_by_tool_id, post_tool_use_failure_hook, post_tool_use_hook, pre_tool_use_hook, subagent_start_hook, subagent_stop_hook, ) from ._messages import ( build_llm_input, flatten_content_blocks, unwrap_message_dicts, ) from ._tools import ( clear_parent_run_tree, get_parent_run_tree, set_parent_run_tree, ) from ._transcripts import LLM_RUN_NAME, reconcile_from_transcripts from ._usage import extract_usage_metadata logger = logging.getLogger(__name__) TRACE_CHAIN_NAME = "claude.conversation" class TurnLifecycle: """Track ongoing model runs so consecutive messages are recorded correctly. The Claude Agent SDK may deliver a single assistant turn as multiple ``AssistantMessage`` events (e.g. one with ``ThinkingBlock``, another with ``TextBlock``/``ToolUseBlock``). Messages that share the same ``message_id`` are accumulated into a single LLM run. """ def __init__(self, query_start_time: Optional[float] = None): self.current_run: Optional[Any] = None self.current_message_id: Optional[str] = None self.next_start_time: Optional[float] = query_start_time # message_id → RunTree for all LLM runs created this conversation. # Used to retroactively set usage from transcripts. self.llm_runs_by_message_id: dict[str, Any] = {} # Runs that have been end()ed but not yet patch()ed. # Deferred so transcript usage can be set before the single patch(). self._pending_patch: list[Any] = [] def start_llm_run( self, message: Any, prompt: Any, history: list[dict[str, Any]], parent: Optional[Any] = None, ) -> Optional[dict[str, Any]]: """Begin or continue a model run for *message*. If *message* has the same ``message_id`` as the current run the output is appended; otherwise a new run is started (ending any previous one first). """ message_id = getattr(message, "message_id", None) start = self.next_start_time or time.time() # Same turn – just accumulate the output blocks and update usage. # Return None so the caller does NOT append a duplicate history # entry; the original entry in ``history`` is updated in place. if message_id and message_id == self.current_message_id and self.current_run: content = flatten_content_blocks(getattr(message, "content", None)) if content and self.current_run.outputs: prev = self.current_run.outputs.get("content", []) if isinstance(prev, list) and isinstance(content, list): merged = prev + content self.current_run.outputs["content"] = merged # Update the existing history entry in place so # subsequent LLM runs see a single merged message. for entry in reversed(history): if entry.get("role") == "assistant": entry["content"] = merged break elif isinstance(content, list): self.current_run.outputs["content"] = content self._set_usage_from_message(message, self.current_run) return None # Different turn – end previous but defer patch() until # transcript usage is available. if self.current_run: self.current_run.end() self._pending_patch.append(self.current_run) final_output, run = begin_llm_run_from_assistant_messages( [message], prompt, history, start_time=start, parent=parent ) self.current_run = run self.current_message_id = message_id self.next_start_time = None if run: if message_id: self.llm_runs_by_message_id[message_id] = run self._set_usage_from_message(message, run) return final_output @staticmethod def _set_usage_from_message(message: Any, run: Any) -> None: """Set usage metadata on a run from a live AssistantMessage. Always overwrites — later chunks in the same turn have more accurate counts. Transcript-based usage will overwrite again if available. """ raw_usage = getattr(message, "usage", None) if not raw_usage: return usage_meta = extract_usage_metadata(raw_usage) if usage_meta: meta = run.extra.setdefault("metadata", {}) meta["usage_metadata"] = usage_meta def mark_next_start(self) -> None: """Mark when the next assistant message will start.""" self.next_start_time = time.time() def close(self) -> None: """End any open run and add to pending patch list.""" if self.current_run: self.current_run.end() self._pending_patch.append(self.current_run) self.current_run = None def flush(self) -> None: """Patch all deferred LLM runs. Call after usage has been set.""" for run in self._pending_patch: try: run.patch() except Exception as e: logger.warning(f"Failed to patch LLM run: {e}") self._pending_patch.clear() def begin_llm_run_from_assistant_messages( messages: list[Any], prompt: Any, history: list[dict[str, Any]], start_time: Optional[float] = None, parent: Optional[Any] = None, ) -> tuple[Optional[dict[str, Any]], Optional[Any]]: """Create a traced model run from assistant messages.""" if not messages or type(messages[-1]).__name__ != "AssistantMessage": return None, None last_msg = messages[-1] model = getattr(last_msg, "model", None) if parent is None: parent = get_parent_run_tree() or get_current_run_tree() if not parent: return None, None inputs = build_llm_input(prompt, history) outputs = [ {"content": flatten_content_blocks(m.content), "role": "assistant"} for m in messages if hasattr(m, "content") ] llm_metadata: dict[str, Any] = {"ls_provider": "anthropic"} if model: llm_metadata["ls_model_name"] = model llm_run = parent.create_child( name=LLM_RUN_NAME, run_type="llm", inputs={"messages": inputs} if inputs else {}, extra={"metadata": llm_metadata}, start_time=datetime.fromtimestamp(start_time, tz=timezone.utc) if start_time else None, ) try: llm_run.post() except Exception as e: logger.warning(f"Failed to post LLM run: {e}") # Set outputs after posting so they are sent with end_time on the patch. llm_run.outputs = outputs[-1] if len(outputs) == 1 else {"content": outputs} final_content = ( {"content": flatten_content_blocks(last_msg.content), "role": "assistant"} if hasattr(last_msg, "content") else None ) return final_content, llm_run def _bind_hook_to_session(hook: Any, session: Optional[SessionState]) -> Any: """Return a hook callable that runs with *session* bound, if provided.""" if session is None: return hook async def _bound(input_data: Any, tool_use_id: Any, context: Any) -> Any: token = _current_session.set(session) try: return await hook(input_data, tool_use_id, context) finally: _current_session.reset(token) return _bound def _inject_tracing_hooks(options: Any, session: Optional[SessionState] = None) -> None: """Inject LangSmith tracing hooks into ClaudeAgentOptions. If *session* is provided, injected hook callables bind that session around each hook invocation. This is important because the Claude SDK may execute hooks in async contexts that do not inherit the ``receive_response`` ContextVar; binding at hook injection time keeps each client isolated. """ if not hasattr(options, "hooks"): return # Initialize hooks dict if not present if options.hooks is None: options.hooks = {} for event in ( "PreToolUse", "PostToolUse", "PostToolUseFailure", "SubagentStart", "SubagentStop", ): if event not in options.hooks: options.hooks[event] = [] try: from claude_agent_sdk import HookMatcher # type: ignore[import-not-found] langsmith_pre_matcher = HookMatcher( matcher=None, hooks=[_bind_hook_to_session(pre_tool_use_hook, session)] ) langsmith_post_matcher = HookMatcher( matcher=None, hooks=[_bind_hook_to_session(post_tool_use_hook, session)] ) langsmith_failure_matcher = HookMatcher( matcher=None, hooks=[_bind_hook_to_session(post_tool_use_failure_hook, session)], ) langsmith_subagent_start_matcher = HookMatcher( matcher=None, hooks=[_bind_hook_to_session(subagent_start_hook, session)] ) langsmith_subagent_stop_matcher = HookMatcher( matcher=None, hooks=[_bind_hook_to_session(subagent_stop_hook, session)] ) options.hooks["PreToolUse"].insert(0, langsmith_pre_matcher) options.hooks["PostToolUse"].insert(0, langsmith_post_matcher) options.hooks["PostToolUseFailure"].insert(0, langsmith_failure_matcher) options.hooks["SubagentStart"].insert(0, langsmith_subagent_start_matcher) options.hooks["SubagentStop"].insert(0, langsmith_subagent_stop_matcher) logger.debug("Injected LangSmith tracing hooks into ClaudeAgentOptions") except ImportError: logger.warning("Failed to import HookMatcher from claude_agent_sdk") except Exception as e: logger.warning(f"Failed to inject tracing hooks: {e}") def _wrap_tool_handler( original_handler: Any, session: Optional[SessionState] = None, tool_name: Optional[str] = None, ) -> Any: """Wrap an MCP tool handler to propagate LangSmith run context. The Claude SDK runs hooks and tool handlers in different async task contexts, so contextvars set in ``PreToolUse`` are invisible to the handler. This wrapper copies the active tool run into the contextvar before calling the original handler, so ``@traceable`` calls inside the handler nest correctly. """ async def _wrapped(args: Any) -> Any: # The most recently added active tool run is the one PreToolUse just # created for this invocation. Prefer an explicitly bound client # session because tool handlers may run in an async context that did # not inherit _current_session. tool_run = _get_last_active_tool_run(session, args=args, tool_name=tool_name) if tool_run: token = _context._PARENT_RUN_TREE_REF.set(weakref.ref(tool_run)) session_token = ( _current_session.set(session) if session is not None else None ) try: return await original_handler(args) finally: if session_token is not None: _current_session.reset(session_token) _context._PARENT_RUN_TREE_REF.reset(token) return await original_handler(args) _wrapped._langsmith_wrapped = True # type: ignore[attr-defined] _wrapped._langsmith_original_handler = original_handler # type: ignore[attr-defined] _wrapped._langsmith_session = session # type: ignore[attr-defined] _wrapped._langsmith_tool_name = tool_name # type: ignore[attr-defined] return _wrapped def _tool_run_matches(run: Any, args: Any, tool_name: Optional[str]) -> bool: """Return whether *run* appears to be for this SDK MCP handler call. Matching is intentionally strict: we require both the tool name and the handler args to line up with what the ``PreToolUse`` hook recorded. This avoids cross-attributing a handler invocation to the wrong client's active tool run under concurrency. """ if not tool_name: return False run_name = str(getattr(run, "name", "")) # SDK MCP tools show up in hook data as e.g. ``mcp__weather__get_weather`` # while the handler only knows its short name ``get_weather``. name_matches = ( tool_name == run_name or tool_name in run_name or run_name in tool_name ) if not name_matches: return False inputs = getattr(run, "inputs", None) if not isinstance(inputs, dict): return False # PreToolUse stores {} when the tool had no inputs, otherwise # {"input": }. Normalise both sides before comparing. recorded = inputs.get("input", {}) if inputs else {} return recorded == (args or {}) def _newest_matching_tool_run( sessions: list[SessionState], args: Any, tool_name: Optional[str] ) -> Any: """Return the most recently created active tool run that matches.""" candidates: list[tuple[float, Any]] = [] for candidate_session in sessions: for run, start_time in candidate_session.active_tool_runs.values(): if _tool_run_matches(run, args, tool_name): candidates.append((start_time, run)) if not candidates: return None return max(candidates, key=lambda item: item[0])[1] def _get_last_active_tool_run( session: Optional[SessionState] = None, *, args: Any = None, tool_name: Optional[str] = None, ) -> Any: """Return the active tool run for an SDK MCP handler, or None. Lookup order: 1. The session explicitly bound to the handler (if any). 2. The session bound to the current ContextVar. 3. The module-level default session (unit tests / unbound callers). 4. Any live client session — only used when the handler is unbound and the SDK invoked it in a detached async context. Requires an exact tool name + args match to avoid cross-attribution across clients. """ from ._hooks import ( _current_session, _current_session_or_default, _registered_sessions, ) # If we have a specific session (explicitly bound, current-context, or the # test default), just return its newest active tool run. There is no # cross-client ambiguity at that point. def _newest_in(s: SessionState) -> Any: if not s.active_tool_runs: return None latest_id = max( s.active_tool_runs, key=lambda tid: s.active_tool_runs[tid][1], ) return s.active_tool_runs[latest_id][0] if session is not None: return _newest_in(session) current_session = _current_session.get() if current_session is not None: run = _newest_in(current_session) if run is not None: return run default_session = _current_session_or_default() if default_session is not current_session: run = _newest_in(default_session) if run is not None: return run # Last resort: the SDK invoked this handler in a detached async context and # the handler object wasn't bound to a session. Require strict tool-name + # args match so concurrent clients can't steal each other's attribution. return _newest_matching_tool_run(_registered_sessions(), args, tool_name) def instrument_claude_client(original_class: Any) -> None: """Patch ``ClaudeSDKClient`` **in place** to trace calls. In-place patching (rather than subclassing + reference replacement) ensures that callers who imported ``ClaudeSDKClient`` *before* ``configure_claude_agent_sdk()`` was called still get instrumented. """ if getattr(original_class, "_langsmith_instrumented", False): return # Already wrapped, avoid double-tracing # ── stash originals ────────────────────────────────────────────── _orig_init = original_class.__init__ _orig_query = original_class.query _orig_receive_response = original_class.receive_response # ── patched __init__ ───────────────────────────────────────────── def _traced_init(self: Any, *args: Any, **kwargs: Any) -> None: options = kwargs.get("options") or (args[0] if args else None) self._ls_session = SessionState() if options: _inject_tracing_hooks(options, self._ls_session) _orig_init(self, *args, **kwargs) self._ls_prompt = None self._ls_start_time = None self._ls_streamed_input = None # ── patched query ──────────────────────────────────────────────── async def _traced_query(self: Any, *args: Any, **kwargs: Any) -> Any: self._ls_start_time = time.time() self._ls_streamed_input = None prompt = args[0] if args else kwargs.get("prompt") if prompt is None: pass elif isinstance(prompt, str): self._ls_prompt = prompt elif isinstance(prompt, AsyncIterable): collector: list[dict[str, Any]] = [] self._ls_streamed_input = collector self._ls_prompt = None async def _gen_wrapper() -> AsyncGenerator[dict[str, Any], None]: async for msg in prompt: collector.append(msg) yield msg if args: args = (_gen_wrapper(),) + args[1:] else: kwargs["prompt"] = _gen_wrapper() else: self._ls_prompt = str(prompt) return await _orig_query(self, *args, **kwargs) # ── patched receive_response ───────────────────────────────────── async def _traced_receive_response(self: Any) -> AsyncGenerator[Any, None]: messages = _orig_receive_response(self) trace_inputs: dict[str, Any] = {} trace_metadata: dict[str, Any] = { "ls_integration": "claude-agent-sdk", "ls_integration_version": get_package_version("claude_agent_sdk"), } awaiting_streamed_input = self._ls_streamed_input is not None if self._ls_prompt: trace_inputs["prompt"] = self._ls_prompt if hasattr(self, "options") and self.options: if hasattr(self.options, "system_prompt") and self.options.system_prompt: system_prompt = self.options.system_prompt if isinstance(system_prompt, str): trace_inputs["system"] = system_prompt elif isinstance(system_prompt, dict): if system_prompt.get("type") == "preset": preset_text = ( f"preset: {system_prompt.get('preset', 'claude_code')}" ) if "append" in system_prompt: preset_text += f"\nappend: {system_prompt['append']}" trace_inputs["system"] = preset_text else: trace_inputs["system"] = system_prompt for attr in ["model", "permission_mode", "max_turns"]: if hasattr(self.options, attr): val = getattr(self.options, attr) if val is not None: trace_metadata[attr] = val config = get_tracing_config() user_metadata = config.get("metadata") or {} trace_kwargs: dict[str, Any] = { "name": config.get("name") or TRACE_CHAIN_NAME, "run_type": "chain", "inputs": trace_inputs, "metadata": { **trace_metadata, **user_metadata, "ls_agent_type": "root", }, } if config.get("project_name"): trace_kwargs["project_name"] = config["project_name"] if config.get("tags"): trace_kwargs["tags"] = config["tags"] async with trace(**trace_kwargs) as run: # Bind this client's state container to the ContextVar so stream # helpers on this SDK event loop pick it up (see # _hooks.SessionState). This keeps concurrent ClaudeSDKClient # instances — eval runs, FastAPI handlers, Celery workers, # asyncio.gather — from corrupting each other's correlation state. session = getattr(self, "_ls_session", None) if session is None: session = self._ls_session = SessionState() session_token = _register_session(session) _set_session_root(session, run) parent_token = set_parent_run_tree(run) tracker = TurnLifecycle(self._ls_start_time) collected_by_ctx: dict[Optional[str], list[dict[str, Any]]] = {None: []} prompt_for_llm: Any = self._ls_prompt try: async for msg in messages: if awaiting_streamed_input and self._ls_streamed_input: unwrapped_messages = unwrap_message_dicts( self._ls_streamed_input ) if unwrapped_messages: run.inputs["messages"] = unwrapped_messages prompt_for_llm = self._ls_streamed_input awaiting_streamed_input = False msg_type = type(msg).__name__ if msg_type == "AssistantMessage": parent_tool_use_id = getattr(msg, "parent_tool_use_id", None) llm_parent = ( get_subagent_run_by_tool_id(parent_tool_use_id) if parent_tool_use_id else None ) ctx_key = parent_tool_use_id ctx_history = collected_by_ctx.setdefault(ctx_key, []) content = tracker.start_llm_run( msg, prompt_for_llm if parent_tool_use_id is None else None, ctx_history, parent=llm_parent, ) if content: ctx_history.append(content) elif msg_type == "UserMessage": parent_tool_use_id = getattr(msg, "parent_tool_use_id", None) ctx_key = parent_tool_use_id ctx_history = collected_by_ctx.setdefault(ctx_key, []) if hasattr(msg, "content"): flattened = flatten_content_blocks(msg.content) if ( isinstance(flattened, list) and flattened and isinstance(flattened[0], dict) and flattened[0].get("type") == "tool_result" ): for block in flattened: tool_use_id = block.get("tool_use_id") ctx_history.append( { "role": "tool", "content": block.get("content", ""), "tool_call_id": tool_use_id, } ) if ( tool_use_id and tool_use_id in session.active_tool_runs ): tool_run, _ = session.active_tool_runs.pop( tool_use_id ) result_content = block.get("content", "") is_error = block.get("is_error", False) tool_run.end( outputs={"output": result_content}, error=str(result_content) if is_error else None, ) try: tool_run.patch() except Exception as e: logger.warning( "Failed to patch" f" orphaned tool run: {e}" ) else: ctx_history.append( { "content": flattened, "role": "user", } ) tracker.mark_next_start() elif msg_type == "ResultMessage": session_id_val = getattr(msg, "session_id", None) meta = { k: v for k, v in { "num_turns": getattr(msg, "num_turns", None), "session_id": session_id_val, "thread_id": session_id_val, "duration_ms": getattr(msg, "duration_ms", None), "duration_api_ms": getattr( msg, "duration_api_ms", None ), "is_error": getattr(msg, "is_error", None), }.items() if v is not None } if meta: run.metadata.update(meta) yield msg main_collected = collected_by_ctx.get(None, []) run.end(outputs=main_collected[-1] if main_collected else None) except Exception: logger.exception("Error while tracing Claude Agent stream") finally: tracker.close() reconcile_from_transcripts(tracker, session=session) tracker.flush() clear_parent_run_tree(parent_token) try: clear_active_tool_runs(session) finally: _unregister_session(session, session_token) # ── apply patches to the class itself ──────────────────────────── original_class.__init__ = _traced_init original_class.query = _traced_query original_class.receive_response = _traced_receive_response original_class._langsmith_instrumented = True def instrument_sdk_mcp_tool(tool_class: Any) -> None: """Patch ``SdkMcpTool.__init__`` to auto-wrap handlers. Wrapping happens at construction time so that any tool created *after* ``configure_claude_agent_sdk()`` automatically gets run-context propagation, regardless of how ``tool`` or ``create_sdk_mcp_server`` were imported. """ if getattr(tool_class, "_langsmith_handler_patched", False): return _orig_init = tool_class.__init__ def _patched_init(self: Any, *args: Any, **kwargs: Any) -> None: _orig_init(self, *args, **kwargs) handler = self.handler if callable(handler) and not getattr(handler, "_langsmith_wrapped", False): self.handler = _wrap_tool_handler( handler, tool_name=getattr(self, "name", None) ) tool_class.__init__ = _patched_init tool_class._langsmith_handler_patched = True