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

[AtulJoshi1206]

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