[AAASM-1685] ✨ (adapters/google_adk): Patch BaseAgent.run_async for SpawnContext propagation#49
Merged
Chisanan232 merged 3 commits intoMay 21, 2026
Conversation
…text wrapper Adds module constants (`_ORIGINAL_AGENT_RUN_ASYNC`, `_AGENT_PATCHED_FLAG`) and three helpers: * `_load_google_adk_base_agent_class()` — imports `google.adk.agents` and returns its `BaseAgent` class, or None if the framework is not installed. * `_apply_agent_run_async_patch(agent_cls, process_agent_id)` — replaces `BaseAgent.run_async` with an async-generator wrapper that pushes a `SpawnContext(spawned_by_tool="google_adk_agent", depth=current+1)` via `spawn_context_scope()` for the lifetime of the generator. Idempotent. * `_revert_agent_run_async_patch(agent_cls)` — restores the original method and clears the patch-flag attribute. No-op when not previously applied. Helpers are not yet called from `GoogleADKPatch.apply()` / `revert()` — that wiring lands in the next commit so this one is bisect-safe in isolation. Refs: AAASM-1685
…ly/revert `apply()` now calls `_apply_agent_run_async_patch()` after the existing tool-patch step (gracefully no-ops if `google.adk.agents` is not importable); `revert()` calls `_revert_agent_run_async_patch()` before the tool revert. Together this gives Google ADK feature parity with `PydanticAIAdapter`, which wraps `Agent.run` / `Agent.run_sync` for the same purpose. The tool-patch returning False short-circuit is preserved: if the `google-adk` framework is not installed at all, both branches no-op and `apply()` returns False as before. Refs: AAASM-1685
…-context patch Nine new tests covering the agent-level patch surface: * `_apply_agent_run_async_patch` is idempotent and rebinds `run_async`, * `_revert_agent_run_async_patch` restores the original and clears the patch-flag attribute (and is a no-op when not previously applied), * the patched generator sets `SpawnContext(spawned_by_tool= "google_adk_agent", parent_agent_id=<process_agent_id>, depth=1)` and the context is observable from inside every `yield` of the original, * the patched generator yields every event from the original (no drops), * `_load_google_adk_base_agent_class` returns None when `google.adk.agents` is missing OR when `BaseAgent` is not a class, * `GoogleADKPatch.apply()` / `revert()` end-to-end with both `BaseTool` and `BaseAgent` modules monkey-patched into place, * `GoogleADKPatch.apply()` still succeeds when only the tool module is importable (agent module missing). All tests hermetic via `monkeypatch.setattr(google_adk_patch.importlib, "import_module", ...)` — no real `google-adk` install required. Refs: AAASM-1685
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.



Description
Closes the one "optionally" item from the AAASM-1550 Story description that was deferred past the original 10/10 AC close: extends
GoogleADKPatchto wrapgoogle.adk.agents.BaseAgent.run_asyncwith aSpawnContextscope so that spawn-lineage (parent agent id + depth +spawned_by_tool="google_adk_agent") propagates through every yieldedEventduring an agent run.Closes a parity gap with
PydanticAIAdapter, which already wrapsAgent.run/Agent.run_syncfor the same purpose.Implementation note
BaseAgent.run_asyncreturnsAsyncIterator[Event](an async generator), not a plain awaitable. The wrapper is therefore itself an async generator:The
withblock stays alive across everyyield(untilStopAsyncIterationoraclose()), so consumers iterating the generator always observe the correct SpawnContext — verified intest_patched_run_async_sets_spawn_context_during_iteration.Type of Change
Breaking Changes
Related Issues
agent_assembly/adapters/pydantic_ai/patch.py:_apply_agent_run_patchTesting
test_google_adk_real_base_tool_class_patch_path_when_availablealready covers the real-class path when google-adk is installed; no new integration test needed)Nine new unit tests in
test/unit/adapters/google_adk/test_google_adk_patch.pycover:test_apply_agent_patches_run_async_and_is_idempotent_apply_*is a no-op; flag set; method reboundtest_revert_agent_patch_restores_run_async_and_clears_flagtest_revert_agent_patch_is_noop_when_not_patchedtest_patched_run_async_sets_spawn_context_during_iteration_SPAWN_CTX.get()non-None at every yield;spawned_by_tool == "google_adk_agent";parent_agent_id == "parent-1";depth == 1test_patched_run_async_yields_all_events_from_originaltest_load_base_agent_returns_none_when_module_missingtest_load_base_agent_returns_none_when_attribute_not_typeBaseAgentnot a class → Nonetest_apply_patches_both_tool_and_agent_when_both_availableGoogleADKPatch.apply()end-to-end with both modules installedtest_apply_proceeds_with_only_tool_when_agent_module_missingLocal run:
```text
test/unit/adapters/google_adk/test_google_adk_patch.py ........ 18 passed (9 original + 9 new)
test/integration/test_google_adk_interception_integration.py 1 passed, 1 skipped (importorskip)
Full suite: 362 passed, 11 skipped (up from 351 — exactly +11 new tests, no regression)
```
mypy agent_assembly/adapters/google_adk→ no issues.isort --check --profile=black→ pass.Checklist
🤖 Generated with Claude Code