diff --git a/AGENTS.md b/AGENTS.md index cc8033c51e69..3e7c7dfad7ae 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -439,6 +439,13 @@ pnpm exec tsx examples/claudeAgents.ts - `Turn.messages` stores input messages, while `Turn.output_messages` stores the terminal agent response. `Turn.record(messages=..., output_messages=...)` replaces the two lists independently. +- Mypy requires explicit `return None` paths in functions annotated with + `T | None`; bare `return` and implicit fallthrough trigger return-value + errors. +- Python `Turn` spans are same-process `invoke_agent` operations and use OTel's + default `SpanKind.INTERNAL`. Do not set `gen_ai.provider.name` on a `Turn`; + set it on child `LLM`/`chat` spans, because one turn can use multiple + providers. - Keep streaming and batch paths aligned: `Turn._build_attrs()` must apply content gating and PII redaction to both lists before passing them to `invoke_agent_attributes()`, and `log_turn()` must accept both fields. @@ -457,6 +464,15 @@ pnpm exec tsx examples/claudeAgents.ts rollups. - Regression coverage must exercise both the calls-based and OTel integrations with nonzero cache-read and cache-creation counts. +- The OTel-selected Claude Agent SDK path composes the Python GenAI + `Conversation`, `Turn`, `LLM`, and `Tool` handles; it must not create raw + OTel spans or call the low-level GenAI attribute builders itself. Use + `set_attributes()` only for semantic fields the typed handles do not expose. + The legacy calls-based path remains separate. +- Stream adapters can create child handles when work starts, then enter them + with a normal `with` when the completion message arrives. Preserve logical + timing with `LLM.started_at` and an explicit `Tool.started_at` instead of + keeping contexts open with `ExitStack`. ### Credential-shaped fields in call inputs and attributes diff --git a/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py b/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py index c51a0434a50a..8012a0110aba 100644 --- a/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py +++ b/tests/integrations/claude_agent_sdk/claude_agent_sdk_otel_test.py @@ -19,13 +19,14 @@ from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from opentelemetry.trace import StatusCode +from opentelemetry.trace import SpanKind, StatusCode from tests.integrations.claude_agent_sdk.conftest import ReplayTransport, load_cassette from weave.conversation import agent_name_override from weave.integrations.claude_agent_sdk.otel_integration import ( get_claude_agent_sdk_otel_patcher, ) +from weave.trace.settings import override_settings @pytest.fixture @@ -55,6 +56,16 @@ def patch_claude_agent_sdk_otel() -> Generator[None]: patcher.undo_patch() +@pytest.fixture(autouse=True) +def disable_capture_info() -> Generator[None]: + """Keep exact span payload assertions independent of host metadata.""" + with override_settings( + capture_client_info=False, + capture_system_info=False, + ): + yield + + # --- helpers ---------------------------------------------------------------- @@ -65,7 +76,7 @@ def get_attrs(span: Any) -> dict[str, Any]: def check_integration_and_strip(attrs: dict[str, Any]) -> dict[str, Any]: """Assert + remove the flattened integration.* provenance keys. - The agent OTel processor stamps integration provenance on every span; pop it + Conversation attributes stamp integration provenance on every span; pop it here so the exact-shape assertions below stay focused on the GenAI semconv keys. """ assert attrs["integration.name"] == "claude_agent_sdk" @@ -116,6 +127,7 @@ async def run_query(cassette: str, prompt: str) -> None: async def test_simple_text_query_otel(otel_spans: InMemorySpanExporter) -> None: await run_query("simple_text_response", "What is 2+2?") spans = otel_spans.get_finished_spans() + assert {span.instrumentation_scope.name for span in spans} == {"weave.conversation"} agent_spans = get_spans_by_op(spans, "invoke_agent") chat_spans = get_spans_by_op(spans, "chat") @@ -126,10 +138,10 @@ async def test_simple_text_query_otel(otel_spans: InMemorySpanExporter) -> None: # glance and any spec drift (added/removed/renamed keys) fails the test. agent_span = agent_spans[0] assert agent_span.name == "invoke_agent claude_agent_sdk" + assert agent_span.kind == SpanKind.INTERNAL assert check_integration_and_strip(get_attrs(agent_span)) == { "gen_ai.operation.name": "invoke_agent", "gen_ai.agent.name": "claude_agent_sdk", - "gen_ai.provider.name": "anthropic", "gen_ai.conversation.id": "s-abc123", "gen_ai.request.model": "claude-sonnet-4-6", "gen_ai.input.messages": ( @@ -358,8 +370,7 @@ async def test_custom_agent_name_query_otel(otel_spans: InMemorySpanExporter) -> assert agent_span.name == "invoke_agent research_agent" attrs = check_integration_and_strip(get_attrs(agent_span)) assert attrs["gen_ai.agent.name"] == "research_agent" - # Only the agent name is overridden — the rest of the GenAI shape is intact. - assert attrs["gen_ai.provider.name"] == "anthropic" + # Only the agent name is overridden — the model remains intact. assert attrs["gen_ai.request.model"] == "claude-sonnet-4-6" diff --git a/weave/integrations/claude_agent_sdk/otel_integration.py b/weave/integrations/claude_agent_sdk/otel_integration.py index 51c91b9a22a0..0bb1b7ed26fa 100644 --- a/weave/integrations/claude_agent_sdk/otel_integration.py +++ b/weave/integrations/claude_agent_sdk/otel_integration.py @@ -1,8 +1,8 @@ -"""Weave OTel tracing for the Claude Agent SDK. +"""Weave GenAI agent tracing for the Claude Agent SDK. -Emits OpenTelemetry GenAI spans to Weave's Agents tab; the sibling -``claude_agent_sdk_integration.py`` emits legacy Weave calls instead. The -dispatcher selects this variant when ``WEAVE_USE_OTEL_V2`` is set. +Uses the Weave Python GenAI agent SDK to emit spans to Weave's Agents tab; the +sibling ``claude_agent_sdk_integration.py`` emits legacy Weave calls instead. +The dispatcher selects this variant when ``WEAVE_USE_OTEL_V2`` is set. Each ``query()`` call / ``ClaudeSDKClient`` turn becomes an ``invoke_agent`` span, with a child ``chat`` span per model response and an ``execute_tool`` @@ -23,6 +23,7 @@ import logging from collections.abc import AsyncIterator from dataclasses import dataclass, field +from datetime import datetime, timezone from functools import wraps from typing import Any @@ -36,17 +37,10 @@ ToolUseBlock, UserMessage, ) -from opentelemetry import context as otel_context -from opentelemetry import trace as otel_trace -from opentelemetry.trace import StatusCode +from weave.conversation import LLM, Conversation, Message, Reasoning, Tool, Turn, Usage from weave.conversation.agent_context import resolve_agent_name -from weave.conversation.conversation_otel import ( - execute_tool_attributes, - invoke_agent_attributes, - llm_attributes, -) -from weave.conversation.types import Message, Reasoning, ToolCallPart, Usage +from weave.conversation.types import ToolCallPart from weave.integrations.claude_agent_sdk.usage import total_input_tokens from weave.integrations.integration_metadata import library_integration from weave.integrations.patcher import MultiPatcher, NoOpPatcher, SymbolPatcher @@ -55,7 +49,6 @@ logger = logging.getLogger(__name__) -_TRACER_NAME = "weave.claude_agent_sdk" _DEFAULT_AGENT_NAME = "claude_agent_sdk" _PROVIDER_NAME = "anthropic" @@ -78,34 +71,6 @@ class _AssistantOutput: reasoning: Reasoning -@dataclass(frozen=True, slots=True) -class _PendingChat: - """A chat span whose end is deferred until the aggregate usage is known. - - ``attrs`` is mutated in place (usage keys added) before the span is ended. - """ - - span: Any - attrs: dict[str, Any] - - -@dataclass(frozen=True, slots=True) -class _OpenTool: - """An in-flight execute_tool span awaiting its tool_result. - - ``arguments`` is the JSON-encoded tool input, captured when the tool_use - block is seen so it can be attached to the span at its tool_result. - """ - - span: Any - name: str - arguments: str - - -def _tracer() -> Any: - return otel_trace.get_tracer(_TRACER_NAME) - - def _usage_from_result(usage: dict[str, Any] | None) -> Usage: """Build a Usage from a ResultMessage's aggregate usage dict.""" raw = usage or {} @@ -154,62 +119,45 @@ class _TurnState: turn, owned by ``_trace_turn``. """ - user_prompt: str | None - # Resolved once at turn start so the span name and gen_ai.agent.name stay - # consistent even if finalization runs outside the context that set it. - agent_name: str = _DEFAULT_AGENT_NAME - conversation_id: str = "" + conversation: Conversation + turn: Turn model: str = "" final_text: str = "" is_error: bool = False accumulated: list[Message] = field(default_factory=list) pending_thinking: list[str] = field(default_factory=list) - # The most recent chat span, whose end is deferred so the aggregate usage - # from ResultMessage can be attached before it closes. - pending_chat: _PendingChat | None = None - # tool_use_id -> in-flight execute_tool span. - open_tool_spans: dict[str, _OpenTool] = field(default_factory=dict) + pending_chat: LLM | None = None + open_tools: dict[str, Tool] = field(default_factory=dict) def _flush_pending_chat(state: _TurnState, *, usage: Usage | None = None) -> None: - """Set attrs (optionally usage) on the deferred chat span and end it.""" + """Record optional aggregate usage on the deferred chat and end it.""" pending = state.pending_chat if pending is None: return - attrs = pending.attrs - if usage is not None: - if usage.input_tokens: - attrs["gen_ai.usage.input_tokens"] = usage.input_tokens - if usage.output_tokens: - attrs["gen_ai.usage.output_tokens"] = usage.output_tokens - if usage.cache_creation_input_tokens: - attrs["gen_ai.usage.cache_creation.input_tokens"] = ( - usage.cache_creation_input_tokens - ) - if usage.cache_read_input_tokens: - attrs["gen_ai.usage.cache_read.input_tokens"] = ( - usage.cache_read_input_tokens - ) - for key, value in attrs.items(): - pending.span.set_attribute(key, value) - pending.span.end() - state.pending_chat = None + try: + with pending: + if usage is not None: + pending.record(usage=usage) + finally: + state.pending_chat = None -def _process_message(msg: Any, tracer: Any, state: _TurnState) -> None: - """Handle one streamed message, creating/closing child spans as needed.""" +def _process_message(msg: Any, state: _TurnState) -> str | None: + """Handle one streamed message and return a newly observed SDK session ID.""" if isinstance(msg, SystemMessage): session_id = (msg.data or {}).get("session_id") - if session_id: - state.conversation_id = session_id - return + if isinstance(session_id, str) and session_id: + state.conversation.conversation_id = session_id + return session_id + return None if isinstance(msg, AssistantMessage): # Buffer thinking-only messages so extended-thinking deltas fold into # the next response's chat span rather than spawning an empty one. if all(isinstance(b, ThinkingBlock) for b in msg.content) and msg.content: state.pending_thinking.extend(b.thinking for b in msg.content) - return + return None # A new response means the previous chat span is done (no usage — only # the final one carries the aggregate usage). @@ -222,28 +170,28 @@ def _process_message(msg: Any, tracer: Any, state: _TurnState) -> None: if output.text: state.final_text = output.text - chat = tracer.start_span(f"chat {msg.model or ''}".rstrip()) - chat_attrs = llm_attributes( + chat = state.turn.start_llm( model=msg.model or "", provider_name=_PROVIDER_NAME, - conversation_id=state.conversation_id, + ) + state.pending_chat = chat + chat.record( input_messages=list(state.accumulated), output_messages=[output.message], reasoning=output.reasoning if output.reasoning.content else None, ) - chat_attrs.update(_INTEGRATION_OTEL_ATTRS) - state.pending_chat = _PendingChat(span=chat, attrs=chat_attrs) state.accumulated.append(output.message) for block in msg.content: if isinstance(block, ToolUseBlock): - tool_span = tracer.start_span(f"execute_tool {block.name}") - state.open_tool_spans[block.id] = _OpenTool( - span=tool_span, + tool = state.turn.start_tool( name=block.name, arguments=json.dumps(block.input, default=str), + tool_call_id=block.id, ) - return + tool.started_at = datetime.now(timezone.utc) + state.open_tools[block.id] = tool + return None if isinstance(msg, UserMessage): content = msg.content if isinstance(msg.content, list) else [] @@ -253,23 +201,17 @@ def _process_message(msg: Any, tracer: Any, state: _TurnState) -> None: state.accumulated.append( Message.tool_result(block.tool_use_id, block.content) ) - open_tool = state.open_tool_spans.pop(block.tool_use_id, None) + open_tool = state.open_tools.get(block.tool_use_id) if open_tool is None: continue - attrs = execute_tool_attributes( - tool_name=open_tool.name, - conversation_id=state.conversation_id, - tool_call_arguments=open_tool.arguments, - tool_call_result=str(block.content), - tool_call_id=block.tool_use_id, - ) - attrs.update(_INTEGRATION_OTEL_ATTRS) - for key, value in attrs.items(): - open_tool.span.set_attribute(key, value) - if block.is_error: - open_tool.span.set_status(StatusCode.ERROR, "tool reported an error") - open_tool.span.end() - return + del state.open_tools[block.tool_use_id] + with open_tool: + open_tool.result = str(block.content) + if block.is_error: + open_tool._record_otel_error( # pyright: ignore[reportPrivateUsage] + RuntimeError("tool reported an error") + ) + return None if isinstance(msg, ResultMessage): _flush_pending_chat(state, usage=_usage_from_result(msg.usage)) @@ -277,34 +219,29 @@ def _process_message(msg: Any, tracer: Any, state: _TurnState) -> None: state.final_text = msg.result if msg.is_error: state.is_error = True - return + return None + + return None -def _finalize_turn(root: Any, state: _TurnState) -> None: - """Close any open child spans and finish the root invoke_agent span.""" +def _finalize_turn(state: _TurnState) -> None: + """Close open children and record terminal fields on the turn.""" _flush_pending_chat(state) - for open_tool in state.open_tool_spans.values(): - open_tool.span.end() - state.open_tool_spans.clear() - - attrs = invoke_agent_attributes( - agent_name=state.agent_name, - conversation_id=state.conversation_id, - provider_name=_PROVIDER_NAME, + for tool in state.open_tools.values(): + with tool: + pass + state.open_tools.clear() + + state.turn.record( model=state.model, - input_messages=[Message(role="user", content=state.user_prompt)] - if state.user_prompt - else None, output_messages=[Message(role="assistant", content=state.final_text)] if state.final_text else None, ) - attrs.update(_INTEGRATION_OTEL_ATTRS) - for key, value in attrs.items(): - root.set_attribute(key, value) if state.is_error: - root.set_status(StatusCode.ERROR, state.final_text or "agent run failed") - root.end() + state.turn._record_otel_error( # pyright: ignore[reportPrivateUsage] + RuntimeError(state.final_text or "agent run failed") + ) async def _trace_turn( @@ -319,32 +256,35 @@ async def _trace_turn( single ``ClaudeSDKClient``: the ``system/init`` message (which holds the session_id) is only sent on the first turn, so later turns must inherit it. """ - tracer = _tracer() - # Resolve once here: this runs on the first __anext__, inside the user's - # ``agent_name_override(...)`` block, so the override contextvar is visible. agent_name = resolve_agent_name(_DEFAULT_AGENT_NAME) - root = tracer.start_span(f"invoke_agent {agent_name}") - token = otel_context.attach(otel_trace.set_span_in_context(root)) - state = _TurnState(user_prompt=user_prompt, agent_name=agent_name) - if conversation_id_holder is not None and conversation_id_holder[0]: - state.conversation_id = conversation_id_holder[0] - try: - async for msg in messages: + conversation_id = ( + conversation_id_holder[0] if conversation_id_holder is not None else "" + ) + with Conversation( + conversation_id=conversation_id, + agent_name=agent_name, + continue_parent_trace=True, + attributes=_INTEGRATION_OTEL_ATTRS, + ) as conversation: + with conversation.start_turn(user_message=user_prompt or "") as turn: + state = _TurnState( + conversation=conversation, + turn=turn, + ) try: - _process_message(msg, tracer, state) - if conversation_id_holder is not None and state.conversation_id: - conversation_id_holder[0] = state.conversation_id - except Exception: - # Never let span bookkeeping break the user's stream. - logger.exception("claude_agent_sdk OTel span processing failed") - yield msg - except Exception as exc: - root.set_status(StatusCode.ERROR, str(exc)) - root.record_exception(exc) - raise - finally: - _finalize_turn(root, state) - otel_context.detach(token) + async for msg in messages: + try: + session_id = _process_message(msg, state) + if conversation_id_holder is not None and session_id: + conversation_id_holder[0] = session_id + except Exception: + # Never let span bookkeeping break the user's stream. + logger.exception( + "claude_agent_sdk GenAI span processing failed" + ) + yield msg + finally: + _finalize_turn(state) def _patched_process_query_wrapper(settings: IntegrationSettings) -> Any: