Skip to content

[AAASM-1685] ✨ (adapters/google_adk): Patch BaseAgent.run_async for SpawnContext propagation#49

Merged
Chisanan232 merged 3 commits into
masterfrom
v0.0.1/AAASM-1685/feat/google_adk_agent_spawn_patch
May 21, 2026
Merged

[AAASM-1685] ✨ (adapters/google_adk): Patch BaseAgent.run_async for SpawnContext propagation#49
Chisanan232 merged 3 commits into
masterfrom
v0.0.1/AAASM-1685/feat/google_adk_agent_spawn_patch

Conversation

@Chisanan232
Copy link
Copy Markdown
Contributor

Description

Closes the one "optionally" item from the AAASM-1550 Story description that was deferred past the original 10/10 AC close: extends GoogleADKPatch to wrap google.adk.agents.BaseAgent.run_async with a SpawnContext scope so that spawn-lineage (parent agent id + depth + spawned_by_tool="google_adk_agent") propagates through every yielded Event during an agent run.

Closes a parity gap with PydanticAIAdapter, which already wraps Agent.run / Agent.run_sync for the same purpose.

Implementation note

BaseAgent.run_async returns AsyncIterator[Event] (an async generator), not a plain awaitable. The wrapper is therefore itself an async generator:

async def patched_run_async(self, *args, **kwargs):
    spawn_ctx = SpawnContext(parent_agent_id=..., depth=..., spawned_by_tool="google_adk_agent")
    with spawn_context_scope(spawn_ctx):
        async for event in original_run_async(self, *args, **kwargs):
            yield event

The with block stays alive across every yield (until StopAsyncIteration or aclose()), so consumers iterating the generator always observe the correct SpawnContext — verified in test_patched_run_async_sets_spawn_context_during_iteration.

Type of Change

  • ✨ New feature
  • 🔧 Bug fix
  • ♻️ Refactoring
  • 🍀 Performance improvement
  • 📚 Documentation update
  • 🚀 Release

Breaking Changes

  • No

Related Issues

  • Related JIRA ticket: AAASM-1685 (parent: AAASM-1550, epic: AAASM-5)
  • Sibling adapter for pattern reference: agent_assembly/adapters/pydantic_ai/patch.py:_apply_agent_run_patch

Testing

  • Unit tests added/updated
  • Integration tests added/updated (the existing test_google_adk_real_base_tool_class_patch_path_when_available already covers the real-class path when google-adk is installed; no new integration test needed)
  • Manual testing performed
  • No tests required

Nine new unit tests in test/unit/adapters/google_adk/test_google_adk_patch.py cover:

Test Verifies
test_apply_agent_patches_run_async_and_is_idempotent second _apply_* is a no-op; flag set; method rebound
test_revert_agent_patch_restores_run_async_and_clears_flag original restored; flag attribute deleted
test_revert_agent_patch_is_noop_when_not_patched revert without apply is harmless
test_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 == 1
test_patched_run_async_yields_all_events_from_original no events dropped by the wrapper
test_load_base_agent_returns_none_when_module_missing ImportError → None
test_load_base_agent_returns_none_when_attribute_not_type BaseAgent not a class → None
test_apply_patches_both_tool_and_agent_when_both_available GoogleADKPatch.apply() end-to-end with both modules installed
test_apply_proceeds_with_only_tool_when_agent_module_missing tool-only path still returns True; revert clean

Local 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

  • Code follows project style guidelines
  • Self-review completed
  • Comments added for complex logic (async-generator wrapping)
  • Documentation updated if needed
  • All tests passing

🤖 Generated with Claude Code

…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
Copy link
Copy Markdown

codecov Bot commented May 21, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@sonarqubecloud
Copy link
Copy Markdown

@Chisanan232 Chisanan232 merged commit 9941634 into master May 21, 2026
24 checks passed
@Chisanan232 Chisanan232 deleted the v0.0.1/AAASM-1685/feat/google_adk_agent_spawn_patch branch May 21, 2026 05:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant