fix(google_adk): support google-adk 2.7.0+ tool dispatch changes [MLOS-864] - #19848
fix(google_adk): support google-adk 2.7.0+ tool dispatch changes [MLOS-864]#19848heyitsgrace996 wants to merge 5 commits into
Conversation
google-adk 2.7.0 removed the internal `__call_tool_live` function and routes live tool execution through `__call_tool_async` instead. The integration patched `__call_tool_live` unconditionally at import time, so applications enabling LLM Observability (which patches with raise_errors=True) crashed on startup with an AttributeError and could not boot at all. Only patch the tool dispatch functions the installed version defines, reusing the check_module_path guard already applied to the code executors. Also fixes three defects the startup crash was masking on 2.7.0+: - The live dispatch path passes every argument by keyword, so reading the tool from `args[0]` raised IndexError. Use get_argument_value instead. - Streaming tools return an async generator from `__call_tool_async`. The span was tagged with the generator object and finished before any chunk was produced. Keep the span open until the stream is exhausted and tag the streamed items, capped at MAX_STREAMED_TOOL_CHUNKS. - The agent-less fallback returned an un-awaited coroutine. `_traced_functions_call_tool_live` also never assigned its accumulated result, so every live tool span on google-adk < 2.7 recorded a null output. Re-enable the tool dispatch assertions in test_google_adk_patch.py, which were commented out due to class-body name mangling, and add coverage for the missing symbol, the keyword-only dispatch path and the streamed output tagging. Pin the riot min version explicitly and regenerate the lockfiles so the latest slot resolves to google-adk 2.7.1 instead of the stale 1.28.1.
Address review feedback on the streaming tool wrappers. `_traced_tool_stream` iterated the generator returned by the tool but never closed it. When a consumer stops early it closes the ddtrace wrapper, leaving the wrapped generator to wait for async generator finalization. google-adk closes the stream itself with `Aclosing`, so mirror that and close it in the wrapper's `finally`. `contextlib.aclosing` is 3.10+, so this calls `aclose()` directly. The pre-2.7 `__call_tool_live` wrapper had the same gap. Add the two streaming lifetime tests that were missing: a tool that raises partway through the stream, and a consumer that stops early and closes it. Both assert the span is still finished. Also skip content-less events in test_agent_run_async. google-adk synthesizes an error event with no content in `workflow/_node_runner.py` before re-raising, so the unguarded loop raised AttributeError on google-adk >= 2.6.3. This is the guard that `test_agent_with_tool_calculation` in the same file and google-adk's own `runners.py` already use; it needs no cassette re-recording.
🎉 All green!🧪 All tests passed 🔄 Datadog auto-retried 5 jobs - 5 passed on retry 🔗 Commit SHA: 92ce7d0 | Docs | View more details | Give us feedback! |
This comment was marked as resolved.
This comment was marked as resolved.
Circular import analysis
|
Dependency direction analysis
|
BenchmarksBenchmark execution time: 2026-08-26 18:08:36 Comparing candidate commit 92ce7d0 in PR branch Found 0 performance improvements and 8 performance regressions! Performance is the same for 578 metrics, 10 unstable metrics, 2 known flaky benchmarks, 16 flaky benchmarks without significant changes.
|
|
This change is marked for backport to 4.14 and it does not conflict with that branch. |
Follow-up review feedback on the streaming tool wrappers. Guard the `aclose()` added in the previous commit. If closing the wrapped generator raised, it both masked the exception the stream was already raising and skipped the tagging and `span.finish()` below it, stranding the span. Build the stream before starting the span in `_traced_functions_call_tool_live`. `integration.trace()` activates by default, so a span created before a call that raises is never finished and reparents everything after it. Evaluating `wrapped()` first means there is no such window rather than a caught one. Fold that wrapper onto `_traced_tool_stream` instead of repeating the accumulate, cap, close, tag and finish sequence. The `with` form is dropped because `_traced_tool_stream` finishes the span itself, and `Span.__exit__` only sets exception info and finishes. This also puts the pre-2.7 path on the implementation the 2.7+ tests exercise, since the streaming tests are gated above 2.7.0. Move MAX_STREAMED_TOOL_CHUNKS above its use sites. Consolidate the release note onto the two symptoms a customer can observe: the failure to start, and streaming tool spans recording no output. The removed entries described internal functions and paths unreachable behind the startup failure.
heyitsgrace996
changed the title
The cap and its omitted-items marker were added earlier in this branch and never shipped. Nothing else in the repo bounds what it retains from a stream, including `_traced_agent_run_async` in this file, so match that and keep every streamed item.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 351d9a8c64
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| try: | ||
| result = await wrapped(*args, **kwargs) | ||
| except Exception: |
There was a problem hiding this comment.
When an ADK tool task is cancelled while wrapped is awaited—for example because of a request timeout or client disconnect—asyncio.CancelledError inherits from BaseException on every supported Python version, so this handler is skipped. Unlike the previous with integration.trace(...) implementation, the span is never tagged or finished and remains active, which can prevent the trace from being emitted and incorrectly parent later spans if execution continues; finish the span in a finally block or explicitly handle cancellation before re-raising it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
fixed - cleanup moved into a finally block with a hand-off flag
| async for item in _traced_tool_stream(agen, span, integration, args, kwargs): | ||
| yield item |
There was a problem hiding this comment.
On google-adk versions before 2.7, if a consumer stops a live tool stream early and closes _traced_functions_call_tool_live, Python does not automatically close the inner async generator used by this async for. Consequently _traced_tool_stream remains suspended at its yield, its finally block does not run, and the tool span remains active and unfinished; explicitly close the delegated iterator from an outer finally block (or use an equivalent async-closing construct).
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed - the stream is now closed in an outer finally
brettlangdon
left a comment
brettlangdon
left a comment
There was a problem hiding this comment.
nit/question on the release note, otherwise release note + riotfile.py changes lgtm
…eardown Address two P1 review findings on the streaming tool wrappers. `asyncio.CancelledError` is a `BaseException`, so `except Exception` skipped it and the span was neither tagged nor finished. Because `integration.trace()` activates the span, it also stayed active and reparented everything after it. Move the cleanup into a `finally` guarded by an explicit hand-off flag, so every exit finishes the span except the streaming one, which `_traced_tool_stream` owns. This is the shape `vertexai` already uses. Cancellation is not recorded as a tool error, matching `django` and `langgraph`, which treat it as ordinary teardown. Close the stream `_traced_functions_call_tool_live` delegates to. Abandoning the outer generator does not close the inner one, so its cleanup waited on async generator finalization: the span took four event loop iterations to finish where `main` took two. Closing it explicitly restores parity. Also address review feedback on the release note: shorten the first entry and state the affected version range on both, since the two fixes apply to opposite sides of 2.7.0 and read as one fix without it.
| # open so the streamed items are tagged instead of the generator object. | ||
| stream = _traced_tool_stream(result, span, integration, args, kwargs) | ||
| stream_handed_off = True |
There was a problem hiding this comment.
Streaming tool span is leaked when the returned generator is never started
stream_handed_off = True skips the finally, so the span is only finished if _traced_tool_stream actually runs. A never-started async generator never runs its finally, neither on aclose() nor on GC, so span.finish() is never called. The span is left unfinished, unflushed, and still activated.
Two reachable paths in google-adk 2.7.1:
- Non-live dispatch.
_execute_single_function_call_asyncawaits__call_tool_asyncatfunctions.py:627with noisasyncgencheck (that check only exists atfunctions.py:1105, on the live path). An async generatorFunctionToolresult goes straight into_normalize_tool_resultas{'result': <agen>}and is never iterated or closed. - Cancellation landing between the
awaitand the firstasync for.
Pre-diff this could not happen, since with integration.trace(...) closed the span on every exit path.
Fix options, roughly in order of preference:
A. Backstop the handoff with a finalizer. Keeps the current shape and covers both paths:
stream = _traced_tool_stream(result, span, integration, args, kwargs)
weakref.finalize(stream, _finish_unstarted_span, span, integration, args, kwargs)where _finish_unstarted_span returns early unless inspect.getasyncgenstate(...) was AGEN_CREATED. Duration is meaningless in that case, but the span is no longer lost. Requires _traced_tool_stream to mark the span as claimed on first iteration so the two do not race.
B. Only hand off on a stream the caller will consume. Check inspect.getasyncgenstate(result) is not AGEN_CREATED at tag time, or gate the async generator branch so the non-live call site falls through to the existing finally and tags the generator object as it did before this PR. Simpler, but path 2 still leaks.
C. Return a proxy instead of a new generator. Wrap result in a wrapt.ObjectProxy that finishes the span on exhaustion, on aclose(), and on finalization. More code, but it also preserves the wrapped stream's identity and attributes, which the current substitution drops.
Worth a regression test for the abandoned-stream case: the three new streaming tests all consume or aclose() the result, so none of them would catch this.
3 participants
Description
google-adk 2.7.0 merged
__call_tool_liveinto__call_tool_asyncand deleted the former. We wrap both by name at startup, so any app with LLM Observability enabled and google-adk 2.7.0+ installed failed to boot:google-adk change ref: google/adk-python@8b9d222
What Changed
This PR does 4 things:
Fixes the crash: Only wrap the tool dispatch functions the installed version actually defines, in both
patch()andunpatch().Fixes underlying tool stream bugs: Three bugs on the tool path nobody could reach while startup was dying: the tool read from
args[0]when ADK calls with keywords only, streaming tools tagged as a generator object with the span closed before anything was produced, and an un-awaited coroutine in the agent-less fallback.Bumps test versioning: Lockfiles were frozen on google-adk 1.28.1, so 2.x was never tested. Regenerates the lockfiles and pins the floor at
==1.0.0alongsidelatest, so the oldest supported version is always covered.Repairs two existing tests: The tool wrapping assertions were commented out with a
# TODO: fix thisand would have caught this.test_agent_run_asyncwas failing on google-adk >= 2.6.3.Testing
monkeypatch, so it holds on any ADK version), keyword-only dispatch, and streaming span lifetime.Risks
Low, confined to the google_adk integration. The streaming path only runs on >= 2.7.0; older versions are unchanged.