Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '561'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61SXU+DMBR9768gfR4LMofG100TP4hTyWJijLnChTWWFttuy1z23y1fG0x9ExJS
7jk9596ebonj0BhEwhIwqOmF82IrjrOtviUmhUFhLNCWbLEAZQ7c+tl21paSLkVsmBQT4Ly3ucEF
5GjrNEPztkYwC1R0cEwClelfNluEyxhK+VLiCYRzpUDETMeSHnF35K+/w/r1YEyV5FVfuUyQt2K7
lkBTJphePCLoxju6n+37prDK7mRWKPletu2eDwPf98bBaeD59h2fjND1AtKaV7Z0qSHDEA3YAGA/
LLUieWEi+YFiIpdVAOOz2qiTVw8PGthIA7yPjAY/VPXUejLejbGTsB0fODObcsbo8jnqZGP1e021
Z0Q6R3nc4j+ZBX0v0iRThzVHpZsbkWFuc3L9oeemHPSiEqQKdSGFxuuk5Ez9NITwFqfhav1p9EzG
XzebB4+SHfkGVtQS/hUDAAA=
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:51 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=530
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
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", "humidity": "45%", "wind": "5 mph NW"}}}],
"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": {"properties": {"location": {"type": "STRING"}},
"required": ["location"], "type": "OBJECT"}}]}], "generationConfig": {}}'
headers:
accept:
- '*/*'
accept-encoding:
- gzip, deflate
connection:
- keep-alive
content-length:
- '881'
content-type:
- application/json
host:
- generativelanguage.googleapis.com
user-agent:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
x-goog-api-client:
- google-genai-sdk/1.31.0 gl-python/3.9.21 google-adk/1.14.1 gl-python/3.9.21
method: POST
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
response:
body:
string: !!binary |
H4sIAAAAAAAC/61R0U7bQBB891esTuobiS5OHELfqqaRUIEisACprdA2Xsen2nfmbl2IovwT38CX
cefUwaGv9YO12pnbGc1sIgCxRJ2pDJmc+Ajf/QZg0/4DZjSTZg90K7+s0fIbd/dterOnMD2FRyIt
CB4JuSALSsM1alhY1EvllgaUA9dovYZHxQUgMFU1WeTGEpgcjuOX58UQwomiqVSmeB2eTJIP4B0D
h9PKD36XQFUXcHE7/KFFz8h2P/88erNvTUnBW2UyKjv6tiOIXGnliitCZ3SgXaffLsUexT+rM7Oq
rfkVEhjIoZQymU1HUo6n8iSeTpKTWRx14q2saByu6JwYfci4j1L4I1XNqflN+rNp2pBnk51Qr5MD
fNzhbBjLA2g0mh39c9fNvaoq+2X1evQBYOlTbYv6cpeKXkh8aKtLKeqF+d7kfxIbvxOL/paz6+uG
rFO7YlZU+aoG8VAO8hJd0V4UllxttKPTLHDmcX6Op+nXi3v19MDu8j6npPkkRbSNXgFas32N/AIA
AA==
headers:
Alt-Svc:
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
Content-Encoding:
- gzip
Content-Type:
- application/json; charset=UTF-8
Date:
- Thu, 18 Sep 2025 20:09:52 GMT
Server:
- scaffolding on HTTPServer2
Server-Timing:
- gfet4t7; dur=612
Transfer-Encoding:
- chunked
Vary:
- Origin
- X-Origin
- Referer
X-Content-Type-Options:
- nosniff
X-Frame-Options:
- SAMEORIGIN
X-XSS-Protection:
- '0'
status:
code: 200
message: OK
version: 1
4 changes: 2 additions & 2 deletions py/src/braintrust/integrations/adk/integration.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,7 +8,7 @@
AgentRunAsyncPatcher,
FlowRunAsyncPatcher,
McpToolPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
ThreadBridgePatcher,
ToolCallAsyncPatcher,
)
Expand All@@ -26,7 +26,7 @@ class ADKIntegration(BaseIntegration):
patchers = (
ThreadBridgePatcher,
AgentRunAsyncPatcher,
RunnerRunSyncPatcher,
RunnerRunPatcher,
FlowRunAsyncPatcher,
ToolCallAsyncPatcher,
McpToolPatcher,
Expand Down
25 changes: 10 additions & 15 deletions py/src/braintrust/integrations/adk/patchers.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -11,7 +11,6 @@
_flow_run_async_wrapper,
_mcp_tool_run_async_wrapper_async,
_runner_run_async_wrapper,
_runner_run_wrapper,
_tool_call_async_wrapper,
)

