fix: prevent contextvars leak across async generators

Merge https://github.com/google/adk-python/pull/5725

Fixes #5722

PiperOrigin-RevId: 967359464
This commit is contained in:
Chris Kinzel
2026-08-19 12:52:16 -07:00
committed by Copybara-Service
parent 8d2f2779e6
commit bb86bdd737
3 changed files with 544 additions and 193 deletions
+60 -30
View File
@@ -32,6 +32,7 @@ from typing import TypeVar
from typing import Union
from google.genai import types
from opentelemetry import context
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
@@ -87,6 +88,21 @@ class BaseAgentState(BaseModel):
AgentState = TypeVar('AgentState', bound=BaseAgentState)
_T = TypeVar('_T')
async def _with_caller_context(
agen: AsyncGenerator[_T, None],
caller_ctx: context.Context,
) -> AsyncGenerator[_T, None]:
"""Wraps an async generator to attach caller_ctx around each yield."""
async with Aclosing(agen) as a:
async for item in a:
token = context.attach(caller_ctx)
try:
yield item
finally:
context.detach(token)
# TODO: drop the explicit abc.ABC base once BaseNode surfaces ABCMeta to
@@ -295,26 +311,33 @@ class BaseAgent(BaseNode, abc.ABC):
Event: the events generated by the agent.
"""
ctx = self._create_invocation_context(parent_context)
async with _instrumentation.record_agent_invocation(ctx, self):
try:
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
caller_ctx = context.get_current()
async with Aclosing(self._run_async_impl(ctx)) as agen:
async for event in agen:
async def _run() -> AsyncGenerator[Event, None]:
ctx = self._create_invocation_context(parent_context)
async with _instrumentation.record_agent_invocation(ctx, self):
try:
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
if ctx.end_invocation:
return
async with Aclosing(self._run_async_impl(ctx)) as agen:
async for event in agen:
yield event
if event := await self._handle_after_agent_callback(ctx):
yield event
except Exception as e:
await self._handle_agent_error_callback(ctx, e)
raise
if ctx.end_invocation:
return
if event := await self._handle_after_agent_callback(ctx):
yield event
except Exception as e:
await self._handle_agent_error_callback(ctx, e)
raise
async with Aclosing(_with_caller_context(_run(), caller_ctx)) as agen:
async for event in agen:
yield event
@override
async def _run_impl(
@@ -350,23 +373,30 @@ class BaseAgent(BaseNode, abc.ABC):
Event: the events generated by the agent.
"""
ctx = self._create_invocation_context(parent_context)
async with _instrumentation.record_agent_invocation(ctx, self):
try:
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
caller_ctx = context.get_current()
async with Aclosing(self._run_live_impl(ctx)) as agen:
async for event in agen:
async def _run() -> AsyncGenerator[Event, None]:
ctx = self._create_invocation_context(parent_context)
async with _instrumentation.record_agent_invocation(ctx, self):
try:
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
if event := await self._handle_after_agent_callback(ctx):
yield event
except Exception as e:
await self._handle_agent_error_callback(ctx, e)
raise
async with Aclosing(self._run_live_impl(ctx)) as agen:
async for event in agen:
yield event
if event := await self._handle_after_agent_callback(ctx):
yield event
except Exception as e:
await self._handle_agent_error_callback(ctx, e)
raise
async with Aclosing(_with_caller_context(_run(), caller_ctx)) as agen:
async for event in agen:
yield event
async def _run_async_impl(
self, ctx: InvocationContext
+182 -163
View File
@@ -33,8 +33,10 @@ from typing import TYPE_CHECKING
import warnings
from google.genai import types
from opentelemetry import context
from typing_extensions import Self
from .agents.base_agent import _with_caller_context
from .agents.base_agent import BaseAgent
from .agents.context_cache_config import ContextCacheConfig
from .agents.invocation_context import InvocationContext
@@ -594,173 +596,181 @@ class Runner:
Events flow through ic._event_queue via NodeRunner.
"""
with _instrumentation.record_invocation(
entrypoint_node=node or self.agent, conversation_id=session_id
):
# 1. Setup
if session is None:
session = await self._get_or_create_session(
user_id=user_id,
session_id=session_id,
get_session_config=(run_config or RunConfig()).get_session_config,
caller_ctx = context.get_current()
async def _run() -> AsyncGenerator[Event, None]:
nonlocal invocation_id, new_message, session
with _instrumentation.record_invocation(
entrypoint_node=node or self.agent, conversation_id=session_id
):
# 1. Setup
if session is None:
session = await self._get_or_create_session(
user_id=user_id,
session_id=session_id,
get_session_config=(run_config or RunConfig()).get_session_config,
)
# Validate and resolve resume inputs
resume_inputs = self._extract_resume_inputs(new_message)
self._validate_new_message(new_message, resume_inputs)
if not invocation_id and new_message:
invocation_id = self._resolve_invocation_id_from_fr(
session, new_message
)
if not invocation_id:
active_scope = _find_active_task_scope(session)
if active_scope:
_, inv_id = active_scope
invocation_id = inv_id
ic = self._new_invocation_context(
session,
new_message=new_message,
run_config=run_config or RunConfig(),
invocation_id=invocation_id,
)
ic._event_queue = asyncio.Queue()
# Validate and resolve resume inputs
resume_inputs = self._extract_resume_inputs(new_message)
self._validate_new_message(new_message, resume_inputs)
# 2. Append user message to session and resolve node_input
node_input = None
if resume_inputs or invocation_id:
# Resume: recover the original user content. new_message here is a
# function response (or None), so it can't populate user_content.
node_input = self._find_original_user_content(
ic.session, ic.invocation_id
)
if node_input:
ic.user_content = node_input
if not node_input:
# Fresh: use user message as node_input
node_input = new_message
if not invocation_id and new_message:
invocation_id = self._resolve_invocation_id_from_fr(
session, new_message
# Failures in the setup hooks below (on_user_message_callback, the
# user-event session append, and before_run_callback) must also notify
# on_run_error_callback: they are part of runner execution even though
# they run before the main event loop. Notification-only; the original
# exception is always re-raised, and after_run stays success-only.
try:
# Run callbacks on user message
if new_message:
modified_user_message = (
await ic.plugin_manager.run_on_user_message_callback(
invocation_context=ic, user_message=new_message
)
)
if modified_user_message is not None:
new_message = modified_user_message
ic.user_content = new_message
# Append user message to session for history
if new_message:
user_event = await self._append_user_event(
ic, new_message, state_delta=state_delta
)
if yield_user_message and user_event:
yield user_event
# Run before_run callbacks
await ic.plugin_manager.run_before_run_callback(invocation_context=ic)
except Exception as e:
await _notify_run_error(ic.plugin_manager, ic, e)
raise
# 3. Start root node in background
from .agents.context import Context
from .workflow._dynamic_node_scheduler import DynamicNodeScheduler
from .workflow._errors import DynamicNodeFailError
from .workflow._errors import NodeInterruptedError
from .workflow._workflow import _LoopState
root_ctx = Context(ic)
root_node = node or self.agent
is_agent = isinstance(self.agent, BaseAgent)
has_sub_agents = is_agent and bool(
getattr(self.agent, 'sub_agents', None)
)
if not invocation_id:
active_scope = _find_active_task_scope(session)
if active_scope:
_, inv_id = active_scope
invocation_id = inv_id
use_scheduler = is_agent and has_sub_agents
ic = self._new_invocation_context(
session,
new_message=new_message,
run_config=run_config or RunConfig(),
invocation_id=invocation_id,
)
ic._event_queue = asyncio.Queue()
# The root chat coordinator's isolation_scope stays None: its own
# events (FCs, text, synthesized FRs from completed task
# delegations) are also unscoped, so the content-builder's
# isolation_scope filter lets the coordinator see all of them
# across user turns. Task sub-agents are scoped under their
# originating function-call id and so remain invisible to the
# coordinator's view.
# 2. Append user message to session and resolve node_input
node_input = None
if resume_inputs or invocation_id:
# Resume: recover the original user content. new_message here is a
# function response (or None), so it can't populate user_content.
node_input = self._find_original_user_content(
ic.session, ic.invocation_id
)
if node_input:
ic.user_content = node_input
if not node_input:
# Fresh: use user message as node_input
node_input = new_message
done_sentinel = object()
# Failures in the setup hooks below (on_user_message_callback, the
# user-event session append, and before_run_callback) must also notify
# on_run_error_callback: they are part of runner execution even though
# they run before the main event loop. Notification-only; the original
# exception is always re-raised, and after_run stays success-only.
try:
# Run callbacks on user message
if new_message:
modified_user_message = (
await ic.plugin_manager.run_on_user_message_callback(
invocation_context=ic, user_message=new_message
async def _drive_root_node() -> None:
try:
if use_scheduler:
# Rehydration warning: DynamicNodeScheduler relies on session.events scanning.
# Stateful live EUC/LRO streams may rehydrate freshly if not yet persisted.
scheduler = DynamicNodeScheduler(state=_LoopState())
root_ctx._workflow_scheduler = scheduler
try:
await root_ctx._run_node_internal(
root_node,
node_input=node_input,
resume_inputs=resume_inputs,
)
)
if modified_user_message is not None:
new_message = modified_user_message
ic.user_content = new_message
except NodeInterruptedError:
# The node was interrupted (e.g. for HITL).
pass
except DynamicNodeFailError as e:
raise e.error
finally:
await ic._event_queue.put((done_sentinel, None))
# Append user message to session for history
if new_message:
user_event = await self._append_user_event(
ic, new_message, state_delta=state_delta
)
if yield_user_message and user_event:
yield user_event
task = asyncio.create_task(_drive_root_node())
# Run before_run callbacks
await ic.plugin_manager.run_before_run_callback(invocation_context=ic)
except Exception as e:
await _notify_run_error(ic.plugin_manager, ic, e)
raise
# 3. Start root node in background
from .agents.context import Context
from .workflow._dynamic_node_scheduler import DynamicNodeScheduler
from .workflow._errors import DynamicNodeFailError
from .workflow._errors import NodeInterruptedError
from .workflow._workflow import _LoopState
root_ctx = Context(ic)
root_node = node or self.agent
is_agent = isinstance(self.agent, BaseAgent)
has_sub_agents = is_agent and bool(
getattr(self.agent, 'sub_agents', None)
)
use_scheduler = is_agent and has_sub_agents
# The root chat coordinator's isolation_scope stays None: its own
# events (FCs, text, synthesized FRs from completed task
# delegations) are also unscoped, so the content-builder's
# isolation_scope filter lets the coordinator see all of them
# across user turns. Task sub-agents are scoped under their
# originating function-call id and so remain invisible to the
# coordinator's view.
done_sentinel = object()
async def _drive_root_node() -> None:
# 4. Main loop: consume events, persist, yield
run_error = None
try:
if use_scheduler:
# Rehydration warning: DynamicNodeScheduler relies on session.events scanning.
# Stateful live EUC/LRO streams may rehydrate freshly if not yet persisted.
scheduler = DynamicNodeScheduler(state=_LoopState())
root_ctx._workflow_scheduler = scheduler
try:
await root_ctx._run_node_internal(
root_node,
node_input=node_input,
resume_inputs=resume_inputs,
)
except NodeInterruptedError:
# The node was interrupted (e.g. for HITL).
pass
except DynamicNodeFailError as e:
raise e.error
async with aclosing(
self._consume_event_queue(ic, done_sentinel)
) as agen:
async for event in agen:
yield event
finally:
# _cleanup_root_task re-raises a root-node Exception (if any) after
# the event stream has drained.
await self._cleanup_root_task(task, self.agent.name)
except Exception as e:
# An unhandled exception escaped runner execution. Notify plugins
# (notification-only) and re-raise. after_run stays success-only.
run_error = e
await _notify_run_error(ic.plugin_manager, ic, e)
raise
finally:
await ic._event_queue.put((done_sentinel, None))
# Success path (also caller early-stop via GeneratorExit, which is not
# an Exception): run after_run and compaction. _cleanup_root_task has
# already run in the inner finally above. A failure in this success
# cleanup (e.g. an after_run plugin raising, which PluginManager
# surfaces as a RuntimeError) is itself an unhandled runner error, so
# notify on_run_error_callback once and re-raise. on_run_error is
# notification-only and never raises, so there is no recursive
# notification.
if run_error is None:
try:
await ic.plugin_manager.run_after_run_callback(
invocation_context=ic
)
await self._run_post_invocation_compaction(
session=session,
skip_token_compaction=ic.token_compaction_checked,
)
except Exception as e:
await _notify_run_error(ic.plugin_manager, ic, e)
raise
task = asyncio.create_task(_drive_root_node())
# 4. Main loop: consume events, persist, yield
run_error = None
try:
try:
async with aclosing(
self._consume_event_queue(ic, done_sentinel)
) as agen:
async for event in agen:
yield event
finally:
# _cleanup_root_task re-raises a root-node Exception (if any) after
# the event stream has drained.
await self._cleanup_root_task(task, self.agent.name)
except Exception as e:
# An unhandled exception escaped runner execution. Notify plugins
# (notification-only) and re-raise. after_run stays success-only.
run_error = e
await _notify_run_error(ic.plugin_manager, ic, e)
raise
finally:
# Success path (also caller early-stop via GeneratorExit, which is not
# an Exception): run after_run and compaction. _cleanup_root_task has
# already run in the inner finally above. A failure in this success
# cleanup (e.g. an after_run plugin raising, which PluginManager
# surfaces as a RuntimeError) is itself an unhandled runner error, so
# notify on_run_error_callback once and re-raise. on_run_error is
# notification-only and never raises, so there is no recursive
# notification.
if run_error is None:
try:
await ic.plugin_manager.run_after_run_callback(
invocation_context=ic
)
await self._run_post_invocation_compaction(
session=session,
skip_token_compaction=ic.token_compaction_checked,
)
except Exception as e:
await _notify_run_error(ic.plugin_manager, ic, e)
raise
async with aclosing(_with_caller_context(_run(), caller_ctx)) as agen:
async for event in agen:
yield event
async def _run_node_live(
self,
@@ -1336,6 +1346,7 @@ class Runner:
new_message: Optional[types.Content] = None,
invocation_id: Optional[str] = None,
) -> AsyncGenerator[Event, None]:
caller_ctx_trace = context.get_current()
with _instrumentation.record_invocation(
entrypoint_node=root_agent, conversation_id=session_id
):
@@ -1417,11 +1428,14 @@ class Runner:
yield event
async with aclosing(
self._exec_with_plugin(
invocation_context=invocation_context,
session=invocation_context.session,
execute_fn=execute,
is_live_call=False,
_with_caller_context(
self._exec_with_plugin(
invocation_context=invocation_context,
session=invocation_context.session,
execute_fn=execute,
is_live_call=False,
),
caller_ctx_trace,
)
) as agen:
async for event in agen:
@@ -1884,6 +1898,8 @@ class Runner:
if run_config.response_modalities is None:
run_config = run_config.model_copy()
run_config.response_modalities = [types.Modality.AUDIO]
caller_ctx = context.get_current()
if session is None and (user_id is None or session_id is None):
raise ValueError(
'Either session or user_id and session_id must be provided.'
@@ -1944,11 +1960,14 @@ class Runner:
yield event
async with aclosing(
self._exec_with_plugin(
invocation_context=invocation_context,
session=invocation_context.session,
execute_fn=execute,
is_live_call=True,
_with_caller_context(
self._exec_with_plugin(
invocation_context=invocation_context,
session=invocation_context.session,
execute_fn=execute,
is_live_call=True,
),
caller_ctx,
)
) as agen:
async for event in agen:
+302
View File
@@ -2681,5 +2681,307 @@ async def test_runner_picks_coordinator_when_has_remote_a2a_task_subagent():
assert called_node == coordinator
@pytest.mark.asyncio
async def test_run_async_does_not_leak_context_base_node():
"""Caller OpenTelemetry context is preserved during run_async iteration for BaseNode."""
from typing import Any
from google.adk.agents.context import Context
from google.adk.workflow._base_node import BaseNode
from opentelemetry import context as otel_context
class _TestEchoNode(BaseNode):
async def _run_impl(
self, *, ctx: Context, node_input: Any
) -> AsyncGenerator[Any, None]:
yield "echo"
session_service = InMemorySessionService()
runner = Runner(
app_name=TEST_APP_ID,
node=_TestEchoNode(name="test_node"),
session_service=session_service,
artifact_service=InMemoryArtifactService(),
auto_create_session=True,
)
test_key = otel_context.create_key("test_key_node")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_node")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(
runner.run_async(
user_id=TEST_USER_ID,
session_id=TEST_SESSION_ID,
new_message=types.Content(
role="user", parts=[types.Part(text="hello")]
),
)
) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert otel_context.get_value(test_key) == "caller_val_node"
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
@pytest.mark.asyncio
async def test_run_async_does_not_leak_context_base_agent():
"""Caller OpenTelemetry context is preserved during run_async iteration for BaseAgent."""
from opentelemetry import context as otel_context
session_service = InMemorySessionService()
runner = Runner(
app_name=TEST_APP_ID,
agent=MockAgent("test_agent"),
session_service=session_service,
artifact_service=InMemoryArtifactService(),
auto_create_session=True,
)
test_key = otel_context.create_key("test_key_agent")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_agent")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(
runner.run_async(
user_id=TEST_USER_ID,
session_id=TEST_SESSION_ID,
new_message=types.Content(
role="user", parts=[types.Part(text="hello")]
),
)
) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert otel_context.get_value(test_key) == "caller_val_agent"
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
@pytest.mark.asyncio
async def test_run_node_async_does_not_leak_context():
"""Caller OpenTelemetry context is preserved during _run_node_async iteration."""
from typing import Any
from google.adk.agents.context import Context
from google.adk.workflow._base_node import BaseNode
from opentelemetry import context as otel_context
class _TestEchoNode(BaseNode):
async def _run_impl(
self, *, ctx: Context, node_input: Any
) -> AsyncGenerator[Any, None]:
yield "echo"
session_service = InMemorySessionService()
runner = Runner(
app_name=TEST_APP_ID,
node=_TestEchoNode(name="test_node"),
session_service=session_service,
artifact_service=InMemoryArtifactService(),
auto_create_session=True,
)
test_key = otel_context.create_key("test_key_run_node_async")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_run_node_async")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(
runner._run_node_async(
user_id=TEST_USER_ID,
session_id=TEST_SESSION_ID,
new_message=types.Content(
role="user", parts=[types.Part(text="hello")]
),
yield_user_message=True,
)
) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert otel_context.get_value(test_key) == "caller_val_run_node_async"
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
@pytest.mark.asyncio
async def test_run_async_does_not_leak_context_llm_agent():
"""Caller OpenTelemetry context is preserved during run_async iteration for LlmAgent."""
from opentelemetry import context as otel_context
session_service = InMemorySessionService()
runner = Runner(
app_name=TEST_APP_ID,
agent=MockLlmAgent("test_llm_agent"),
session_service=session_service,
artifact_service=InMemoryArtifactService(),
auto_create_session=True,
)
test_key = otel_context.create_key("test_key_llm")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_llm")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(
runner.run_async(
user_id=TEST_USER_ID,
session_id=TEST_SESSION_ID,
new_message=types.Content(
role="user", parts=[types.Part(text="hello")]
),
)
) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert otel_context.get_value(test_key) == "caller_val_llm"
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
@pytest.mark.asyncio
async def test_run_live_does_not_leak_context():
"""Caller OpenTelemetry context is preserved during run_live iteration."""
from google.adk.agents.live_request_queue import LiveRequestQueue
from opentelemetry import context as otel_context
session_service = InMemorySessionService()
runner = Runner(
app_name=TEST_APP_ID,
agent=MockLiveAgent("test_live_agent"),
session_service=session_service,
artifact_service=InMemoryArtifactService(),
auto_create_session=True,
)
live_queue = LiveRequestQueue()
test_key = otel_context.create_key("test_key_live")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_live")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(
runner.run_live(
user_id=TEST_USER_ID,
session_id=TEST_SESSION_ID,
live_request_queue=live_queue,
)
) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert otel_context.get_value(test_key) == "caller_val_live"
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
@pytest.mark.asyncio
async def test_base_agent_run_async_does_not_leak_context():
"""Caller OpenTelemetry context is preserved during BaseAgent.run_async iteration."""
from google.adk.plugins.plugin_manager import PluginManager
from opentelemetry import context as otel_context
agent = MockAgent("test_agent")
session_service = InMemorySessionService()
session = await session_service.create_session(
app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
)
inv_ctx = InvocationContext(
session=session,
session_service=session_service,
plugin_manager=PluginManager(),
agent=agent,
invocation_id="inv_test",
)
test_key = otel_context.create_key("test_key_base_agent_run_async")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_base_agent_run_async")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(agent.run_async(parent_context=inv_ctx)) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert (
otel_context.get_value(test_key)
== "caller_val_base_agent_run_async"
)
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
@pytest.mark.asyncio
async def test_base_agent_run_live_does_not_leak_context():
"""Caller OpenTelemetry context is preserved during BaseAgent.run_live iteration."""
from google.adk.plugins.plugin_manager import PluginManager
from opentelemetry import context as otel_context
agent = MockLiveAgent("test_live_agent")
session_service = InMemorySessionService()
session = await session_service.create_session(
app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
)
inv_ctx = InvocationContext(
session=session,
session_service=session_service,
plugin_manager=PluginManager(),
agent=agent,
invocation_id="inv_test_live",
)
test_key = otel_context.create_key("test_key_base_agent_run_live")
token = otel_context.attach(
otel_context.set_value(test_key, "caller_val_base_agent_run_live")
)
caller_ctx = otel_context.get_current()
try:
events = []
async with aclosing(agent.run_live(parent_context=inv_ctx)) as agen:
async for event in agen:
assert otel_context.get_current() == caller_ctx
assert (
otel_context.get_value(test_key) == "caller_val_base_agent_run_live"
)
events.append(event)
assert events
assert otel_context.get_current() == caller_ctx
finally:
otel_context.detach(token)
if __name__ == "__main__":
pytest.main([__file__])