diff --git a/AGENTS.md b/AGENTS.md index 4b3c547b508e..604b6781a723 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -452,6 +452,10 @@ 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. +- `SubAgent` owns its input/output messages and originating tool-call + arguments/result. Integrations populate them through `SubAgent.record()` so + `_build_attrs()` applies `include_content` and PII redaction; do not write + those content attributes directly with `set_attributes()`. - Mypy requires explicit `return None` paths in functions annotated with `T | None`; bare `return` and implicit fallthrough trigger return-value errors. @@ -478,13 +482,26 @@ pnpm exec tsx examples/claudeAgents.ts - 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. + `Conversation`, `Turn`, `LLM`, `Tool`, and `SubAgent` 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. - Tap Python `AsyncIterable[dict]` prompts without consuming or cloning them. Map Claude text and base64/URL image blocks to GenAI text, blob, and URI parts; media payloads remain data for `content_refs`, not chat prose. +- Treat `Agent` and legacy `Task` tool calls as subagents keyed by tool-use ID. + Route nested assistant messages and tools through `parent_tool_use_id`. + Synchronous calls close on their matching tool result; background launch + acknowledgements stay open until `task_notification`. Dispatch task events + through `SystemMessage.subtype` and `data`, map `task_id` to tool-use ID from + `task_started` / `task_progress`, and do not import typed task messages that + are absent from older supported Claude Agent SDK versions. +- Start those subagents with `SubAgent.start(set_current=False)`, never + `__enter__`. Parallel delegations close in completion order, and `end()` + detaches through `ContextVar.reset`, so an out-of-LIFO close leaves the + ambient OTel context pointing at an ended span — user code running between + streamed messages would nest under it. Children nest either way, because + `start_llm` / `start_tool` / `start_subagent` thread an explicit parent. - 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 diff --git a/tests/conversation/test_conversation_otel.py b/tests/conversation/test_conversation_otel.py index 7428eb8f8f93..e190d1090707 100644 --- a/tests/conversation/test_conversation_otel.py +++ b/tests/conversation/test_conversation_otel.py @@ -2318,6 +2318,12 @@ def test_sets_all_fields(self) -> None: name="research-bot", model="gpt-4o-mini", system_instructions=["sys"], + input_messages=[Message.user("question")], + output_messages=[Message.assistant("answer")], + tool_name="Agent", + tool_call_id="toolu_1", + tool_call_arguments='{"prompt": "question"}', + tool_call_result="answer", agent_id="id-1", agent_description="desc", agent_version="v1", @@ -2325,6 +2331,12 @@ def test_sets_all_fields(self) -> None: assert sa.name == "research-bot" assert sa.model == "gpt-4o-mini" assert sa.system_instructions == ["sys"] + assert sa.input_messages == [Message.user("question")] + assert sa.output_messages == [Message.assistant("answer")] + assert sa.tool_name == "Agent" + assert sa.tool_call_id == "toolu_1" + assert sa.tool_call_arguments == '{"prompt": "question"}' + assert sa.tool_call_result == "answer" assert sa.agent_id == "id-1" assert sa.agent_description == "desc" assert sa.agent_version == "v1" diff --git a/tests/conversation/test_conversation_settings.py b/tests/conversation/test_conversation_settings.py index fab2e5696eee..608afec5bd71 100644 --- a/tests/conversation/test_conversation_settings.py +++ b/tests/conversation/test_conversation_settings.py @@ -19,6 +19,7 @@ from __future__ import annotations +import json import platform import sys from collections.abc import Callable @@ -34,6 +35,7 @@ LLM, Message, Reasoning, + SubAgent, Tool, Turn, log_conversation, @@ -289,6 +291,92 @@ def test_redact_pii_applied( assert _REDACTED_EMAIL in attrs[key], key +def _streaming_subagent_with_content(include_content: bool) -> None: + with start_conversation( + conversation_id="convo-subagent-content", + include_content=include_content, + ) as session: + with session.start_turn() as turn: + with turn.start_subagent(name="researcher") as subagent: + subagent.record( + input_messages=[Message.user(f"Research {_PII_EMAIL}")], + output_messages=[Message.assistant(f"Found {_PII_EMAIL}")], + tool_name="Agent", + tool_call_id="toolu_01", + tool_call_arguments=json.dumps({"email": _PII_EMAIL}), + tool_call_result=f"Found {_PII_EMAIL}", + ) + + +def _batch_subagent_with_content(include_content: bool) -> None: + log_turn( + conversation_id="convo-subagent-content", + agent_name="orchestrator", + include_content=include_content, + spans=[ + SubAgent( + name="researcher", + input_messages=[Message.user(f"Research {_PII_EMAIL}")], + output_messages=[Message.assistant(f"Found {_PII_EMAIL}")], + tool_name="Agent", + tool_call_id="toolu_01", + tool_call_arguments=json.dumps({"email": _PII_EMAIL}), + tool_call_result=f"Found {_PII_EMAIL}", + ) + ], + ) + + +@pytest.mark.parametrize( + "emit_subagent", + [ + pytest.param(_streaming_subagent_with_content, id="streaming"), + pytest.param(_batch_subagent_with_content, id="batch"), + ], +) +def test_subagent_content_is_pii_redacted( + emit_subagent: Callable[[bool], None], + otel_spans: InMemorySpanExporter, + fake_presidio: None, +) -> None: + with override_settings( + redact_pii=True, + capture_client_info=False, + capture_system_info=False, + ): + emit_subagent(True) + + spans = _spans_with_prefix(otel_spans, "invoke_agent researcher") + assert len(spans) == 1 + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.name": "researcher", + "gen_ai.conversation.id": "convo-subagent-content", + "gen_ai.input.messages": json.dumps( + [ + { + "role": "user", + "parts": [ + {"type": "text", "content": f"Research {_REDACTED_EMAIL}"} + ], + } + ] + ), + "gen_ai.output.messages": json.dumps( + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": f"Found {_REDACTED_EMAIL}"}], + } + ] + ), + "gen_ai.tool.name": "Agent", + "gen_ai.tool.call.id": "toolu_01", + "gen_ai.tool.call.arguments": json.dumps({"email": _REDACTED_EMAIL}), + "gen_ai.tool.call.result": f"Found {_REDACTED_EMAIL}", + } + + # --------------------------------------------------------------------------- # include_content=False — content dropped at source, Presidio never loaded # --------------------------------------------------------------------------- @@ -358,6 +446,37 @@ def test_include_content_false_skips_presidio( mock_get_engines.assert_not_called() +@pytest.mark.parametrize( + "emit_subagent", + [ + pytest.param(_streaming_subagent_with_content, id="streaming"), + pytest.param(_batch_subagent_with_content, id="batch"), + ], +) +def test_subagent_include_content_false_drops_content_before_redaction( + emit_subagent: Callable[[bool], None], + otel_spans: InMemorySpanExporter, +) -> None: + with override_settings( + redact_pii=True, + capture_client_info=False, + capture_system_info=False, + ): + with patch("weave.utils.pii_redaction._get_engines") as mock_get_engines: + emit_subagent(False) + + mock_get_engines.assert_not_called() + spans = _spans_with_prefix(otel_spans, "invoke_agent researcher") + assert len(spans) == 1 + assert dict(spans[0].attributes) == { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.name": "researcher", + "gen_ai.conversation.id": "convo-subagent-content", + "gen_ai.tool.name": "Agent", + "gen_ai.tool.call.id": "toolu_01", + } + + def _reasoning_content_off_streaming() -> None: with start_conversation( conversation_id="convo-reasoning-off-streaming", include_content=False diff --git a/tests/conversation/test_subagent_nesting.py b/tests/conversation/test_subagent_nesting.py index 9e2e2248ebd5..571f1819a0d8 100644 --- a/tests/conversation/test_subagent_nesting.py +++ b/tests/conversation/test_subagent_nesting.py @@ -18,6 +18,7 @@ import threading import pytest +from opentelemetry import trace as otel_trace from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from weave.conversation.conversation import ( @@ -333,3 +334,59 @@ def worker() -> None: llm_span = _by_prefix(spans, "chat") assert llm_span.parent is not None assert llm_span.parent.span_id == sa_span.context.span_id + + +# --------------------------------------------------------------------------- +# Concurrent sub-agents: SubAgent.start(set_current=False) +# --------------------------------------------------------------------------- + + +class TestConcurrentSubagents: + """``end()`` detaches via ``ContextVar.reset``, so overlapping sub-agents + that finish out of LIFO order corrupt the ambient OTel context. Adapters + that run sub-agents concurrently start them with ``set_current=False``. + """ + + @pytest.mark.parametrize("close_order", [("a", "b"), ("b", "a")]) + def test_ambient_context_survives_any_close_order( + self, otel_spans: InMemorySpanExporter, close_order: tuple[str, str] + ) -> None: + with Conversation(conversation_id="s"), Turn(agent_name="bot") as turn: + subagents = { + "a": turn.start_subagent(name="alpha").start(set_current=False), + "b": turn.start_subagent(name="beta").start(set_current=False), + } + for key in close_order: + subagents[key].end() + assert ( + otel_trace.get_current_span().get_span_context().span_id + == turn._otel_span.get_span_context().span_id + ) + + spans = otel_spans.get_finished_spans() + turn_span = _by_agent_name(spans, "bot") + for name in ("alpha", "beta"): + assert _by_agent_name(spans, name).parent.span_id == ( + turn_span.context.span_id + ) + + def test_children_still_nest_under_a_non_current_subagent( + self, otel_spans: InMemorySpanExporter + ) -> None: + with Conversation(conversation_id="s"), Turn(agent_name="bot") as turn: + first = turn.start_subagent(name="alpha").start(set_current=False) + second = turn.start_subagent(name="beta").start(set_current=False) + with first.start_llm(model="gpt-4o"): + pass + with second.start_tool(name="grep"): + pass + first.end() + second.end() + + spans = otel_spans.get_finished_spans() + assert _by_prefix(spans, "chat").parent.span_id == ( + _by_agent_name(spans, "alpha").context.span_id + ) + assert _by_prefix(spans, "execute_tool").parent.span_id == ( + _by_agent_name(spans, "beta").context.span_id + ) diff --git a/tests/integrations/claude_agent_sdk/cassettes/background_subagent_response.json b/tests/integrations/claude_agent_sdk/cassettes/background_subagent_response.json new file mode 100644 index 000000000000..ea84dfb2c7b4 --- /dev/null +++ b/tests/integrations/claude_agent_sdk/cassettes/background_subagent_response.json @@ -0,0 +1,163 @@ +[ + { + "type": "system", + "subtype": "init", + "session_id": "s-background-subagent001" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "I will launch a background researcher." + }, + { + "type": "tool_use", + "id": "toolu_background_agent_01", + "name": "Agent", + "input": { + "subagent_type": "researcher", + "description": "Research while the root agent continues", + "prompt": "Find the capital of France.", + "model": "claude-haiku-4-5", + "run_in_background": true + } + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_background_agent_01", + "content": [ + { + "type": "text", + "text": "{\"isAsync\": true, \"status\": \"async_launched\", \"agentId\": \"task-background-01\", \"description\": \"Research while the root agent continues\", \"outputFile\": \"/tmp/task-background-01.output\"}" + } + ] + } + ] + }, + "parent_tool_use_id": null + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "While that runs, I will prepare the response." + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "system", + "subtype": "task_started", + "task_id": "task-background-01", + "description": "Research while the root agent continues", + "uuid": "uuid-task-started-01", + "session_id": "s-background-subagent001", + "tool_use_id": "toolu_background_agent_01", + "task_type": "agent" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "I will verify the capital." + }, + { + "type": "tool_use", + "id": "toolu_background_bash_01", + "name": "Bash", + "input": { + "command": "printf Paris" + } + } + ], + "model": "claude-haiku-4-5" + }, + "parent_tool_use_id": "toolu_background_agent_01" + }, + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_background_bash_01", + "content": "Paris" + } + ] + }, + "parent_tool_use_id": "toolu_background_agent_01" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "The delegated research is complete." + } + ], + "model": "claude-haiku-4-5" + }, + "parent_tool_use_id": "toolu_background_agent_01" + }, + { + "type": "system", + "subtype": "task_notification", + "task_id": "task-background-01", + "status": "completed", + "output_file": "/tmp/task-background-01.output", + "summary": "The capital of France is Paris.", + "uuid": "uuid-task-notification-01", + "session_id": "s-background-subagent001", + "usage": { + "total_tokens": 42, + "tool_uses": 1, + "duration_ms": 800 + } + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "The background researcher confirmed Paris." + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "result", + "subtype": "result", + "duration_ms": 2200, + "duration_api_ms": 1700, + "is_error": false, + "num_turns": 1, + "session_id": "s-background-subagent001", + "total_cost_usd": 0.007, + "usage": { + "input_tokens": 140, + "output_tokens": 55 + }, + "result": "The background researcher confirmed Paris." + } +] diff --git a/tests/integrations/claude_agent_sdk/cassettes/parallel_subagent_response.json b/tests/integrations/claude_agent_sdk/cassettes/parallel_subagent_response.json new file mode 100644 index 000000000000..918accd89931 --- /dev/null +++ b/tests/integrations/claude_agent_sdk/cassettes/parallel_subagent_response.json @@ -0,0 +1,122 @@ +[ + { + "type": "system", + "subtype": "init", + "session_id": "s-parallel001" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "I will delegate both questions." + }, + { + "type": "tool_use", + "id": "toolu_agent_a", + "name": "Agent", + "input": { + "subagent_type": "geographer", + "description": "Answer geography questions", + "prompt": "Find the capital of France." + } + }, + { + "type": "tool_use", + "id": "toolu_agent_b", + "name": "Agent", + "input": { + "subagent_type": "astronomer", + "description": "Answer astronomy questions", + "prompt": "Find the largest planet." + } + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ], + "model": "claude-haiku-4-5" + }, + "parent_tool_use_id": "toolu_agent_a" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "The largest planet is Jupiter." + } + ], + "model": "claude-haiku-4-5" + }, + "parent_tool_use_id": "toolu_agent_b" + }, + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_agent_a", + "content": "The capital of France is Paris.", + "is_error": false + } + ] + }, + "parent_tool_use_id": null + }, + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_agent_b", + "content": "The largest planet is Jupiter.", + "is_error": false + } + ] + }, + "parent_tool_use_id": null + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "Paris and Jupiter." + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "result", + "subtype": "success", + "duration_ms": 12, + "duration_api_ms": 10, + "is_error": false, + "num_turns": 1, + "session_id": "s-parallel001", + "total_cost_usd": 0.001, + "usage": { + "input_tokens": 20, + "output_tokens": 10 + }, + "result": "Paris and Jupiter." + } +] diff --git a/tests/integrations/claude_agent_sdk/cassettes/subagent_response.json b/tests/integrations/claude_agent_sdk/cassettes/subagent_response.json new file mode 100644 index 000000000000..2fcacf4ae041 --- /dev/null +++ b/tests/integrations/claude_agent_sdk/cassettes/subagent_response.json @@ -0,0 +1,123 @@ +[ + { + "type": "system", + "subtype": "init", + "session_id": "s-subagent001" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "I will delegate this research." + }, + { + "type": "tool_use", + "id": "toolu_agent_01", + "name": "Agent", + "input": { + "subagent_type": "researcher", + "description": "Research factual questions", + "prompt": "Find the capital of France." + } + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "I will verify the answer." + }, + { + "type": "tool_use", + "id": "toolu_bash_01", + "name": "Bash", + "input": { + "command": "printf Paris" + } + } + ], + "model": "claude-haiku-4-5" + }, + "parent_tool_use_id": "toolu_agent_01" + }, + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_bash_01", + "content": "Paris" + } + ] + }, + "parent_tool_use_id": "toolu_agent_01" + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ], + "model": "claude-haiku-4-5" + }, + "parent_tool_use_id": "toolu_agent_01" + }, + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_agent_01", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ] + } + ] + }, + "parent_tool_use_id": null + }, + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "The researcher confirmed that Paris is the capital of France." + } + ], + "model": "claude-sonnet-4-6" + }, + "parent_tool_use_id": null + }, + { + "type": "result", + "subtype": "result", + "duration_ms": 1800, + "duration_api_ms": 1400, + "is_error": false, + "num_turns": 1, + "session_id": "s-subagent001", + "total_cost_usd": 0.006, + "usage": { + "input_tokens": 120, + "output_tokens": 45 + }, + "result": "The researcher confirmed that Paris is the capital of France." + } +] 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 723cd3e2f43a..b9a00aa983ad 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 @@ -27,6 +27,10 @@ get_claude_agent_sdk_otel_patcher, ) from weave.trace.settings import override_settings +from weave.utils import pii_redaction + +_PII_EMAIL = "alice@example.com" +_REDACTED_EMAIL = "" @pytest.fixture @@ -213,6 +217,507 @@ async def test_tool_use_query_otel(otel_spans: InMemorySpanExporter) -> None: assert len(tool_call_chats) == 1 +# --- query(): subagent delegation ------------------------------------------- + + +@pytest.mark.asyncio +async def test_subagent_query_otel(otel_spans: InMemorySpanExporter) -> None: + await run_query("subagent_response", "Find the capital of France") + spans = otel_spans.get_finished_spans() + + agent_spans = get_spans_by_op(spans, "invoke_agent") + assert len(agent_spans) == 2 + root_span = next( + span + for span in agent_spans + if get_attrs(span)["gen_ai.agent.name"] == "claude_agent_sdk" + ) + subagent_span = next( + span + for span in agent_spans + if get_attrs(span)["gen_ai.agent.name"] == "researcher" + ) + assert subagent_span.name == "invoke_agent researcher" + assert subagent_span.parent.span_id == root_span.context.span_id + assert check_integration_and_strip(get_attrs(subagent_span)) == { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.name": "researcher", + "gen_ai.conversation.id": "s-subagent001", + "gen_ai.request.model": "claude-haiku-4-5", + "gen_ai.agent.description": "Research factual questions", + "gen_ai.tool.call.id": "toolu_agent_01", + "gen_ai.tool.name": "Agent", + "gen_ai.tool.call.arguments": json.dumps( + { + "subagent_type": "researcher", + "description": "Research factual questions", + "prompt": "Find the capital of France.", + } + ), + "gen_ai.input.messages": json.dumps( + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "Find the capital of France.", + } + ], + } + ] + ), + "gen_ai.tool.call.result": "The capital of France is Paris.", + "gen_ai.output.messages": json.dumps( + [ + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "The capital of France is Paris.", + } + ], + } + ] + ), + } + + chat_spans = get_spans_by_op(spans, "chat") + assert len(chat_spans) == 4 + assert [ + span.parent.span_id == subagent_span.context.span_id for span in chat_spans + ].count(True) == 2 + assert [ + span.parent.span_id == root_span.context.span_id for span in chat_spans + ].count(True) == 2 + + tool_spans = get_spans_by_op(spans, "execute_tool") + assert len(tool_spans) == 1 + tool_span = tool_spans[0] + assert tool_span.name == "execute_tool Bash" + assert tool_span.parent.span_id == subagent_span.context.span_id + assert check_integration_and_strip(get_attrs(tool_span)) == { + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": "Bash", + "gen_ai.tool.call.id": "toolu_bash_01", + "gen_ai.tool.call.arguments": '{"command": "printf Paris"}', + "gen_ai.tool.call.result": "Paris", + "gen_ai.conversation.id": "s-subagent001", + } + + +@pytest.mark.asyncio +async def test_subagent_query_content_is_pii_redacted_otel( + otel_spans: InMemorySpanExporter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages = load_cassette("subagent_response") + messages[1]["message"]["content"][1]["input"]["prompt"] = f"Research {_PII_EMAIL}" + messages[5]["message"]["content"][0]["content"][0]["text"] = f"Found {_PII_EMAIL}" + + def redact_messages(messages: list[Any]) -> list[Any]: + return [ + message.model_copy( + update={"content": message.content.replace(_PII_EMAIL, _REDACTED_EMAIL)} + ) + for message in messages + ] + + monkeypatch.setattr(pii_redaction, "redact_messages", redact_messages) + monkeypatch.setattr( + pii_redaction, + "redact_pii_string", + lambda value: value.replace(_PII_EMAIL, _REDACTED_EMAIL), + ) + + with override_settings(redact_pii=True): + async for _ in query( + prompt="Delegate sensitive research", + options=ClaudeAgentOptions(), + transport=ReplayTransport(messages), + ): + pass + + subagent_span = next( + span + for span in get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent") + if get_attrs(span)["gen_ai.agent.name"] == "researcher" + ) + assert check_integration_and_strip(get_attrs(subagent_span)) == { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.name": "researcher", + "gen_ai.conversation.id": "s-subagent001", + "gen_ai.request.model": "claude-haiku-4-5", + "gen_ai.agent.description": "Research factual questions", + "gen_ai.input.messages": json.dumps( + [ + { + "role": "user", + "parts": [ + {"type": "text", "content": f"Research {_REDACTED_EMAIL}"} + ], + } + ] + ), + "gen_ai.output.messages": json.dumps( + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": f"Found {_REDACTED_EMAIL}"}], + } + ] + ), + "gen_ai.tool.name": "Agent", + "gen_ai.tool.call.id": "toolu_agent_01", + "gen_ai.tool.call.arguments": json.dumps( + { + "subagent_type": "researcher", + "description": "Research factual questions", + "prompt": f"Research {_REDACTED_EMAIL}", + } + ), + "gen_ai.tool.call.result": f"Found {_REDACTED_EMAIL}", + } + + +@pytest.mark.parametrize( + ( + "status", + "summary", + "tool_name", + "notification_has_tool_use_id", + "launch_has_task_id", + "expected_status_code", + "expected_status_description", + "expected_error_type", + ), + [ + pytest.param( + "completed", + "The capital of France is Paris.", + "Agent", + False, + True, + StatusCode.UNSET, + None, + None, + id="completed-agent-task-id-fallback", + ), + pytest.param( + "failed", + "The background researcher failed.", + "Agent", + True, + True, + StatusCode.ERROR, + "background subagent failed", + "claude_agent_sdk.task_failed", + id="failed-agent-direct-tool-use-id", + ), + pytest.param( + "stopped", + "The background researcher was stopped.", + "Task", + False, + False, + StatusCode.ERROR, + "background subagent stopped", + "claude_agent_sdk.task_stopped", + id="stopped-legacy-task-task-id-fallback", + ), + ], +) +@pytest.mark.asyncio +async def test_background_subagent_closes_on_task_notification_otel( + status: str, + summary: str, + tool_name: str, + notification_has_tool_use_id: bool, + launch_has_task_id: bool, + expected_status_code: StatusCode, + expected_status_description: str | None, + expected_error_type: str | None, + otel_spans: InMemorySpanExporter, +) -> None: + messages = load_cassette("background_subagent_response") + messages[1]["message"]["content"][1]["name"] = tool_name + + launch_text = messages[2]["message"]["content"][0]["content"][0]["text"] + launch_payload = json.loads(launch_text) + if not launch_has_task_id: + del launch_payload["agentId"] + messages[2]["message"]["content"][0]["content"][0]["text"] = json.dumps( + launch_payload + ) + + notification = messages[8] + notification["status"] = status + notification["summary"] = summary + if notification_has_tool_use_id: + notification["tool_use_id"] = "toolu_background_agent_01" + + async for _ in query( + prompt="Research in the background", + options=ClaudeAgentOptions(), + transport=ReplayTransport(messages), + ): + pass + + spans = otel_spans.get_finished_spans() + agent_spans = get_spans_by_op(spans, "invoke_agent") + assert len(agent_spans) == 2 + root_span = next( + span + for span in agent_spans + if get_attrs(span)["gen_ai.agent.name"] == "claude_agent_sdk" + ) + subagent_span = next( + span + for span in agent_spans + if get_attrs(span)["gen_ai.agent.name"] == "researcher" + ) + + expected_subagent_attrs = { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.name": "researcher", + "gen_ai.conversation.id": "s-background-subagent001", + "gen_ai.request.model": "claude-haiku-4-5", + "gen_ai.agent.id": "task-background-01", + "gen_ai.agent.description": "Research while the root agent continues", + "gen_ai.input.messages": json.dumps( + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "Find the capital of France.", + } + ], + } + ] + ), + "gen_ai.output.messages": json.dumps( + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": summary}], + } + ] + ), + "gen_ai.tool.name": tool_name, + "gen_ai.tool.call.id": "toolu_background_agent_01", + "gen_ai.tool.call.arguments": json.dumps( + { + "subagent_type": "researcher", + "description": "Research while the root agent continues", + "prompt": "Find the capital of France.", + "model": "claude-haiku-4-5", + "run_in_background": True, + } + ), + "gen_ai.tool.call.result": summary, + "claude_agent_sdk.task.status": status, + } + if expected_error_type is not None: + expected_subagent_attrs["error.type"] = expected_error_type + assert check_integration_and_strip(get_attrs(subagent_span)) == ( + expected_subagent_attrs + ) + assert subagent_span.name == "invoke_agent researcher" + assert subagent_span.parent.span_id == root_span.context.span_id + assert subagent_span.status.status_code == expected_status_code + assert subagent_span.status.description == expected_status_description + + chat_spans = get_spans_by_op(spans, "chat") + assert [ + ( + span.parent.span_id, + get_all_text(get_messages(span, "gen_ai.output.messages")), + ) + for span in chat_spans + ] == [ + (root_span.context.span_id, "I will launch a background researcher."), + (root_span.context.span_id, "While that runs, I will prepare the response."), + (subagent_span.context.span_id, "I will verify the capital."), + (subagent_span.context.span_id, "The delegated research is complete."), + (root_span.context.span_id, "The background researcher confirmed Paris."), + ] + child_chat_spans = [ + span + for span in chat_spans + if span.parent.span_id == subagent_span.context.span_id + ] + assert [span.end_time <= subagent_span.end_time for span in child_chat_spans] == [ + True, + True, + ] + + tool_spans = get_spans_by_op(spans, "execute_tool") + assert len(tool_spans) == 1 + assert tool_spans[0].parent.span_id == subagent_span.context.span_id + assert check_integration_and_strip(get_attrs(tool_spans[0])) == { + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": "Bash", + "gen_ai.tool.call.id": "toolu_background_bash_01", + "gen_ai.tool.call.arguments": '{"command": "printf Paris"}', + "gen_ai.tool.call.result": "Paris", + "gen_ai.conversation.id": "s-background-subagent001", + } + + +@pytest.mark.asyncio +async def test_background_subagent_summary_is_pii_redacted_otel( + otel_spans: InMemorySpanExporter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages = load_cassette("background_subagent_response") + messages[1]["message"]["content"][1]["input"]["prompt"] = f"Research {_PII_EMAIL}" + messages[8]["summary"] = f"Found {_PII_EMAIL}" + + def redact_messages(messages: list[Any]) -> list[Any]: + return [ + message.model_copy( + update={"content": message.content.replace(_PII_EMAIL, _REDACTED_EMAIL)} + ) + for message in messages + ] + + monkeypatch.setattr(pii_redaction, "redact_messages", redact_messages) + monkeypatch.setattr( + pii_redaction, + "redact_pii_string", + lambda value: value.replace(_PII_EMAIL, _REDACTED_EMAIL), + ) + + with override_settings(redact_pii=True): + async for _ in query( + prompt="Delegate sensitive background research", + options=ClaudeAgentOptions(), + transport=ReplayTransport(messages), + ): + pass + + subagent_span = next( + span + for span in get_spans_by_op(otel_spans.get_finished_spans(), "invoke_agent") + if get_attrs(span)["gen_ai.agent.name"] == "researcher" + ) + assert check_integration_and_strip(get_attrs(subagent_span)) == { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.agent.name": "researcher", + "gen_ai.conversation.id": "s-background-subagent001", + "gen_ai.request.model": "claude-haiku-4-5", + "gen_ai.agent.id": "task-background-01", + "gen_ai.agent.description": "Research while the root agent continues", + "gen_ai.input.messages": json.dumps( + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": f"Research {_REDACTED_EMAIL}", + } + ], + } + ] + ), + "gen_ai.output.messages": json.dumps( + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": f"Found {_REDACTED_EMAIL}"}], + } + ] + ), + "gen_ai.tool.name": "Agent", + "gen_ai.tool.call.id": "toolu_background_agent_01", + "gen_ai.tool.call.arguments": json.dumps( + { + "subagent_type": "researcher", + "description": "Research while the root agent continues", + "prompt": f"Research {_REDACTED_EMAIL}", + "model": "claude-haiku-4-5", + "run_in_background": True, + } + ), + "gen_ai.tool.call.result": f"Found {_REDACTED_EMAIL}", + "claude_agent_sdk.task.status": "completed", + } + + +@pytest.mark.asyncio +async def test_parallel_subagents_close_out_of_order_otel( + otel_spans: InMemorySpanExporter, +) -> None: + """Two delegations open A, B and close A, B — completion, not LIFO, order. + + Spans the integration creates itself are pinned to an explicit parent, so + the risk is to the ambient OTel context that surrounds user code running + between streamed messages. ``user_code`` stands in for that. + """ + tracer = otel_trace.get_tracer("test.user_code") + async for _ in query( + prompt="Find the capital of France and the largest planet", + options=ClaudeAgentOptions(), + transport=ReplayTransport(load_cassette("parallel_subagent_response")), + ): + tracer.start_span("user_code").end() + + spans = otel_spans.get_finished_spans() + agent_spans = get_spans_by_op(spans, "invoke_agent") + assert len(agent_spans) == 3 + by_name = {get_attrs(span)["gen_ai.agent.name"]: span for span in agent_spans} + root = by_name["claude_agent_sdk"] + geographer = by_name["geographer"] + astronomer = by_name["astronomer"] + + assert geographer.parent.span_id == root.context.span_id + assert astronomer.parent.span_id == root.context.span_id + assert get_attrs(geographer)["gen_ai.tool.call.result"] == ( + "The capital of France is Paris." + ) + assert get_attrs(astronomer)["gen_ai.tool.call.result"] == ( + "The largest planet is Jupiter." + ) + + # Each delegation's own chat nests under it, not under its sibling. + chats_by_parent: dict[int, list[Any]] = {} + for chat in get_spans_by_op(spans, "chat"): + chats_by_parent.setdefault(chat.parent.span_id, []).append(chat) + assert ( + get_all_text( + get_messages( + chats_by_parent[geographer.context.span_id][0], "gen_ai.output.messages" + ) + ) + == "The capital of France is Paris." + ) + assert ( + get_all_text( + get_messages( + chats_by_parent[astronomer.context.span_id][0], "gen_ai.output.messages" + ) + ) + == "The largest planet is Jupiter." + ) + + # The ambient context must never hand user code an ended span as its + # parent: detaching out of LIFO order restores whichever span was current + # when the later token was attached, and that sibling has already ended. + ended_subagents = {geographer.context.span_id, astronomer.context.span_id} + orphaned = [ + span + for span in spans + if span.name == "user_code" + and span.parent is not None + and span.parent.span_id in ended_subagents + ] + assert orphaned == [] + + # --- query(): prompt caching ------------------------------------------------ diff --git a/weave/conversation/conversation.py b/weave/conversation/conversation.py index 6134c2cda8ff..c01e420a4959 100644 --- a/weave/conversation/conversation.py +++ b/weave/conversation/conversation.py @@ -166,6 +166,7 @@ def _start_otel_span( *, new_trace: bool = False, start_time_ns: int | None = None, + set_current: bool = True, ) -> None: """Create an OTel span and attach it to the current context. @@ -177,6 +178,9 @@ def _start_otel_span( When ``_parent_otel_context`` is set (a child threaded from its parent Turn/SubAgent), that context wins over ambient. ``new_trace`` is only honored when no explicit parent was provided. + + ``set_current=False`` starts the span without pushing it onto the + ambient OTel context stack — see ``SubAgent.start``. """ if not _OTEL_AVAILABLE or should_disable_weave(): return @@ -189,9 +193,10 @@ def _start_otel_span( elif new_trace: kwargs["context"] = Context() self._otel_span = tracer.start_span(name, **kwargs) - self._otel_token = otel_context.attach( - otel_trace.set_span_in_context(self._otel_span) - ) + if set_current: + self._otel_token = otel_context.attach( + otel_trace.set_span_in_context(self._otel_span) + ) # Stamp the active conversation's attributes on every span. Read from the # conversation contextvar (not OTel context) so they reach the root turn # span too, which starts in a fresh OTel Context to force a new trace. @@ -761,12 +766,20 @@ class SubAgent(_SpanBase): Maps to a nested invoke_agent OTel span in the same trace. """ + model_config = ConfigDict(validate_assignment=True) + name: str = "" model: str = "" agent_id: str = "" agent_description: str = "" agent_version: str = "" system_instructions: list[str] = Field(default_factory=list) + input_messages: list[Message] = Field(default_factory=list) + output_messages: list[Message] = Field(default_factory=list) + tool_name: str = "" + tool_call_id: str = "" + tool_call_arguments: JSONString = "" + tool_call_result: JSONString = "" started_at: datetime | None = None ended_at: datetime | None = None @@ -847,6 +860,12 @@ def record( name: str | None = None, model: str | None = None, system_instructions: list[str] | None = None, + input_messages: list[Message] | None = None, + output_messages: list[Message] | None = None, + tool_name: str | None = None, + tool_call_id: str | None = None, + tool_call_arguments: str | None = None, + tool_call_result: str | None = None, agent_id: str | None = None, agent_description: str | None = None, agent_version: str | None = None, @@ -871,6 +890,18 @@ def record( self.model = model if system_instructions is not None: self.system_instructions = system_instructions + if input_messages is not None: + self.input_messages = input_messages + if output_messages is not None: + self.output_messages = output_messages + if tool_name is not None: + self.tool_name = tool_name + if tool_call_id is not None: + self.tool_call_id = tool_call_id + if tool_call_arguments is not None: + self.tool_call_arguments = tool_call_arguments + if tool_call_result is not None: + self.tool_call_result = tool_call_result if agent_id is not None: self.agent_id = agent_id if agent_description is not None: @@ -885,29 +916,55 @@ def _build_attrs( """Build the full OTel attribute dict for this sub-agent span. Shared between streaming (``end``) and batch (``_attrs_for_span``). - ``system_instructions`` is the only content-bearing field — it is - gated by ``include_content`` and PII-redacted, mirroring ``Turn``; - the identifiers are always emitted. + Content-bearing fields are gated by ``include_content`` and + PII-redacted, mirroring ``Turn`` and ``Tool``. Identifiers are always + emitted. """ + input_messages: list[Message] | None + output_messages: list[Message] | None system_instructions: list[str] | None if include_content: + input_messages = self.input_messages + output_messages = self.output_messages system_instructions = self.system_instructions + tool_call_arguments = self.tool_call_arguments + tool_call_result = self.tool_call_result if should_redact_pii(): + input_messages = pii_redaction.redact_messages(input_messages) + output_messages = pii_redaction.redact_messages(output_messages) system_instructions = pii_redaction.redact_system_instructions( system_instructions ) + tool_call_arguments = pii_redaction.redact_pii_string( + tool_call_arguments + ) + tool_call_result = pii_redaction.redact_pii_string(tool_call_result) else: + input_messages = None + output_messages = None system_instructions = None + tool_call_arguments = "" + tool_call_result = "" attrs = invoke_agent_attributes( agent_name=self.name, model=self.model, conversation_id=conversation_id, conversation_name=conversation_name, + input_messages=input_messages, + output_messages=output_messages, system_instructions=system_instructions, agent_id=self.agent_id, agent_description=self.agent_description, agent_version=self.agent_version, ) + if self.tool_name: + attrs["gen_ai.tool.name"] = self.tool_name + if self.tool_call_id: + attrs["gen_ai.tool.call.id"] = self.tool_call_id + if tool_call_arguments: + attrs["gen_ai.tool.call.arguments"] = tool_call_arguments + if tool_call_result: + attrs["gen_ai.tool.call.result"] = tool_call_result attrs.update(_capture_info_attrs()) return attrs @@ -926,13 +983,30 @@ def end(self) -> None: ) self._end_otel_span(attrs, end_time_ns=_to_ns(self.ended_at)) - def __enter__(self) -> Self: + def start(self, *, set_current: bool = True) -> Self: + """Start this sub-agent's span. ``__enter__`` without the ``with``. + + Pass ``set_current=False`` when sub-agents can be in flight + concurrently. ``end()`` detaches via ``ContextVar.reset``, which + silently corrupts the ambient context stack when overlapping spans + end out of LIFO order — the surviving sibling stops being current, + and the last detach restores an already-ended span. Children created + through ``start_llm`` / ``start_tool`` / ``start_subagent`` nest under + this span either way, because those factories thread an explicit + parent context. + """ if self.started_at is None: self.started_at = datetime.now(timezone.utc) - start_ns = int(self.started_at.timestamp() * 1_000_000_000) - self._start_otel_span(f"invoke_agent {self.name}", start_time_ns=start_ns) + self._start_otel_span( + f"invoke_agent {self.name}", + start_time_ns=_to_ns(self.started_at), + set_current=set_current, + ) return self + def __enter__(self) -> Self: + return self.start() + def __exit__( self, exc_type: type[BaseException] | None, diff --git a/weave/integrations/claude_agent_sdk/otel_integration.py b/weave/integrations/claude_agent_sdk/otel_integration.py index a6e2b390abbd..8e53dcd0f35b 100644 --- a/weave/integrations/claude_agent_sdk/otel_integration.py +++ b/weave/integrations/claude_agent_sdk/otel_integration.py @@ -38,7 +38,16 @@ UserMessage, ) -from weave.conversation import LLM, Conversation, Message, Reasoning, Tool, Turn, Usage +from weave.conversation import ( + LLM, + Conversation, + Message, + Reasoning, + SubAgent, + Tool, + Turn, + Usage, +) from weave.conversation.agent_context import resolve_agent_name from weave.conversation.types import ( BlobPart, @@ -57,6 +66,9 @@ _DEFAULT_AGENT_NAME = "claude_agent_sdk" _PROVIDER_NAME = "anthropic" +_SUBAGENT_TOOL_NAMES = {"Agent", "Task"} +_TASK_STATUS_ATTRIBUTE = "claude_agent_sdk.task.status" +_TASK_ERROR_STATUSES = {"failed", "stopped"} _claude_agent_sdk_otel_patcher: MultiPatcher | None = None @@ -169,7 +181,15 @@ def _assistant_output_message( thinking_chunks.append(block.thinking) elif isinstance(block, ToolUseBlock): tool_calls.append( - ToolCallPart(id=block.id, name=block.name, arguments=block.input) + ToolCallPart( + id=block.id, + name=block.name, + arguments=json.dumps( + block.input, + ensure_ascii=False, + default=str, + ), + ) ) text = "\n".join(text_chunks) reasoning = Reasoning( @@ -179,6 +199,51 @@ def _assistant_output_message( return _AssistantOutput(message=message, text=text, reasoning=reasoning) +def _tool_result_text(content: Any) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + text = "\n".join( + str(block.get("text", "")) + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ) + if text: + return text + return json.dumps(content, ensure_ascii=False, default=str) + + +def _async_launch_payload(content: Any) -> dict[str, Any] | None: + candidates = content if isinstance(content, list) else [content] + for candidate in candidates: + payload = candidate + if isinstance(candidate, dict) and candidate.get("type") == "text": + payload = candidate.get("text") + if isinstance(payload, str): + try: + payload = json.loads(payload) + except (json.JSONDecodeError, TypeError): + continue + if isinstance(payload, dict) and payload.get("status") == "async_launched": + return payload + return None + + +def _nonempty_string(value: Any) -> str: + return value if isinstance(value, str) and value else "" + + +_SpanParent = Turn | SubAgent + + +@dataclass(slots=True) +class _OpenSubagent: + span: SubAgent + background: bool = False + + @dataclass(slots=True) class _TurnState: """Mutable per-turn accumulator (one invoke_agent span and its children). @@ -192,10 +257,12 @@ class _TurnState: model: str = "" final_text: str = "" is_error: bool = False - accumulated: list[Message] = field(default_factory=list) - pending_thinking: list[str] = field(default_factory=list) + accumulated: dict[str | None, list[Message]] = field(default_factory=dict) + pending_thinking: dict[str | None, list[str]] = field(default_factory=dict) pending_chat: LLM | None = None open_tools: dict[str, Tool] = field(default_factory=dict) + open_subagents: dict[str, _OpenSubagent] = field(default_factory=dict) + task_tool_use_ids: dict[str, str] = field(default_factory=dict) def _flush_pending_chat(state: _TurnState, *, usage: Usage | None = None) -> None: @@ -211,48 +278,219 @@ def _flush_pending_chat(state: _TurnState, *, usage: Usage | None = None) -> Non state.pending_chat = None +def _start_subagent( + block: ToolUseBlock, + parent: _SpanParent, + state: _TurnState, +) -> None: + raw_input = block.input if isinstance(block.input, dict) else {} + name = next( + ( + value + for key in ("subagent_type", "name") + if isinstance(value := raw_input.get(key), str) and value + ), + "subagent", + ) + model = raw_input.get("model") + description = raw_input.get("description") + prompt = raw_input.get("prompt") + subagent = parent.start_subagent( + name=name, + model=model if isinstance(model, str) else "", + ) + subagent.record( + agent_description=description if isinstance(description, str) else None, + input_messages=[Message.user(prompt)] + if isinstance(prompt, str) and prompt + else [], + tool_name=block.name, + tool_call_id=block.id, + tool_call_arguments=json.dumps( + block.input, + ensure_ascii=False, + default=str, + ), + ) + # Parallel delegations finish in completion order, not launch order. + subagent.start(set_current=False) + state.open_subagents[block.id] = _OpenSubagent( + span=subagent, + background=raw_input.get("run_in_background") is True, + ) + + +def _span_parent(msg: AssistantMessage, state: _TurnState) -> _SpanParent: + parent_tool_use_id = msg.parent_tool_use_id + if parent_tool_use_id is None: + return state.turn + open_subagent = state.open_subagents.get(parent_tool_use_id) + if open_subagent is None: + subagent = state.turn.start_subagent( + name="subagent", + model=msg.model or "", + ) + subagent.start(set_current=False) + subagent.set_attributes({"gen_ai.tool.call.id": parent_tool_use_id}) + state.open_subagents[parent_tool_use_id] = _OpenSubagent(span=subagent) + elif msg.model: + subagent = open_subagent.span + subagent.record(model=msg.model) + else: + subagent = open_subagent.span + return subagent + + +def _task_tool_use_id(data: dict[str, Any], state: _TurnState) -> str: + task_id = _nonempty_string(data.get("task_id")) + tool_use_id = _nonempty_string(data.get("tool_use_id")) + if task_id and tool_use_id: + state.task_tool_use_ids[task_id] = tool_use_id + elif task_id: + tool_use_id = state.task_tool_use_ids.get(task_id, "") + return tool_use_id + + +def _record_task_identity(data: dict[str, Any], state: _TurnState) -> None: + task_id = _nonempty_string(data.get("task_id")) + tool_use_id = _task_tool_use_id(data, state) + if not tool_use_id: + return + open_subagent = state.open_subagents.get(tool_use_id) + if open_subagent is None: + return + open_subagent.background = True + if task_id: + open_subagent.span.record(agent_id=task_id) + + +def _discard_task_mappings(tool_use_id: str, state: _TurnState) -> None: + state.task_tool_use_ids = { + task_id: mapped_tool_use_id + for task_id, mapped_tool_use_id in state.task_tool_use_ids.items() + if mapped_tool_use_id != tool_use_id + } + + +def _finish_task_notification(data: dict[str, Any], state: _TurnState) -> None: + task_id = _nonempty_string(data.get("task_id")) + tool_use_id = _task_tool_use_id(data, state) + if not tool_use_id: + return + open_subagent = state.open_subagents.get(tool_use_id) + if open_subagent is None: + return + + _flush_pending_chat(state) + summary = _nonempty_string(data.get("summary")) + status = _nonempty_string(data.get("status")) + open_subagent.span.record( + agent_id=task_id or None, + output_messages=[Message.assistant(summary)] if summary else [], + tool_call_result=summary, + ) + if status: + attributes = {_TASK_STATUS_ATTRIBUTE: status} + if status in _TASK_ERROR_STATUSES: + attributes["error.type"] = f"claude_agent_sdk.task_{status}" + open_subagent.span.set_attributes(attributes) + if status in _TASK_ERROR_STATUSES: + open_subagent.span._record_otel_error( # pyright: ignore[reportPrivateUsage] + RuntimeError(f"background subagent {status}") + ) + open_subagent.span.end() + del state.open_subagents[tool_use_id] + _discard_task_mappings(tool_use_id, state) + + +def _process_system_message(msg: SystemMessage, state: _TurnState) -> str | None: + data = msg.data if isinstance(msg.data, dict) else {} + session_id = _nonempty_string(data.get("session_id")) + if session_id: + state.conversation.conversation_id = session_id + + if msg.subtype in {"task_started", "task_progress"}: + _record_task_identity(data, state) + elif msg.subtype == "task_notification": + _finish_task_notification(data, state) + return session_id or None + + +def _record_background_launch( + tool_use_id: str, + payload: dict[str, Any] | None, + state: _TurnState, +) -> None: + open_subagent = state.open_subagents.get(tool_use_id) + if open_subagent is None: + return + open_subagent.background = True + if payload is None: + return + task_id = next( + ( + value + for key in ("agentId", "taskId", "task_id") + if (value := _nonempty_string(payload.get(key))) + ), + "", + ) + if task_id: + state.task_tool_use_ids[task_id] = tool_use_id + open_subagent.span.record(agent_id=task_id) + + 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 isinstance(session_id, str) and session_id: - state.conversation.conversation_id = session_id - return session_id - return None + return _process_system_message(msg, state) if isinstance(msg, AssistantMessage): + parent_tool_use_id = msg.parent_tool_use_id # 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) + thinking_blocks = [ + block for block in msg.content if isinstance(block, ThinkingBlock) + ] + if thinking_blocks and len(thinking_blocks) == len(msg.content): + state.pending_thinking.setdefault(parent_tool_use_id, []).extend( + block.thinking for block in thinking_blocks + ) return None # A new response means the previous chat span is done (no usage — only # the final one carries the aggregate usage). _flush_pending_chat(state) - if msg.model: + if parent_tool_use_id is None and msg.model: state.model = msg.model - output = _assistant_output_message(msg, state.pending_thinking) - state.pending_thinking.clear() - if output.text: + output = _assistant_output_message( + msg, + state.pending_thinking.pop(parent_tool_use_id, []), + ) + if parent_tool_use_id is None and output.text: state.final_text = output.text - chat = state.turn.start_llm( + parent = _span_parent(msg, state) + chat = parent.start_llm( model=msg.model or "", provider_name=_PROVIDER_NAME, ) state.pending_chat = chat + accumulated = state.accumulated.setdefault(parent_tool_use_id, []) chat.record( - input_messages=list(state.accumulated), + input_messages=list(accumulated), output_messages=[output.message], reasoning=output.reasoning if output.reasoning.content else None, ) - state.accumulated.append(output.message) + accumulated.append(output.message) for block in msg.content: if isinstance(block, ToolUseBlock): - tool = state.turn.start_tool( + if block.name in _SUBAGENT_TOOL_NAMES: + _start_subagent(block, parent, state) + continue + tool = parent.start_tool( name=block.name, arguments=json.dumps(block.input, default=str), tool_call_id=block.id, @@ -263,18 +501,44 @@ def _process_message(msg: Any, state: _TurnState) -> str | None: if isinstance(msg, UserMessage): content = msg.content if isinstance(msg.content, list) else [] + accumulated = state.accumulated.setdefault(msg.parent_tool_use_id, []) for block in content: if not isinstance(block, ToolResultBlock): continue - state.accumulated.append( - Message.tool_result(block.tool_use_id, block.content) - ) + result_text = _tool_result_text(block.content) + accumulated.append(Message.tool_result(block.tool_use_id, result_text)) + open_subagent = state.open_subagents.get(block.tool_use_id) + if open_subagent is not None: + launch_payload = _async_launch_payload(block.content) + if not block.is_error and ( + open_subagent.background or launch_payload is not None + ): + _record_background_launch( + block.tool_use_id, + launch_payload, + state, + ) + continue + open_subagent.span.record( + output_messages=[Message.assistant(result_text)] + if result_text + else [], + tool_call_result=result_text, + ) + if block.is_error: + open_subagent.span._record_otel_error( # pyright: ignore[reportPrivateUsage] + RuntimeError("subagent reported an error") + ) + open_subagent.span.end() + del state.open_subagents[block.tool_use_id] + _discard_task_mappings(block.tool_use_id, state) + continue open_tool = state.open_tools.get(block.tool_use_id) if open_tool is None: continue del state.open_tools[block.tool_use_id] with open_tool: - open_tool.result = str(block.content) + open_tool.result = result_text if block.is_error: open_tool._record_otel_error( # pyright: ignore[reportPrivateUsage] RuntimeError("tool reported an error") @@ -299,6 +563,10 @@ def _finalize_turn(state: _TurnState) -> None: with tool: pass state.open_tools.clear() + for open_subagent in state.open_subagents.values(): + open_subagent.span.end() + state.open_subagents.clear() + state.task_tool_use_ids.clear() state.turn.record( model=state.model,