feat(adk): trace workflow nodes and tool calls

py/src/braintrust/integrations/adk/__init__.py

@@ -11,6 +11,7 @@
     wrap_flow,
     wrap_mcp_tool,
     wrap_runner,
+    wrap_workflow,
 )
 
 
@@ -22,6 +23,7 @@
     "setup_braintrust",
     "wrap_agent",
     "wrap_runner",
+    "wrap_workflow",
     "wrap_flow",
     "wrap_mcp_tool",
 ]
@@ -41,7 +43,7 @@ def setup_adk(
     """
     Setup Braintrust integration with Google ADK. Will automatically patch Google ADK agents, runners, flows, and MCP tools for automatic tracing.
 
-    If you prefer manual patching take a look at `wrap_agent`, `wrap_runner`, `wrap_flow`, and `wrap_mcp_tool`.
+    If you prefer manual patching take a look at `wrap_agent`, `wrap_runner`, `wrap_workflow`, `wrap_flow`, and `wrap_mcp_tool`.
 
     Args:
         api_key (Optional[str]): Braintrust API key.

py/src/braintrust/integrations/adk/integration.py

@@ -11,6 +11,8 @@
     RunnerRunPatcher,
     ThreadBridgePatcher,
     ToolCallAsyncPatcher,
+    WorkflowNodeRunPatcher,
+    WorkflowToolNodeRunImplPatcher,
 )
 
 
@@ -29,5 +31,7 @@ class ADKIntegration(BaseIntegration):
         RunnerRunPatcher,
         FlowRunAsyncPatcher,
         ToolCallAsyncPatcher,
+        WorkflowNodeRunPatcher,
+        WorkflowToolNodeRunImplPatcher,
         McpToolPatcher,
     )

py/src/braintrust/integrations/adk/patchers.py

@@ -1,5 +1,6 @@
 """ADK patchers — one patcher per coherent patch target."""
 
+from importlib import import_module
 from typing import Any, ClassVar
 
 from braintrust.integrations.base import CompositeFunctionWrapperPatcher, FunctionWrapperPatcher
@@ -12,6 +13,8 @@
     _mcp_tool_run_async_wrapper_async,
     _runner_run_async_wrapper,
     _tool_call_async_wrapper,
+    _workflow_node_run_wrapper,
+    _workflow_tool_node_run_impl_wrapper,
 )
 
 
@@ -85,6 +88,29 @@ class FlowRunAsyncPatcher(CompositeFunctionWrapperPatcher):
     sub_patchers = (_FlowRunAsyncSubPatcher, _FlowCallLlmAsyncSubPatcher)
 
 
+# ---------------------------------------------------------------------------
+# ADK 2.x workflow patchers
+# ---------------------------------------------------------------------------
+
+
+class WorkflowNodeRunPatcher(FunctionWrapperPatcher):
+    """Patch ``BaseNode.run`` to trace workflow and graph node execution."""
+
+    name = "adk.workflow.node.run"
+    target_module = "google.adk.workflow._base_node"
+    target_path = "BaseNode.run"
+    wrapper = _workflow_node_run_wrapper
+
+
+class WorkflowToolNodeRunImplPatcher(FunctionWrapperPatcher):
+    """Patch ``_ToolNode._run_impl`` to trace direct workflow tool execution."""
+
+    name = "adk.workflow.tool_node.run_impl"
+    target_module = "google.adk.workflow._tool_node"
+    target_path = "_ToolNode._run_impl"
+    wrapper = _workflow_tool_node_run_impl_wrapper
+
+
 # ---------------------------------------------------------------------------
 # Tool patcher
 # ---------------------------------------------------------------------------
@@ -203,6 +229,36 @@ def wrap_runner(Runner: Any) -> Any:
     return RunnerRunPatcher.wrap_target(Runner)
 
 
+def wrap_workflow(Workflow: Any) -> Any:
+    """Manually patch ADK 2.x workflow nodes for tracing.
+
+    Pass the ``Workflow`` class from ``google.adk.workflow``. The base node
+    patch covers workflow and graph nodes; the tool-node patch adds spans for
+    tools executed directly by the workflow graph.
+    """
+    base_node = next(
+        (
+            cls
+            for cls in Workflow.__mro__
+            if cls.__name__ == "BaseNode" and cls.__module__.startswith("google.adk.workflow")
+        ),
+        None,
+    )
+    if base_node is None:
+        raise TypeError("wrap_workflow expects an ADK Workflow class")
+
+    WorkflowNodeRunPatcher.wrap_target(base_node)
+    try:
+        tool_node_module = import_module("google.adk.workflow._tool_node")
+    except ImportError:
+        return Workflow
+
+    tool_node = getattr(tool_node_module, "_ToolNode", None)
+    if tool_node is not None:
+        WorkflowToolNodeRunImplPatcher.wrap_target(tool_node)
+    return Workflow
+
+
 def wrap_flow(Flow: Any) -> Any:
     """Manually patch a flow class for tracing."""
     return FlowRunAsyncPatcher.wrap_target(Flow)

py/src/braintrust/integrations/adk/test_adk.py

@@ -525,6 +525,90 @@ def get_weather(location: str):
     assert tool_span["output"]["temperature"] == "72°F"
 
 
+@pytest.mark.asyncio
+@pytest.mark.skipif(ADK_VERSION < (2, 10, 0), reason="Workflow nodes require ADK 2.10+")
+async def test_adk_workflow_tool_node_is_traced(memory_logger):
+    """A graph ToolNode should emit tool and workflow node spans."""
+    from google.adk.tools.load_artifacts_tool import LoadArtifactsTool
+    from google.adk.workflow import START, Workflow
+
+    assert not memory_logger.pop()
+
+    tool = LoadArtifactsTool()
+    workflow = Workflow(
+        name="weather_workflow",
+        edges=[(START, tool)],
+    )
+    app_name = "workflow_weather_app"
+    user_id = "test-user"
+    session_id = "test-session-workflow"
+    runner = await _create_runner(workflow, app_name=app_name, user_id=user_id, session_id=session_id)
+
+    events = [
+        event
+        async for event in runner.run_async(
+            user_id=user_id,
+            session_id=session_id,
+            new_message=types.Content(role="user", parts=[types.Part(text='{"artifact_names": ["report.txt"]}')]),
+        )
+    ]
+
+    assert events
+    spans = memory_logger.pop()
+    invocation = next(row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]")
+    workflow_span = next(row for row in spans if row["span_attributes"]["name"] == "workflow [weather_workflow]")
+    tool_span = next(row for row in spans if row["span_attributes"]["type"] == "tool")
+
+    assert workflow_span["span_attributes"]["type"] == "task"
+    assert workflow_span["span_parents"] == [invocation["span_id"]]
+    assert tool_span["span_attributes"]["name"] == "tool [load_artifacts]"
+    assert tool_span["input"]["arguments"] == {"artifact_names": ["report.txt"]}
+    assert tool_span["output"]["artifact_names"] == ["report.txt"]
+    assert "temporarily inserted and removed" in tool_span["output"]["status"]
+    tool_node_span = next(row for row in spans if row["span_id"] == tool_span["span_parents"][0])
+    assert tool_node_span["span_attributes"]["name"] == "workflow_node [load_artifacts]"
+    assert tool_node_span["output"]["artifact_names"] == ["report.txt"]
+
+
+@pytest.mark.vcr
+@pytest.mark.asyncio
+@pytest.mark.skipif(ADK_VERSION < (2, 10, 0), reason="Workflow nodes require ADK 2.10+")
+async def test_adk_workflow_root_traces_recorded_agent_tool_calls(memory_logger):
+    """A real ADK LLM/tool turn stays nested under the Workflow task span."""
+    from google.adk.workflow import START, Workflow
+
+    assert not memory_logger.pop()
+
+    weather_agent = Agent(
+        name="weather_agent",
+        model=ADK_MODEL,
+        instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
+        tools=[get_weather],
+    )
+    workflow = Workflow(name="weather_workflow", edges=[(START, weather_agent)])
+    app_name = "workflow_weather_app"
+    user_id = "test-user"
+    session_id = "test-session-workflow-recorded"
+    runner = await _create_runner(workflow, app_name=app_name, user_id=user_id, session_id=session_id)
+    response = await _run_final_responses(
+        runner,
+        user_id=user_id,
+        session_id=session_id,
+        new_message=types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")]),
+    )
+
+    assert response
+    spans = memory_logger.pop()
+    invocation = next(row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]")
+    workflow_span = next(row for row in spans if row["span_attributes"]["name"] == "workflow [weather_workflow]")
+    agent_span = next(row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]")
+    tool_span = next(row for row in spans if row["span_attributes"]["name"] == "tool [get_weather]")
+    assert workflow_span["span_parents"] == [invocation["span_id"]]
+    assert agent_span["span_parents"] == [workflow_span["span_id"]]
+    assert tool_span["span_attributes"]["type"] == "tool"
+    assert tool_span["input"]["arguments"] == {"location": "San Francisco"}
+
+
 @pytest.mark.vcr
 @pytest.mark.asyncio
 async def test_adk_max_tokens_captures_content(memory_logger):
@@ -1102,3 +1186,13 @@ def test_auto_instrument_adk(self):
         from braintrust.integrations.test_utils import verify_autoinstrument_script
 
         verify_autoinstrument_script("test_auto_adk.py")
+
+
+class TestManualWrapADK:
+    """Tests the public manual wrapping helpers in an isolated process."""
+
+    @pytest.mark.skipif(ADK_VERSION < (2, 10, 0), reason="Workflow nodes require ADK 2.10+")
+    def test_manual_wrap_adk_workflow(self):
+        from braintrust.integrations.test_utils import verify_autoinstrument_script
+
+        verify_autoinstrument_script("test_manual_wrap_adk_workflow.py")