From 6ac39ce4d05bda18f037462a99fc7a76f2677794 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 8 Oct 2026 12:51:20 -0400 Subject: [PATCH 1/2] feat(adk): trace workflow nodes and tool calls --- ...root_traces_recorded_agent_tool_calls.yaml | 126 ++++++++++++++++++ .../integrations/adk/integration.py | 4 + .../braintrust/integrations/adk/patchers.py | 25 ++++ .../braintrust/integrations/adk/test_adk.py | 83 ++++++++++++ py/src/braintrust/integrations/adk/tracing.py | 98 ++++++++++++++ 5 files changed, 336 insertions(+) create mode 100644 py/src/braintrust/integrations/adk/cassettes/latest/test_adk_workflow_root_traces_recorded_agent_tool_calls.yaml diff --git a/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_workflow_root_traces_recorded_agent_tool_calls.yaml b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_workflow_root_traces_recorded_agent_tool_calls.yaml new file mode 100644 index 000000000..c0039ccd8 --- /dev/null +++ b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_workflow_root_traces_recorded_agent_tool_calls.yaml @@ -0,0 +1,126 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "What''s the weather in San Francisco?"}], + "role": "user"}], "systemInstruction": {"parts": [{"text": "You are a helpful + weather assistant. Use the get_weather tool to answer questions about weather.\n\nYou + are an agent. Your internal name is \"weather_agent\"."}], "role": "user"}, + "tools": [{"functionDeclarations": [{"description": "Get the weather for a location.", + "name": "get_weather", "parameters_json_schema": {"properties": {"location": + {"title": "Location", "type": "string"}}, "required": ["location"], "title": + "get_weatherParams", "type": "object"}}]}], "generationConfig": {}}' + headers: + Content-Type: + - application/json + user-agent: + - google-genai-sdk/2.28.0 gl-python/3.14.7 google-adk/2.11.0 gl-python/3.14.7 + x-goog-api-client: + - google-genai-sdk/2.28.0 gl-python/3.14.7 google-adk/2.11.0 gl-python/3.14.7 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"functionCall\": {\n \"name\": \"get_weather\",\n + \ \"args\": {\n \"location\": \"San Francisco\"\n + \ }\n }\n }\n ],\n \"role\": + \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n + \ \"finishMessage\": \"Model generated function call(s).\"\n }\n ],\n + \ \"usageMetadata\": {\n \"promptTokenCount\": 84,\n \"candidatesTokenCount\": + 16,\n \"totalTokenCount\": 100,\n \"promptTokensDetails\": [\n {\n + \ \"modality\": \"TEXT\",\n \"tokenCount\": 84\n }\n ],\n + \ \"serviceTier\": \"standard\"\n },\n \"modelVersion\": \"gemini-2.5-flash-lite\",\n + \ \"responseId\": \"DbjDav6UGPm639IP5bCO-Qk\"\n}\n" + headers: + Alt-Svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + Content-Type: + - application/json; charset=UTF-8 + Date: + - Mon, 05 Oct 2026 14:45:33 GMT + Server: + - scaffolding on HTTPServer2 + Server-Timing: + - gfet4t7; dur=285 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + X-Content-Type-Options: + - nosniff + X-Frame-Options: + - SAMEORIGIN + X-Gemini-Service-Tier: + - standard + X-XSS-Protection: + - '0' + content-length: + - '751' + status: + code: 200 + message: OK +- request: + body: '{"contents": [{"parts": [{"text": "What''s the weather in San Francisco?"}], + "role": "user"}, {"parts": [{"functionCall": {"args": {"location": "San Francisco"}, + "name": "get_weather"}}], "role": "model"}, {"parts": [{"functionResponse": + {"name": "get_weather", "response": {"location": "San Francisco", "temperature": + "72\u00b0F", "condition": "sunny"}}}], "role": "user"}], "systemInstruction": + {"parts": [{"text": "You are a helpful weather assistant. Use the get_weather + tool to answer questions about weather.\n\nYou are an agent. Your internal name + is \"weather_agent\"."}], "role": "user"}, "tools": [{"functionDeclarations": + [{"description": "Get the weather for a location.", "name": "get_weather", "parameters_json_schema": + {"properties": {"location": {"title": "Location", "type": "string"}}, "required": + ["location"], "title": "get_weatherParams", "type": "object"}}]}], "generationConfig": + {}}' + headers: + Content-Type: + - application/json + user-agent: + - google-genai-sdk/2.28.0 gl-python/3.14.7 google-adk/2.11.0 gl-python/3.14.7 + x-goog-api-client: + - google-genai-sdk/2.28.0 gl-python/3.14.7 google-adk/2.11.0 gl-python/3.14.7 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"text\": \"The weather in San Francisco is sunny + with a temperature of 72\xB0F.\"\n }\n ],\n \"role\": + \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n + \ }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 131,\n \"candidatesTokenCount\": + 17,\n \"totalTokenCount\": 148,\n \"promptTokensDetails\": [\n {\n + \ \"modality\": \"TEXT\",\n \"tokenCount\": 131\n }\n ],\n + \ \"serviceTier\": \"standard\"\n },\n \"modelVersion\": \"gemini-2.5-flash-lite\",\n + \ \"responseId\": \"DbjDaq_AKt_y-8YP87HZ6Q0\"\n}\n" + headers: + Alt-Svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + Content-Type: + - application/json; charset=UTF-8 + Date: + - Mon, 05 Oct 2026 14:45:33 GMT + Server: + - scaffolding on HTTPServer2 + Server-Timing: + - gfet4t7; dur=308 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + X-Content-Type-Options: + - nosniff + X-Frame-Options: + - SAMEORIGIN + X-Gemini-Service-Tier: + - standard + X-XSS-Protection: + - '0' + content-length: + - '617' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/adk/integration.py b/py/src/braintrust/integrations/adk/integration.py index ca4e5a4f7..4d6a2b923 100644 --- a/py/src/braintrust/integrations/adk/integration.py +++ b/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, ) diff --git a/py/src/braintrust/integrations/adk/patchers.py b/py/src/braintrust/integrations/adk/patchers.py index 7e6f302de..100e36e41 100644 --- a/py/src/braintrust/integrations/adk/patchers.py +++ b/py/src/braintrust/integrations/adk/patchers.py @@ -12,6 +12,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 +87,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 # --------------------------------------------------------------------------- diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 8907196f2..15803e081 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -525,6 +525,89 @@ 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]" + + +@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): diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index bf02aa057..e455a3aef 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -3,6 +3,7 @@ import asyncio import contextvars import inspect +import json import logging import time from contextlib import aclosing, contextmanager @@ -367,6 +368,103 @@ async def _agent_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwarg agent_span.log(output=last_event) +async def _workflow_node_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): + """Trace ADK 2.x workflow nodes, including the Workflow root node.""" + if any( + cls.__name__ == "BaseAgent" and cls.__module__.startswith("google.adk.agents") + for cls in instance.__class__.__mro__ + ): + # Newer ADK agent classes share BaseNode.run; their existing agent span + # is the canonical task span and must remain the direct child of Runner. + async with aclosing(wrapped(*args, **kwargs)) as agen: + async for event in agen: + yield event + return + + node_name = getattr(instance, "name", instance.__class__.__name__) + node_input = kwargs.get("node_input", args[1] if len(args) > 1 else None) + is_workflow = any( + cls.__name__ == "Workflow" and cls.__module__.startswith("google.adk.workflow") + for cls in instance.__class__.__mro__ + ) + span_name = f"workflow [{node_name}]" if is_workflow else f"workflow_node [{node_name}]" + + async def _trace(): + with _start_stream_span( + name=span_name, + type=SpanTypeAttribute.TASK, + input=node_input, + metadata={"node_name": node_name, "node_class": instance.__class__.__name__}, + ) as node_span: + last_event = None + async with aclosing(wrapped(*args, **kwargs)) as agen: + async for event in agen: + last_event = event + yield event + if last_event is not None: + node_span.log(output=getattr(last_event, "output", last_event)) + + async with aclosing(_trace()) as agen: + async for event in agen: + yield event + + +async def _workflow_tool_node_run_impl_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): + """Trace tools executed directly by ADK 2.x workflow tool nodes.""" + tool = getattr(instance, "tool", None) + if tool is None: + async with aclosing(wrapped(*args, **kwargs)) as agen: + async for event in agen: + yield event + return + + # MCP tools already have their own wrapper around run_async. + if getattr(tool.__class__, "__module__", "").startswith("google.adk.tools.mcp_tool"): + async with aclosing(wrapped(*args, **kwargs)) as agen: + async for event in agen: + yield event + return + + tool_name = getattr(tool, "name", tool.__class__.__name__) + node_input = kwargs.get("node_input", args[1] if len(args) > 1 else None) + tool_args = node_input + if hasattr(tool_args, "parts"): + tool_args = "".join(getattr(part, "text", "") or "" for part in (tool_args.parts or [])) + if hasattr(tool_args, "model_dump"): + tool_args = tool_args.model_dump() + if isinstance(tool_args, str): + try: + tool_args = json.loads(tool_args) + except json.JSONDecodeError: + pass + if not isinstance(tool_args, dict): + tool_args = {} + + async def _trace(): + with _start_stream_span( + name=f"tool [{tool_name}]", + type=SpanTypeAttribute.TOOL, + input={"tool_name": tool_name, "arguments": tool_args}, + metadata={"tool_class": tool.__class__.__name__}, + ) as tool_span: + last_output = None + try: + async with aclosing(wrapped(*args, **kwargs)) as agen: + async for event in agen: + if getattr(event, "output", None) is not None: + last_output = event.output + yield event + except Exception as error: + tool_span.log(error=error) + raise + if last_output is not None: + tool_span.log(output=last_output) + + async with aclosing(_trace()) as agen: + async for event in agen: + yield event + + async def _flow_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): async def _trace(): with _start_stream_span( From 1a91a89320b11a98418500264aba3b4270fff91f Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 8 Oct 2026 13:28:20 -0400 Subject: [PATCH 2/2] fix(adk): address workflow tracing review --- .../braintrust/integrations/adk/__init__.py | 4 +- .../braintrust/integrations/adk/patchers.py | 31 ++++++++++++ .../braintrust/integrations/adk/test_adk.py | 11 +++++ py/src/braintrust/integrations/adk/tracing.py | 10 ++-- .../test_manual_wrap_adk_workflow.py | 47 +++++++++++++++++++ 5 files changed, 98 insertions(+), 5 deletions(-) create mode 100644 py/src/braintrust/integrations/auto_test_scripts/test_manual_wrap_adk_workflow.py diff --git a/py/src/braintrust/integrations/adk/__init__.py b/py/src/braintrust/integrations/adk/__init__.py index bec4d7114..5421ac1f9 100644 --- a/py/src/braintrust/integrations/adk/__init__.py +++ b/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. diff --git a/py/src/braintrust/integrations/adk/patchers.py b/py/src/braintrust/integrations/adk/patchers.py index 100e36e41..606f01e1e 100644 --- a/py/src/braintrust/integrations/adk/patchers.py +++ b/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 @@ -228,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) diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index 15803e081..79588a46f 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -567,6 +567,7 @@ async def test_adk_workflow_tool_node_is_traced(memory_logger): 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 @@ -1185,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") diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index e455a3aef..ce6f4b352 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -396,13 +396,15 @@ async def _trace(): input=node_input, metadata={"node_name": node_name, "node_class": instance.__class__.__name__}, ) as node_span: - last_event = None + last_output = None async with aclosing(wrapped(*args, **kwargs)) as agen: async for event in agen: - last_event = event + event_output = getattr(event, "output", None) + if event_output is not None: + last_output = event_output yield event - if last_event is not None: - node_span.log(output=getattr(last_event, "output", last_event)) + if last_output is not None: + node_span.log(output=last_output) async with aclosing(_trace()) as agen: async for event in agen: diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_manual_wrap_adk_workflow.py b/py/src/braintrust/integrations/auto_test_scripts/test_manual_wrap_adk_workflow.py new file mode 100644 index 000000000..ea72b3a97 --- /dev/null +++ b/py/src/braintrust/integrations/auto_test_scripts/test_manual_wrap_adk_workflow.py @@ -0,0 +1,47 @@ +"""Verify manual ADK wrapping covers Workflow nodes and direct tool calls.""" + +import asyncio + +from braintrust import logger +from braintrust.integrations.adk import wrap_runner, wrap_workflow +from braintrust.integrations.adk.patchers import WorkflowNodeRunPatcher +from braintrust.test_helpers import init_test_logger +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.adk.tools.load_artifacts_tool import LoadArtifactsTool +from google.adk.workflow import START, BaseNode, Workflow +from google.genai import types + + +async def main(): + app_name = "manual_workflow_app" + user_id = "test-user" + session_id = "manual-workflow-session" + session_service = InMemorySessionService() + await session_service.create_session(app_name=app_name, user_id=user_id, session_id=session_id) + runner = Runner( + agent=Workflow(name="manual_workflow", edges=[(START, LoadArtifactsTool())]), + app_name=app_name, + session_service=session_service, + ) + async for _ 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"]}')]), + ): + pass + + +wrap_runner(Runner) +wrap_workflow(Workflow) +assert getattr(BaseNode, WorkflowNodeRunPatcher.patch_marker_attr(), False) + +init_test_logger("manual-adk-workflow") +with logger._internal_with_memory_background_logger() as memory_logger: + asyncio.run(main()) + spans = memory_logger.pop() + +assert any(span["span_attributes"]["name"] == "invocation [manual_workflow_app]" for span in spans) +assert any(span["span_attributes"]["name"] == "workflow [manual_workflow]" for span in spans) +assert any(span["span_attributes"]["name"] == "tool [load_artifacts]" for span in spans) +print("SUCCESS")