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:
committed by
Copybara-Service
parent
8d2f2779e6
commit
bb86bdd737
@@ -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
@@ -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:
|
||||
|
||||
@@ -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__])
|
||||
|
||||
Reference in New Issue
Block a user