Expand All@@ -31,19 +30,10 @@ class AgentRunAsyncPatcher(FunctionWrapperPatcher):


# ---------------------------------------------------------------------------
# Runner patchers (sync + async)
# Runner patchers
# ---------------------------------------------------------------------------


class _RunnerRunSyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync generator)."""

name = "adk.runner.run.sync"
target_module = "google.adk.runners"
target_path = "Runner.run"
wrapper = _runner_run_wrapper


class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
"""Patch ``Runner.run_async`` (async generator)."""

Expand All@@ -53,11 +43,16 @@ class _RunnerRunAsyncSubPatcher(FunctionWrapperPatcher):
wrapper = _runner_run_async_wrapper


class RunnerRunSyncPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run`` (sync) and ``Runner.run_async`` for tracing."""
class RunnerRunPatcher(CompositeFunctionWrapperPatcher):
"""Patch ``Runner.run_async`` for tracing.

``Runner.run()`` already delegates to ``run_async()`` in supported ADK
versions, so tracing the async surface alone gives sync and async callers a
single invocation span with the same child structure.
"""

name = "adk.runner.run"
sub_patchers = (_RunnerRunSyncSubPatcher, _RunnerRunAsyncSubPatcher)
sub_patchers = (_RunnerRunAsyncSubPatcher,)


# ---------------------------------------------------------------------------
Expand DownExpand Up@@ -161,7 +156,7 @@ def wrap_agent(Agent: Any) -> Any:

def wrap_runner(Runner: Any) -> Any:
"""Manually patch a runner class for tracing."""
return RunnerRunSyncPatcher.wrap_target(Runner)
return RunnerRunPatcher.wrap_target(Runner)


def wrap_flow(Flow: Any) -> Any:
Expand Down
77 changes: 67 additions & 10 deletions py/src/braintrust/integrations/adk/test_adk.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -61,6 +61,17 @@ async def _create_runner(agent: Agent, *, app_name: str, user_id: str, session_i
return Runner(agent=agent, app_name=app_name, session_service=session_service)


def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}


def _extract_text_parts(contents):
texts = []
for content in contents or []:
Expand DownExpand Up@@ -333,21 +344,67 @@ async def generate_content_async(self, llm_request: LlmRequest, stream: bool = F
assert llm_document_part["image_url"]["url"].reference["content_type"] == "application/pdf"


@pytest.mark.vcr
def test_adk_sync_runner_run_does_not_duplicate_invocation_spans(memory_logger):
"""Runner.run() should emit a single invocation span even though it delegates to run_async()."""
import asyncio

from braintrust.util import LazyValue

assert not memory_logger.pop()

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
instruction="You are a helpful weather assistant. Use the get_weather tool to answer questions about weather.",
tools=[get_weather],
)

app_name = "weather_app"
user_id = "test-user"
session_id = "test-session"

runner = asyncio.run(_create_runner(agent, app_name=app_name, user_id=user_id, session_id=session_id))
user_msg = types.Content(role="user", parts=[types.Part(text="What's the weather in San Francisco?")])

