fix(agents): close the inner run_async when an agent node stops early

[AtulJoshi1206]

Update (Oct 17): rebased onto current main, plus one more commit.

• Since 10c694f (stop holding OTel context across async generator yields), the Failed to detach context traceback described below no longer appears. The underlying issue is still there: when _run_impl is closed early, run_async is not closed with it, so its finally blocks and after_agent_callback unwinding are deferred to the asyncgen finalizer hook. test_run_impl_closes_run_async_when_consumer_stops_early still fails on current main (assert [] == [True]).
• ManagedAgent._run_impl and AntigravityAgent._run_impl mirror BaseAgent._run_impl so they can thread node_input through, and they had the same unclosed iteration. The second commit applies the same Aclosing there, with regression tests for both the classic path and the node_input path.
• Tests: tests/unittests/agents, workflow and runners all pass (2183 passed). For the Antigravity tests I used google-antigravity==0.1.7 with protobuf 6.33: the newest google-antigravity needs protobuf >= 7.35, and the test extra caps protobuf below 7, so it doesn't import in the standard dev env. In that env, 48 tests in tests/unittests/labs/antigravity also fail on main (Step has no trajectory_id). This change leaves all of them unchanged, and its 2 new tests fail on main and pass with the fix.

---

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

No existing issue; described below following the bug-report structure.

Describe the Bug:

BaseAgent._run_impl, the node-mode entry point, iterates run_async without closing it:

async for event in self.run_async(
    parent_context=ctx.get_invocation_context()
):
  ...
  yield event

Every other consumer of an agent generator in the repo wraps it in Aclosing, including BaseNode.run, which is this method's only caller, and LlmAgent._run_impl, the override of this very method. Aclosing is already imported in base_agent.py.

When BaseNode.run's Aclosing closes _run_impl early, GeneratorExit is raised at the yield, the frame unwinds, and the inner run_async generator is dropped rather than closed. That generator is suspended inside record_agent_invocation, which holds an OpenTelemetry span, and inside the try/except asyncio.CancelledError that runs after_agent_callback.

Impact:

Its cleanup is deferred to the async-generator finalizer hook, which resumes it as a separate task in a different contextvars context. That is not merely late, it fails:

ERROR opentelemetry.context: Failed to detach context
Traceback (most recent call last):
  File ".../telemetry/_instrumentation.py", line 561, in record_agent_invocation
    yield scope
  File ".../agents/base_agent.py", line 393, in _run
    yield event
GeneratorExit

During handling of the above exception, another exception occurred:
  File ".../opentelemetry/context/contextvars_context.py", line 42, in detach
    self._current_context.reset(token)
ValueError: <Token var=<ContextVar name='current_context' ...>> was created in a different Context

So an early-terminated agent node leaves its OTel span un-detached and runs its after_agent_callback and exit-stack teardown out of band. Early termination is ordinary: a workflow finishing, a node being interrupted, or an error upstream.

Steps to Reproduce:

On current main (665ec9835), consume one event from _run_impl and close it:

agen = agent._run_impl(ctx=Context(parent_ctx), node_input=None)
await agen.__anext__()
await agen.aclose()
assert cleaned_up == [True]   # the agent's own `finally` marker

Observed Behavior:

assert [] == [True]

plus the Failed to detach context traceback above on stderr. The agent's finally runs only later, from the finalizer hook.

Expected Behavior:

aclose() returns only after the inner generator has been closed, so cleanup ordering is deterministic and the OTel span is detached in the context that attached it.

Solution:

Wrap the inner generator in the Aclosing helper already imported in this module, matching BaseNode.run and LlmAgent._run_impl.

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Added test_run_impl_closes_run_async_when_consumer_stops_early, which asserts the inner generator's finally has run by the time aclose() returns. Verified it fails on an unmodified base_agent.py with assert [] == [True] and passes with this change.

$ pytest tests/unittests/agents tests/unittests/workflow -q
1729 passed, 1 skipped, 7 xfailed

$ pytest tests/unittests -q
15226 passed, 86 skipped, 27 xfailed, 2 xpassed in 358.66s (0:05:58)
# 1 unrelated failure, test_import_loading.py::...[agent], is an artifact of the
# git-worktree checkout I built this on; it fails there on an unmodified main
# and passes in a normal clone.

$ pre-commit run --files src/google/adk/agents/base_agent.py tests/unittests/agents/test_base_agent.py
# all hooks pass

Manual End-to-End (E2E) Tests:

Driven through _run_impl directly, since this is the node-mode adapter. Consuming one event from a three-event agent and closing:

before (main):   1. consumer stops early
                 2. outer _run_impl generator is closed      <- aclose() returns
                 3. agent impl cleaned up                    <- deferred, plus the OTel traceback

after (this PR): 1. consumer stops early
                 2. agent impl cleaned up
                 3. outer _run_impl generator is closed      <- aclose() returns

The ordering flips and the Failed to detach context traceback disappears.

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules. (none)

Additional context

No public API change, and the fully-consumed path is unchanged.

Two sibling adapters have the same bare async for over run_async and are deliberately left out of this PR to keep it to one file: agents/_managed_agent.py and labs/antigravity/_antigravity_agent.py, whose comment already says "Keep in sync with BaseAgent._run_impl". Happy to follow up on both, or fold them in here if you would rather have one change.

🤖 Generated with Claude Code

[AtulJoshi1206]

Rebased onto current main to clear the conflict. The conflict came from an unrelated runtime-config.json change that slipped into the original commit; that's dropped, so the diff is now only base_agent.py and its test. The fix merges cleanly with the callback refactor in c7fc421.

tests/unittests/agents: 855 passed, 2 xfailed. The new test still fails on unpatched main, so the bug is still present there.

@DeanChensj ready for review whenever you have a moment.

[AtulJoshi1206]

Rebased on latest main and pushed one more commit. After the telemetry change in 10c694f, the OTel detach error from the description is gone, but the early-close cleanup is still deferred (the test still fails on main). I also applied the same fix to ManagedAgent and AntigravityAgent, since their _run_impl mirrors BaseAgent._run_impl. Details are in the updated description.

@DeanChensj whenever you get a chance, thanks!