# The memory_logger fixture overrides via thread-local (_override_bg_logger),
# but Runner.run() dispatches to a background thread where that's invisible.
# We must also set _global_bg_logger so spans emitted on the worker thread
# are captured.
original_global_bg_logger = logger._state._global_bg_logger
logger._state._global_bg_logger = LazyValue(lambda: memory_logger, use_mutex=False)
try:
responses = [
event
for event in runner.run(user_id=user_id, session_id=session_id, new_message=user_msg)
if event.is_final_response()
]
finally:
logger._state._global_bg_logger = original_global_bg_logger

assert responses
spans = memory_logger.pop()

invocation_spans = [row for row in spans if row["span_attributes"]["name"] == f"invocation [{app_name}]"]
assert len(invocation_spans) == 1, (
f"expected exactly one invocation span for Runner.run(), got {len(invocation_spans)}: "
f"{[span['span_id'] for span in invocation_spans]}"
)

invocation_span = invocation_spans[0]
agent_spans = [row for row in spans if row["span_attributes"]["name"] == "agent_run [weather_agent]"]
assert len(agent_spans) == 1
assert invocation_span["span_id"] in agent_spans[0].get("span_parents", []), (
f"agent span should be parented to the single sync invocation span {invocation_span['span_id']}, "
f"got parents {agent_spans[0].get('span_parents')}"
)


@pytest.mark.vcr
@pytest.mark.asyncio
async def test_adk_braintrust_integration(memory_logger):
assert not memory_logger.pop()

def get_weather(location: str):
"""Get the weather for a location."""
return {
"location": location,
"temperature": "72°F",
"condition": "sunny",
"humidity": "45%",
"wind": "5 mph NW",
}

agent = Agent(
name="weather_agent",
model="gemini-2.0-flash",
Expand Down
32 changes: 0 additions & 32 deletions py/src/braintrust/integrations/adk/tracing.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -478,38 +478,6 @@ async def _trace():
yield event


def _runner_run_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
new_message = kwargs.get("new_message")

# Serialize new_message before any dict conversion to handle binary data
serialized_message = _serialize_content(new_message) if new_message else None

def _trace():
with start_span(
name=f"invocation [{instance.app_name}]",
type=SpanTypeAttribute.TASK,
input={"new_message": serialized_message},
metadata=bt_safe_deep_copy(
{
"user_id": user_id,
"session_id": session_id,
**_omit(kwargs, ["user_id", "session_id", "new_message"]),
}
),
) as runner_span:
last_event = None
for event in wrapped(*args, **kwargs):
if event.is_final_response():
last_event = event
yield event
if last_event:
runner_span.log(output=last_event)

yield from _trace()


async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any):
user_id = kwargs.get("user_id")
session_id = kwargs.get("session_id")
Expand Down
Original file line numberDiff line numberDiff line change
Expand Up@@ -7,7 +7,6 @@
from braintrust.integrations.adk.patchers import (
AgentRunAsyncPatcher,
_RunnerRunAsyncSubPatcher,
_RunnerRunSyncSubPatcher,
_ThreadBridgePlatformSubPatcher,
_ThreadBridgeRunnersSubPatcher,
)
Expand All@@ -27,7 +26,6 @@ def is_patched(target, patcher):

# 1. Verify ADK surfaces are not patched initially.
assert not is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert not is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert not is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert not is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)
Expand All@@ -38,16 +36,16 @@ def is_patched(target, patcher):

# 3. Verify the imported google.adk surfaces are patched.
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)
assert is_patched(platform_thread.create_thread, _ThreadBridgePlatformSubPatcher)
assert is_patched(adk_runners.create_thread, _ThreadBridgeRunnersSubPatcher)

# 4. Idempotent.
results2 = auto_instrument()
assert results2.get("adk") == True, "auto_instrument should still return True on second call"
assert is_patched(BaseAgent.run_async, AgentRunAsyncPatcher)
assert is_patched(Runner.run, _RunnerRunSyncSubPatcher)
assert is_patched(Runner.run_async, _RunnerRunAsyncSubPatcher)
assert not is_patched(Runner.run, _RunnerRunAsyncSubPatcher)

print("SUCCESS")
Loading