From 044cd7b5bae9c949dfd0c6ce7fa6e253084fa4c1 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Tue, 30 Jun 2026 15:17:53 -0700 Subject: [PATCH 1/6] feat(telemetry): add gen_ai.invoke_agent.{inference,tool}_calls metrics Adds the two per-invocation call-count metrics from semconv #336 (open-telemetry/semantic-conventions-genai#336): * gen_ai.invoke_agent.inference_calls: model calls in one invoke_agent span. * gen_ai.invoke_agent.tool_calls: tool calls in one invoke_agent span. Both are counted on the per-span TelemetryContext at the call sites (record_inference_telemetry, record_tool_execution) and flushed in record_agent_invocation, keyed by gen_ai.agent.name. One point per invoke_agent span, so a multi-agent invocation does not double count. Failed calls still count; only client-side tool calls count. Always emitted (a no-op without a MeterProvider). PiperOrigin-RevId: 940684383 --- src/google/adk/agents/invocation_context.py | 10 ++++ src/google/adk/telemetry/_instrumentation.py | 55 ++++++++++++++++++- src/google/adk/telemetry/_metrics.py | 50 +++++++++++++++++ .../telemetry/functional_node_test_cases.py | 12 ++++ .../telemetry/functional_test_cases.py | 12 ++++ .../telemetry/functional_test_helpers.py | 10 ++++ .../telemetry/test_instrumentation.py | 5 +- 7 files changed, 149 insertions(+), 5 deletions(-) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 6afa0d3573e..5dac20f6231 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -18,6 +18,7 @@ from typing import Any from typing import cast from typing import Optional +from typing import TYPE_CHECKING from google.adk.platform import uuid as platform_uuid from google.genai import types @@ -47,6 +48,9 @@ from .run_config import RunConfig from .transcription_entry import TranscriptionEntry +if TYPE_CHECKING: + from google.adk.telemetry._instrumentation import TelemetryContext + class LlmCallsLimitExceededError(Exception): """Error thrown when the number of LLM calls exceed the limit.""" @@ -270,6 +274,12 @@ class InvocationContext(BaseModel): of this invocation. """ + _invoke_agent_telemetry_context: Optional[TelemetryContext] = PrivateAttr( + default=None + ) + """TelemetryContext of the active ``invoke_agent`` span, if any. + """ + @property def is_resumable(self) -> bool: """Returns whether the current invocation is resumable.""" diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 2284fb839dc..e12b7c9b054 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -27,11 +27,11 @@ from . import _metrics from . import tracing -from ..events import event as event_lib if TYPE_CHECKING: from ..agents.base_agent import BaseAgent from ..agents.invocation_context import InvocationContext + from ..events import event as event_lib from ..models.llm_request import LlmRequest from ..models.llm_response import LlmResponse from ..tools.base_tool import BaseTool @@ -78,11 +78,31 @@ class TelemetryContext: error_type: str | None = None span: tracing.GenerateContentSpan | trace.Span | None = None _llm_responses: list[LlmResponse] = dataclasses.field(default_factory=list) + _inference_call_count: int = 0 + _tool_call_count: int = 0 @property def llm_responses(self) -> list[LlmResponse]: return self._llm_responses + @property + def inference_call_count(self) -> int: + """Number of model calls counted against this invoke_agent span.""" + return self._inference_call_count + + def increment_inference_calls(self) -> None: + """Counts one model call against this span (including calls that raised).""" + self._inference_call_count += 1 + + def increment_tool_calls(self) -> None: + """Counts one tool call against this invoke_agent span.""" + self._tool_call_count += 1 + + @property + def tool_call_count(self) -> int: + """Number of tool calls counted against this invoke_agent span.""" + return self._tool_call_count + def record_llm_response( self, invocation_context: InvocationContext, response: LlmResponse ) -> None: @@ -110,6 +130,30 @@ def _record_agent_metrics( logger.exception("Failed to record agent metrics for agent %s", agent_name) +def _flush_invoke_agent_metrics( + tel_ctx: TelemetryContext, agent_name: str +) -> None: + """Flushes this span's accumulated inference/tool-call metrics.""" + _metrics.record_invoke_agent_inference_calls( + agent_name, tel_ctx.inference_call_count + ) + _metrics.record_invoke_agent_tool_calls(agent_name, tel_ctx.tool_call_count) + + +def _accumulate_invoke_agent_tool_call(ctx: InvocationContext) -> None: + """Counts one tool call against the active invoke_agent span.""" + span_tel_ctx = ctx._invoke_agent_telemetry_context # pylint: disable=protected-access + if span_tel_ctx is not None: + span_tel_ctx.increment_tool_calls() + + +def _accumulate_invoke_agent_inference_call(ctx: InvocationContext) -> None: + """Counts one model call against the active invoke_agent span.""" + span_tel_ctx = ctx._invoke_agent_telemetry_context # pylint: disable=protected-access + if span_tel_ctx is not None: + span_tel_ctx.increment_inference_calls() + + @contextlib.asynccontextmanager async def record_agent_invocation( ctx: InvocationContext, agent: BaseAgent @@ -119,11 +163,13 @@ async def record_agent_invocation( caught_error: Exception | None = None span: trace.Span | None = None span_name = f"invoke_agent {agent.name}" + tel_ctx = TelemetryContext() try: with tracing.tracer.start_as_current_span(span_name) as s: span = s tracing.trace_agent_invocation(span, agent, ctx) - tel_ctx = TelemetryContext(otel_context=context_api.get_current()) + tel_ctx.otel_context = context_api.get_current() + ctx._invoke_agent_telemetry_context = tel_ctx # pylint: disable=protected-access yield tel_ctx except Exception as e: caught_error = e @@ -136,6 +182,7 @@ async def record_agent_invocation( getattr(getattr(ctx, "session", None), "events", []), caught_error, ) + _flush_invoke_agent_metrics(tel_ctx, agent.name) @contextlib.asynccontextmanager @@ -143,7 +190,7 @@ async def record_tool_execution( tool: BaseTool, agent: BaseAgent, function_args: dict[str, object], - invocation_context: InvocationContext | None = None, + invocation_context: InvocationContext, ) -> AsyncIterator[TelemetryContext]: """Unified context manager for consolidated tool execution telemetry.""" start_time = time.monotonic() @@ -171,6 +218,7 @@ async def record_tool_execution( invocation_context=invocation_context, error_type=tel_ctx.error_type, ) + _accumulate_invoke_agent_tool_call(invocation_context) finally: try: _metrics.record_tool_execution_duration( @@ -205,6 +253,7 @@ async def record_inference_telemetry( yield tel_ctx finally: inference_error = sys.exc_info()[1] + _accumulate_invoke_agent_inference_call(invocation_context) agent = invocation_context.agent elapsed_s = _get_elapsed_s(tel_ctx.span, start_time) try: diff --git a/src/google/adk/telemetry/_metrics.py b/src/google/adk/telemetry/_metrics.py index eac65b48ca0..7be66bc7a8c 100644 --- a/src/google/adk/telemetry/_metrics.py +++ b/src/google/adk/telemetry/_metrics.py @@ -146,6 +146,44 @@ gen_ai_metrics.create_gen_ai_client_operation_duration(meter) ) _client_token_usage = gen_ai_metrics.create_gen_ai_client_token_usage(meter) +_invoke_agent_inference_calls = meter.create_histogram( + "gen_ai.invoke_agent.inference_calls", + unit="1", + description="Number of inference (model) calls per agent invocation.", + explicit_bucket_boundaries_advisory=[ + 1, + 2, + 3, + 4, + 5, + 6, + 8, + 12, + 16, + 24, + 32, + 64, + ], +) +_invoke_agent_tool_calls = meter.create_histogram( + "gen_ai.invoke_agent.tool_calls", + unit="1", + description="Number of tool calls per agent invocation.", + explicit_bucket_boundaries_advisory=[ + 1, + 2, + 3, + 4, + 5, + 6, + 8, + 12, + 16, + 24, + 32, + 64, + ], +) def record_agent_invocation_duration( @@ -160,6 +198,18 @@ def record_agent_invocation_duration( _agent_invocation_duration.record(elapsed_s, attributes=attrs) +def record_invoke_agent_inference_calls(agent_name: str, count: int): + """Records the number of inference (model) calls in an agent invocation.""" + attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name} + _invoke_agent_inference_calls.record(count, attributes=attrs) + + +def record_invoke_agent_tool_calls(agent_name: str, count: int): + """Records the number of tool calls in an agent invocation.""" + attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name} + _invoke_agent_tool_calls.record(count, attributes=attrs) + + def record_agent_request_size( agent_name: str, user_content: types.Content | None ): diff --git a/tests/unittests/telemetry/functional_node_test_cases.py b/tests/unittests/telemetry/functional_node_test_cases.py index 56e58a7fa7a..67774235786 100644 --- a/tests/unittests/telemetry/functional_node_test_cases.py +++ b/tests/unittests/telemetry/functional_node_test_cases.py @@ -2900,6 +2900,12 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_agent.inference_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), + }), + "gen_ai.invoke_agent.tool_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), + }), } @@ -2947,6 +2953,12 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_agent.inference_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), + }), + "gen_ai.invoke_agent.tool_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), + }), } diff --git a/tests/unittests/telemetry/functional_test_cases.py b/tests/unittests/telemetry/functional_test_cases.py index 44c352a4c88..0bb7b9b3695 100644 --- a/tests/unittests/telemetry/functional_test_cases.py +++ b/tests/unittests/telemetry/functional_test_cases.py @@ -2294,6 +2294,12 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_agent.inference_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), + }), + "gen_ai.invoke_agent.tool_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), + }), } @@ -2341,6 +2347,12 @@ value=NON_DETERMINISTIC, ), }), + "gen_ai.invoke_agent.inference_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), + }), + "gen_ai.invoke_agent.tool_calls": frozenset({ + MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), + }), } diff --git a/tests/unittests/telemetry/functional_test_helpers.py b/tests/unittests/telemetry/functional_test_helpers.py index e4ae30a3607..82f997b6979 100644 --- a/tests/unittests/telemetry/functional_test_helpers.py +++ b/tests/unittests/telemetry/functional_test_helpers.py @@ -339,6 +339,16 @@ class HistogramSpec(NamedTuple): attr="_client_token_usage", metric_name="gen_ai.client.token.usage", ), + HistogramSpec( + module=_metrics, + attr="_invoke_agent_inference_calls", + metric_name="gen_ai.invoke_agent.inference_calls", + ), + HistogramSpec( + module=_metrics, + attr="_invoke_agent_tool_calls", + metric_name="gen_ai.invoke_agent.tool_calls", + ), ) diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index 8711aa979ef..92fc33b050e 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/tests/unittests/telemetry/test_instrumentation.py @@ -15,6 +15,7 @@ # pylint: disable=protected-access import time +import types from unittest import mock from google.adk.telemetry import _instrumentation @@ -94,8 +95,8 @@ async def test_record_agent_invocation_tolerates_minimal_context(): """ agent = mock.MagicMock() agent.name = "test_agent" - # Bare object without `user_content` and without `session`. - bare_ctx = object() + # Context-like without `user_content` and without `session`. + bare_ctx = types.SimpleNamespace() with ( mock.patch.object( From 7329f7d302b85e001e643e1a46c3d057aca40ed8 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Tue, 30 Jun 2026 15:23:13 -0700 Subject: [PATCH 2/6] chore: Extract transfer loop out of dynamic scheduler This moves the transfer loop from _dynamic_node_scheduler to Context.run_node, making run_node the unified entrypoint for transfers in both static workflows and dynamic agent calls. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 940687129 --- .../adk/workflow/_dynamic_node_scheduler.py | 187 +++----- src/google/adk/workflow/_llm_agent_wrapper.py | 30 +- .../workflow/test_dynamic_node_scheduler.py | 450 ------------------ 3 files changed, 80 insertions(+), 587 deletions(-) diff --git a/src/google/adk/workflow/_dynamic_node_scheduler.py b/src/google/adk/workflow/_dynamic_node_scheduler.py index 110117c015f..6cdd45e6c14 100644 --- a/src/google/adk/workflow/_dynamic_node_scheduler.py +++ b/src/google/adk/workflow/_dynamic_node_scheduler.py @@ -153,138 +153,77 @@ async def __call__( Returns: Child Context with output, route, and interrupt_ids set. """ - curr_node = node - curr_name = node_name or node.name - curr_run_id = run_id - curr_input = node_input - curr_parent_ctx: Context | None = ctx - - while True: - curr_parent_path = curr_parent_ctx.node_path if curr_parent_ctx else None - base_path_builder = ( - _NodePathBuilder.from_string(curr_parent_path) - if curr_parent_path - else _NodePathBuilder([]) - ) - node_path = str(base_path_builder.append(curr_name, curr_run_id)) - - # Rehydration chronological sequence barrier setup for the parent path - parent_path = curr_parent_ctx.node_path if curr_parent_ctx else '' - if parent_path and parent_path not in self._parent_sequence_barriers: - seq = self._scan_parent_child_sequence(curr_parent_ctx, parent_path) - self._parent_sequence_barriers[parent_path] = ReplaySequenceBarrier(seq) - - # Runtime schema validation. - if curr_input is not None: - try: - curr_input = curr_node._validate_input_data(curr_input) - except ValidationError as e: - raise ValueError( - 'Runtime schema validation failed for dynamic node' - f" '{curr_name}'. Input does not match input_schema: {e}" - ) from e - - logger.debug('node %s schedule start.', node_path) - - # Phase 1: Lazy rehydration from session events. - if node_path not in self._state.runs: - self._rehydrate_from_events(curr_parent_ctx, node_path) + curr_parent_path = ctx.node_path if ctx else None + base_path_builder = ( + _NodePathBuilder.from_string(curr_parent_path) + if curr_parent_path + else _NodePathBuilder([]) + ) + node_path = str(base_path_builder.append(node_name or node.name, run_id)) + + # Rehydration chronological sequence barrier setup for the parent path + parent_path = ctx.node_path if ctx else '' + if parent_path and parent_path not in self._parent_sequence_barriers: + seq = self._scan_parent_child_sequence(ctx, parent_path) + self._parent_sequence_barriers[parent_path] = ReplaySequenceBarrier(seq) + + # Runtime schema validation. + if node_input is not None: + try: + node_input = node._validate_input_data(node_input) + except ValidationError as e: + raise ValueError( + 'Runtime schema validation failed for dynamic node' + f" '{node_name or node.name}'. Input does not match" + f' input_schema: {e}' + ) from e + + logger.debug('node %s schedule start.', node_path) + + # Phase 1: Lazy rehydration from session events. + if node_path not in self._state.runs: + self._rehydrate_from_events(ctx, node_path) + + # Check existing run and determine if fresh execution is needed. + child_ctx, run_completed = await self._check_existing_run( + ctx, + node, + node_name or node.name, + node_path, + run_id, + node_input, + use_as_output, + use_sub_branch, + override_branch, + override_isolation_scope=override_isolation_scope, + ) - # Check existing run and determine if fresh execution is needed. - child_ctx, run_completed = await self._check_existing_run( - curr_parent_ctx, - curr_node, - curr_name, + if not run_completed: + # Phase 3: Fresh execution. + logger.debug('node %s schedule: Fresh execution.', node_path) + child_ctx = await self._run_node_internal( + ctx, + node, + node_name or node.name, node_path, - curr_run_id, - curr_input, + run_id, + node_input, use_as_output, - use_sub_branch, - override_branch, + is_fresh=True, + use_sub_branch=use_sub_branch, + override_branch=override_branch, override_isolation_scope=override_isolation_scope, ) - if not run_completed: - # Phase 3: Fresh execution. - logger.debug('node %s schedule: Fresh execution.', node_path) - child_ctx = await self._run_node_internal( - curr_parent_ctx, - curr_node, - curr_name, - node_path, - curr_run_id, - curr_input, - use_as_output, - is_fresh=True, - use_sub_branch=use_sub_branch, - override_branch=override_branch, - override_isolation_scope=override_isolation_scope, - ) - - logger.debug('node %s schedule end.', node_path) - - # Advance chronological sequence for this parent path and key - parent_path = curr_parent_ctx.node_path if curr_parent_ctx else '' - key = f'{curr_name}@{curr_run_id}' - if parent_path in self._parent_sequence_barriers: - self._parent_sequence_barriers[parent_path].check_and_advance(key) - - # Check for transfer_to_agent signal. - transfer_to_agent = ( - child_ctx.actions.transfer_to_agent if child_ctx else None - ) - if isinstance(transfer_to_agent, str): - target_name = transfer_to_agent - root_agent = getattr(curr_node, 'root_agent', None) - if not root_agent: - raise ValueError(f'Cannot find root_agent on node {curr_node.name}') - - # Local import to avoid runtime circular dependencies with Context - from .utils._transfer_utils import resolve_and_derive_transfer_context - - target_agent, next_parent_ctx = resolve_and_derive_transfer_context( - target_name=target_name, - current_agent=curr_node, - root_agent=root_agent, - curr_ctx=child_ctx, - curr_parent_ctx=curr_parent_ctx, - ) - if not target_agent: - raise ValueError(f"Transfer target agent '{target_name}' not found.") - if not next_parent_ctx: - available = [] - if hasattr(curr_node, '_get_available_agent_names'): - available = curr_node._get_available_agent_names() - available_str = ( - f"\nAvailable agents: {', '.join(available)}" if available else '' - ) - raise ValueError( - f"Cannot transfer from '{curr_name}' to unrelated agent" - f" '{target_name}'.{available_str}" - ) - curr_parent_ctx = next_parent_ctx - - # Set up parameters for next iteration. - curr_node = target_agent - curr_name = target_agent.name - - if not curr_parent_ctx: - raise AssertionError( - 'curr_parent_ctx cannot be None during active workflow execution' - ) - - curr_parent_ctx._child_run_counters[target_agent.name] = ( - curr_parent_ctx._child_run_counters.get(target_agent.name, 0) + 1 - ) - curr_run_id = str( - curr_parent_ctx._child_run_counters[target_agent.name] - ) - curr_input = None # Input for transfer target is usually empty. + logger.debug('node %s schedule end.', node_path) - # Loop continues to execute the next agent - continue + # Advance chronological sequence for this parent path and key + parent_path = ctx.node_path if ctx else '' + key = f'{node_name or node.name}@{run_id}' + if parent_path in self._parent_sequence_barriers: + self._parent_sequence_barriers[parent_path].check_and_advance(key) - return child_ctx + return child_ctx async def _check_existing_run( self, diff --git a/src/google/adk/workflow/_llm_agent_wrapper.py b/src/google/adk/workflow/_llm_agent_wrapper.py index 62920db2912..a4f0dc900ed 100644 --- a/src/google/adk/workflow/_llm_agent_wrapper.py +++ b/src/google/adk/workflow/_llm_agent_wrapper.py @@ -394,19 +394,23 @@ async def run_llm_agent_as_node( break # close this run_iter; outer loop re-enters if event.actions.transfer_to_agent: target_name = event.actions.transfer_to_agent - if target_name != agent.name: - from ..agents.llm_agent import LlmAgent - - if ( - isinstance(agent, LlmAgent) - and ctx._invocation_context.is_resumable - ): - ctx._invocation_context.set_agent_state( - agent.name, end_of_agent=True - ) - yield agent._create_agent_state_event(ctx._invocation_context) - transferred = True - break + if target_name == agent.name: + raise ValueError( + f"Agent '{target_name}' cannot transfer to itself." + ) + + from ..agents.llm_agent import LlmAgent + + if ( + isinstance(agent, LlmAgent) + and ctx._invocation_context.is_resumable + ): + ctx._invocation_context.set_agent_state( + agent.name, end_of_agent=True + ) + yield agent._create_agent_state_event(ctx._invocation_context) + transferred = True + break if not had_task_fc or transferred: # LLM finished without delegating (or transferred away); # nothing more for this wrapper to do. diff --git a/tests/unittests/workflow/test_dynamic_node_scheduler.py b/tests/unittests/workflow/test_dynamic_node_scheduler.py index 6ef1b8f77f2..27244caeafb 100644 --- a/tests/unittests/workflow/test_dynamic_node_scheduler.py +++ b/tests/unittests/workflow/test_dynamic_node_scheduler.py @@ -698,456 +698,6 @@ async def test_runtime_schema_validation_content_fallback(): # Should not raise -# ========================================================================= -# __call__ — Agent Transfer logic -# ========================================================================= - - -@pytest.mark.asyncio -async def test_scheduler_handles_child_transfer(): - """Scheduler processes CHILD relationship by nesting next context.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - target = LlmAgent(name='target') - current = LlmAgent(name='current', sub_agents=[target]) - root = LlmAgent(name='root', sub_agents=[current]) - - ctx, _ = _make_parent_ctx() - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'wf/parent/current@1' - child_ctx1.parent_ctx = ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'wf/parent/current@1/target@1' - child_ctx2.parent_ctx = child_ctx1 - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # First call: current node - args1, kwargs1 = calls[0] - assert args1[0] is ctx - assert args1[1] is current - assert args1[2] == 'current' - assert args1[4] == '1' - - # Second call: target node (transferred CHILD) - args2, kwargs2 = calls[1] - assert args2[0] is child_ctx1 - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '1' - - -@pytest.mark.asyncio -async def test_scheduler_handles_sibling_transfer(): - """Scheduler processes SIBLING relationship by sharing parent context.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - target = LlmAgent(name='target') - root = LlmAgent(name='root', sub_agents=[current, target]) - - ctx, _ = _make_parent_ctx() - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'wf/parent/current@1' - child_ctx1.parent_ctx = ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'wf/parent/target@1' - child_ctx2.parent_ctx = ctx - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # Second call: target node (transferred SIBLING) - args2, kwargs2 = calls[1] - assert args2[0] is ctx - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '1' - - -@pytest.mark.asyncio -async def test_scheduler_handles_parent_transfer(): - """Scheduler processes PARENT relationship by truncating parent context.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - target = LlmAgent(name='target', sub_agents=[current]) - root = LlmAgent(name='root', sub_agents=[target]) - - ctx, _ = _make_parent_ctx() - ctx._child_run_counters = {'target': 1} - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - target_ctx = MagicMock() - target_ctx.node_path = 'wf/parent/target@1' - target_ctx.node = target - target_ctx.parent_ctx = ctx - target_ctx._invocation_context = ctx._invocation_context - target_ctx._child_run_counters = {} - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'wf/parent/target@1/current@1' - child_ctx1.node = current - child_ctx1.parent_ctx = target_ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'wf/parent/target@2' - child_ctx2.parent_ctx = ctx - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - target_ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # Second call: target node (transferred ANCESTOR) - args2, kwargs2 = calls[1] - assert args2[0] is ctx - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '2' - - -@pytest.mark.asyncio -async def test_scheduler_raises_value_error_on_self_transfer(): - """Scheduler raises ValueError when agent attempts to transfer to itself.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - root = LlmAgent(name='root', sub_agents=[current]) - - ctx, _ = _make_parent_ctx() - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - child_ctx = MagicMock() - child_ctx.node_path = 'wf/parent/current@1' - child_ctx.parent_ctx = ctx - child_ctx.actions = EventActions(transfer_to_agent='current') - child_ctx.error = None - child_ctx.interrupt_ids = set() - child_ctx._invocation_context = ctx._invocation_context - child_ctx._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(return_value=child_ctx) - - # Act & Assert - with pytest.raises(ValueError, match='cannot transfer to itself'): - await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - -@pytest.mark.asyncio -async def test_scheduler_handles_parent_transfer_bypassed_on_resume(): - """Scheduler processes PARENT relationship when parent was bypassed on resume.""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - current = LlmAgent(name='current') - target = LlmAgent(name='target', sub_agents=[current]) - root = LlmAgent(name='root', sub_agents=[target]) - - current.parent_agent = target - target.parent_agent = root - - ctx, _ = _make_parent_ctx() - ctx.node = None - ctx._child_run_counters = {'target': 1} - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - target_ctx = MagicMock() - target_ctx.node_path = 'current@1' - target_ctx.node = current - target_ctx.parent_ctx = ctx - target_ctx._invocation_context = ctx._invocation_context - target_ctx._child_run_counters = {} - - child_ctx1 = MagicMock() - child_ctx1.node_path = 'current@1' - child_ctx1.node = current - child_ctx1.parent_ctx = ctx - child_ctx1.actions = EventActions(transfer_to_agent='target') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ctx._invocation_context - child_ctx1._child_run_counters = {} - - child_ctx2 = MagicMock() - child_ctx2.node_path = 'target@2' - child_ctx2.parent_ctx = ctx - child_ctx2.actions = EventActions() - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ctx._invocation_context - child_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock(side_effect=[child_ctx1, child_ctx2]) - - # Act - final_ctx = await scheduler( - ctx, - current, - 'input', - node_name='current', - run_id='1', - ) - - # Assert - assert final_ctx is child_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 2 - - # First call: current node - args1, kwargs1 = calls[0] - assert args1[0] is ctx - assert args1[1] is current - assert args1[2] == 'current' - assert args1[4] == '1' - - # Second call: target node (transferred PARENT) - args2, kwargs2 = calls[1] - assert args2[0] is ctx - assert args2[1] is target - assert args2[2] == 'target' - assert args2[4] == '2' - - -@pytest.mark.asyncio -async def test_scheduler_handles_three_layer_agent_transfers_round_trip(): - """Verify 3-layer agent transfers (Root -> Child -> Grandchild -> Child -> Root).""" - # Arrange - from google.adk.agents.llm_agent import LlmAgent - from google.adk.events.event_actions import EventActions - from google.adk.workflow._dynamic_node_scheduler import DynamicNodeScheduler - from google.adk.workflow._workflow import _LoopState - - grandchild = LlmAgent(name='grandchild') - child = LlmAgent(name='child', sub_agents=[grandchild]) - root = LlmAgent(name='root', sub_agents=[child]) - - grandchild.parent_agent = child - child.parent_agent = root - - # ctx is the root context (node = None, parent_ctx = None) - ctx = MagicMock() - ctx.node = None - ctx.parent_ctx = None - ctx.node_path = '' - ctx._child_run_counters = {'child': 1} - ctx._output_for_ancestors = [] - - ic = MagicMock() - ic.invocation_id = 'inv-1' - ic.session = MagicMock() - ic.session.state = {} - ic.session.events = [] - ic.run_config = None - ctx._invocation_context = ic - - # root_ctx is the parent context of child (node = root, parent_ctx = ctx) - root_ctx = MagicMock() - root_ctx.node = root - root_ctx.parent_ctx = ctx - root_ctx.node_path = 'root@1' - root_ctx._child_run_counters = {'child': 1} - root_ctx._invocation_context = ic - - ls = _LoopState() - scheduler = DynamicNodeScheduler(state=ls) - - # Step 1: root delegates to child - child_ctx1 = MagicMock() - child_ctx1.node_path = 'root@1/child@1' - child_ctx1.node = child - child_ctx1.parent_ctx = root_ctx - child_ctx1.actions = EventActions(transfer_to_agent='grandchild') - child_ctx1.error = None - child_ctx1.interrupt_ids = set() - child_ctx1._invocation_context = ic - child_ctx1._child_run_counters = {} - - # Step 2: child delegates to grandchild - grandchild_ctx1 = MagicMock() - grandchild_ctx1.node_path = 'root@1/child@1/grandchild@1' - grandchild_ctx1.node = grandchild - grandchild_ctx1.parent_ctx = child_ctx1 - grandchild_ctx1.actions = EventActions(transfer_to_agent='child') - grandchild_ctx1.error = None - grandchild_ctx1.interrupt_ids = set() - grandchild_ctx1._invocation_context = ic - grandchild_ctx1._child_run_counters = {} - - # Step 3: grandchild transfers back to child - child_ctx2 = MagicMock() - child_ctx2.node_path = 'root@1/child@2' - child_ctx2.node = child - child_ctx2.parent_ctx = root_ctx - child_ctx2.actions = EventActions(transfer_to_agent='root') - child_ctx2.error = None - child_ctx2.interrupt_ids = set() - child_ctx2._invocation_context = ic - child_ctx2._child_run_counters = {} - - # Step 4: child transfers back to root - ctx._child_run_counters = {'root': 1} - root_ctx2 = MagicMock() - root_ctx2.node_path = 'root@2' - root_ctx2.node = root - root_ctx2.parent_ctx = ctx - root_ctx2.actions = EventActions() - root_ctx2.error = None - root_ctx2.interrupt_ids = set() - root_ctx2._invocation_context = ic - root_ctx2._child_run_counters = {} - - scheduler._run_node_internal = AsyncMock( - side_effect=[child_ctx1, grandchild_ctx1, child_ctx2, root_ctx2] - ) - - # Act - final_ctx = await scheduler( - root_ctx, - child, - 'input', - node_name='child', - run_id='1', - ) - - # Assert - assert final_ctx is root_ctx2 - calls = scheduler._run_node_internal.call_args_list - assert len(calls) == 4 - - # 1st call: child (scheduled by root) - args1, kwargs1 = calls[0] - assert args1[0] is root_ctx - assert args1[1] is child - assert args1[2] == 'child' - assert args1[4] == '1' - - # 2nd call: grandchild (transferred CHILD from child) - args2, kwargs2 = calls[1] - assert args2[0] is child_ctx1 - assert args2[1] is grandchild - assert args2[2] == 'grandchild' - assert args2[4] == '1' - - # 3rd call: child (transferred PARENT from grandchild) - args3, kwargs3 = calls[2] - assert args3[0] is root_ctx - assert args3[1] is child - assert args3[2] == 'child' - assert args3[4] == '2' - - # 4th call: root (transferred PARENT from child) - args4, kwargs4 = calls[3] - assert args4[0] is ctx - assert args4[1] is root - assert args4[2] == 'root' - assert args4[4] == '2' - - # ========================================================================= # Replay Sequence Ordering preservation for Dynamic Nodes # ========================================================================= From 9d0a5a8615505cfddf358d811f8c172b2314b167 Mon Sep 17 00:00:00 2001 From: Haran Rajkumar Date: Tue, 30 Jun 2026 15:30:59 -0700 Subject: [PATCH 3/6] test(samples): add replay tests for the interactions_api sample Add replay test fixtures covering the documented flows of the contributing/samples/models/interactions_api sample (basic text, Google Search grounding, multi-turn stateful recall, and a custom function tool), so the sample is exercised by tests/unittests/test_samples.py. Also fix the shared sample test harness (cli/agent_test_runner.py) to support the Interactions API: - Exclude the volatile interaction_id (a server-issued token) and turn_complete fields from fixture comparison and from rebuilt fixtures, matching how other non-reproducible fields (timestamps, usage metadata, etc.) are already handled. The Interactions API stamps these onto every model response and the replay MockModel cannot reproduce them. - Drive all turns of a fixture rebuild on a single persistent event loop. The sync Runner.run() uses asyncio.run() per call, which closes the loop the model's cached async api_client is bound to, so subsequent turns fail with "Event loop is closed" and corrupt multi-turn fixtures. Reusing one loop keeps the cached client valid across the conversation. Co-authored-by: Haran Rajkumar PiperOrigin-RevId: 940691339 --- .../interactions_api/tests/basic_text.json | 37 +++++ .../tests/custom_function_weather.json | 86 ++++++++++ .../tests/google_search_1984.json | 40 +++++ .../tests/google_search_france.json | 40 +++++ .../interactions_api/tests/multi_turn.json | 152 ++++++++++++++++++ src/google/adk/cli/agent_test_runner.py | 50 +++++- 6 files changed, 402 insertions(+), 3 deletions(-) create mode 100644 contributing/samples/models/interactions_api/tests/basic_text.json create mode 100644 contributing/samples/models/interactions_api/tests/custom_function_weather.json create mode 100644 contributing/samples/models/interactions_api/tests/google_search_1984.json create mode 100644 contributing/samples/models/interactions_api/tests/google_search_france.json create mode 100644 contributing/samples/models/interactions_api/tests/multi_turn.json diff --git a/contributing/samples/models/interactions_api/tests/basic_text.json b/contributing/samples/models/interactions_api/tests/basic_text.json new file mode 100644 index 00000000000..f200e01e94b --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/basic_text.json @@ -0,0 +1,37 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "Hello! What can you help me with?" + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "Hello! I am \"interactions_test_agent,\" an agent designed for testing the Interactions API integration.\n\nI can help you with two main tasks:\n1. **Searching the web:** If you have questions that require up-to-date information, I can search the web for you.\n2. **Checking the weather:** I can provide the current weather for a specific city.\n\nHow can I assist you today?" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/custom_function_weather.json b/contributing/samples/models/interactions_api/tests/custom_function_weather.json new file mode 100644 index 00000000000..f5f5bcad550 --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/custom_function_weather.json @@ -0,0 +1,86 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "What's the weather like in Tokyo?" + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "city": "Tokyo" + }, + "id": "fc-1", + "name": "get_current_weather" + } + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "longRunningToolIds": [], + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-1", + "name": "get_current_weather", + "response": { + "city": "Tokyo", + "condition": "Partly Cloudy", + "humidity": 60, + "temperature_f": 68 + } + } + } + ], + "role": "user" + }, + "id": "e-3", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "I have retrieved the weather information for Tokyo. It is currently 68\u00b0F and partly cloudy with 60% humidity." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-4", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/google_search_1984.json b/contributing/samples/models/interactions_api/tests/google_search_1984.json new file mode 100644 index 00000000000..317116c7e1d --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/google_search_1984.json @@ -0,0 +1,40 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "Use google search to find out who wrote the novel '1984'." + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "search_suggestions='\\n
\\n
\\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n
\\n
\\n \\n
'" + }, + { + "text": "The novel *1984* was written by the English author **George Orwell** (whose real name was Eric Arthur Blair). It was published in 1949." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/google_search_france.json b/contributing/samples/models/interactions_api/tests/google_search_france.json new file mode 100644 index 00000000000..f2ca266d82a --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/google_search_france.json @@ -0,0 +1,40 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "Search for the capital of France." + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "search_suggestions='\\n
\\n
\\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n \\n
\\n
\\n \\n
'" + }, + { + "text": "I have searched for the capital of France and confirmed that it is Paris." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/contributing/samples/models/interactions_api/tests/multi_turn.json b/contributing/samples/models/interactions_api/tests/multi_turn.json new file mode 100644 index 00000000000..aeb4d0494b6 --- /dev/null +++ b/contributing/samples/models/interactions_api/tests/multi_turn.json @@ -0,0 +1,152 @@ +{ + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "My favorite color is blue. Just acknowledge this, don't use any tools." + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "I acknowledge that your favorite color is blue." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-2", + "invocationId": "i-1", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "user", + "content": { + "parts": [ + { + "text": "What's the weather like in London?" + } + ], + "role": "user" + }, + "id": "e-3", + "invocationId": "i-2", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "city": "London" + }, + "id": "fc-1", + "name": "get_current_weather" + } + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-4", + "invocationId": "i-2", + "longRunningToolIds": [], + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-1", + "name": "get_current_weather", + "response": { + "city": "London", + "condition": "Cloudy", + "humidity": 78, + "temperature_f": 59 + } + } + } + ], + "role": "user" + }, + "id": "e-5", + "invocationId": "i-2", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "I have checked the weather in London for you. It is currently cloudy with a temperature of 59\u00b0F and 78% humidity." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-6", + "invocationId": "i-2", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + }, + { + "author": "user", + "content": { + "parts": [ + { + "text": "What is my favorite color that I mentioned earlier in our conversation?" + } + ], + "role": "user" + }, + "id": "e-7", + "invocationId": "i-3", + "nodeInfo": { + "path": "" + } + }, + { + "author": "interactions_test_agent", + "content": { + "parts": [ + { + "text": "Your favorite color, which you mentioned earlier, is blue." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "id": "e-8", + "invocationId": "i-3", + "nodeInfo": { + "path": "interactions_test_agent@1" + } + } + ] +} diff --git a/src/google/adk/cli/agent_test_runner.py b/src/google/adk/cli/agent_test_runner.py index 60fa5b9c49a..5312b94d50e 100644 --- a/src/google/adk/cli/agent_test_runner.py +++ b/src/google/adk/cli/agent_test_runner.py @@ -51,6 +51,17 @@ "citation_metadata", } +# The Interactions API stamps these volatile, non-reproducible fields onto every +# model response (interaction_id is a server-issued token; turn_complete is only +# emitted by the live API), so they are excluded from fixture comparison too. +# They are kept in a separate, private constant so the value of the public +# EXCLUDED_EVENT_FIELDS stays stable for the API breaking-change detector. +_EXTRA_EXCLUDED_EVENT_FIELDS = frozenset({"interaction_id", "turn_complete"}) + +_ALL_EXCLUDED_EVENT_FIELDS = ( + EXCLUDED_EVENT_FIELDS | _EXTRA_EXCLUDED_EVENT_FIELDS +) + # Read target folder from environment def get_test_files( @@ -175,7 +186,7 @@ def normalize_events(events, is_json=False): for e in events: if is_json: d = dict(e) - for k in EXCLUDED_EVENT_FIELDS: + for k in _ALL_EXCLUDED_EVENT_FIELDS: d.pop(k, None) d.pop(alias_generators.to_camel(k), None) d = {k: v for k, v in d.items() if v is not None} @@ -183,7 +194,7 @@ def normalize_events(events, is_json=False): d = e.model_dump( mode="json", by_alias=True, - exclude=EXCLUDED_EVENT_FIELDS, + exclude=_ALL_EXCLUDED_EVENT_FIELDS, exclude_none=True, ) @@ -706,6 +717,7 @@ def mock_randint(a, b): def rebuild_tests(path: str): """Discovers test files and rebuilds them by running the agent live.""" + import asyncio import json import sys @@ -734,6 +746,7 @@ def rebuild_tests(path: str): # Add agent_dir.parent to sys.path so relative imports work sys_path_saved = list(sys.path) sys.path.insert(0, str(agent_dir.parent)) + rebuild_loop = None try: import random @@ -774,6 +787,29 @@ def rebuild_tests(path: str): else InMemoryRunner(root_agent=agent_or_app) ) + # Drive every turn of this fixture on a single, persistent event loop. + # The sync Runner.run() spins up a fresh loop per call via asyncio.run() + # and closes it afterwards. For multi-turn fixtures that closes the loop + # the model's cached async api_client was bound to, so subsequent turns + # raise "Event loop is closed" (e.g. with the Interactions API). Reusing + # one loop for all turns keeps the client valid across the conversation. + rebuild_loop = asyncio.new_event_loop() + + def run_turn(content): + session = runner.session + + async def _collect(): + events = [] + async for event in runner.runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=content, + ): + events.append(event) + return events + + return rebuild_loop.run_until_complete(_collect()) + new_events = [] inv_counter = 1 @@ -857,7 +893,7 @@ def get_next_fc_id(): content=msg, ) - run_events = runner.run(msg) + run_events = run_turn(msg) # Build mapping from old IDs to new agent IDs for e in run_events: @@ -891,6 +927,10 @@ def get_next_fc_id(): "cache_metadata", "logprobs_result", "citation_metadata", + # Volatile/non-replayable Interactions API state; keeping + # these out keeps rebuilt fixtures deterministic. + "interaction_id", + "turn_complete", }, ) for e in new_events @@ -931,6 +971,10 @@ def get_next_fc_id(): except Exception as e: print(f"Error rebuilding {test_file}: {e}") finally: + # Always close the per-fixture event loop, even if a turn raised, so we + # don't leak unclosed loops and emit resource warnings. + if rebuild_loop is not None: + rebuild_loop.close() sys.path = sys_path_saved From 38700324134f8d392a577f0cda78e4c4400e69d2 Mon Sep 17 00:00:00 2001 From: Kathy Wu Date: Tue, 30 Jun 2026 16:28:17 -0700 Subject: [PATCH 4/6] fix: in mcp mtls logic, check for authorization header case insensitively httpx internally normalizes HTTP header keys to lowercase, so the check should be case-insensitive closes https://github.com/google/adk-python/issues/6200 Co-authored-by: Kathy Wu PiperOrigin-RevId: 940719476 --- .../adk/tools/mcp_tool/mcp_session_manager.py | 2 +- .../mcp_tool/test_mcp_session_manager.py | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 59b0dab478a..8407ee7ab80 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -391,7 +391,7 @@ async def before_request( ) return - if 'Authorization' in headers: + if any(k.lower() == 'authorization' for k in headers): logger.debug('Authorization header already present, not overwriting') return diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index e37fdfe67ce..b769eca55e1 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -1298,6 +1298,28 @@ def mock_refresh(req): assert headers["Authorization"] == "Bearer refreshed_token" + @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported") + @pytest.mark.parametrize( + "existing_header_key", + ["Authorization", "authorization", "AUTHORIZATION", "authORIZATION"], + ) + @pytest.mark.asyncio + async def test_before_request_skips_refresh_if_authorization_header_exists_case_insensitive( + self, existing_header_key + ): + mock_creds = Mock() + mock_creds.expired = True + mock_creds.token = "new_token" + mock_creds.refresh = Mock() + + credentials = _RefreshableAsyncCredentials(mock_creds) + headers = {existing_header_key: "Bearer existing_token"} + + await credentials.before_request(None, "GET", "http://example.com", headers) + + mock_creds.refresh.assert_not_called() + assert headers == {existing_header_key: "Bearer existing_token"} + class TestGoogleAuthAsyncByteStream: From 196f7708f82cca3f500185c2d413afcc3580edae Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Tue, 30 Jun 2026 16:46:40 -0700 Subject: [PATCH 5/6] chore: Rebase static graph dispatch onto dynamic scheduler Use Context.run_node from within _workflow.py so static graphs also use the unified transfer loop. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 940728866 --- src/google/adk/workflow/_workflow.py | 29 +++++----- .../unittests/workflow/test_workflow_hitl.py | 11 ++-- .../unittests/workflow/test_workflow_live.py | 57 +++++++++++-------- .../workflow/test_workflow_schema.py | 8 ++- 4 files changed, 58 insertions(+), 47 deletions(-) diff --git a/src/google/adk/workflow/_workflow.py b/src/google/adk/workflow/_workflow.py index dc9527cff0e..d022dcd8528 100644 --- a/src/google/adk/workflow/_workflow.py +++ b/src/google/adk/workflow/_workflow.py @@ -37,7 +37,6 @@ from ._dynamic_node_scheduler import DynamicNodeState from ._graph import EdgeItem from ._graph import Graph -from ._node_runner import NodeRunner from ._node_state import NodeState from ._node_status import NodeStatus from ._trigger import Trigger @@ -536,7 +535,7 @@ def _start_node_task( node_name: str, trigger: Trigger, ) -> None: - """Create NodeRunner and start asyncio task for a node.""" + """Start asyncio task for scheduling and executing a node.""" assert self.graph is not None @@ -602,22 +601,24 @@ async def return_ctx(): key ].isolation_scope - runner = NodeRunner( - node=node, - parent_ctx=ctx, - run_id=run_id, - use_as_output=is_terminal, - use_sub_branch=trigger.use_sub_branch, - override_branch=trigger.branch, - override_isolation_scope=self._compute_isolation_scope_for_node( - node, trigger, ctx, run_id - ), - ) resume_inputs = ( dict(node_state.resume_inputs) if node_state.resume_inputs else None ) loop_state.pending_tasks[node_name] = asyncio.create_task( - runner.run(node_input=trigger.input, resume_inputs=resume_inputs) + ctx._run_node_internal( + node, + node_input=trigger.input, + use_sub_branch=trigger.use_sub_branch, + override_branch=trigger.branch, + override_isolation_scope=self._compute_isolation_scope_for_node( + node, trigger, ctx, run_id + ), + return_ctx=True, + resume_inputs=resume_inputs, + run_id=run_id, + use_as_output=is_terminal, + skip_run_id_validation=True, + ) ) def _make_schedule_dynamic_node( diff --git a/tests/unittests/workflow/test_workflow_hitl.py b/tests/unittests/workflow/test_workflow_hitl.py index 8a2558b1533..dbe485aca54 100644 --- a/tests/unittests/workflow/test_workflow_hitl.py +++ b/tests/unittests/workflow/test_workflow_hitl.py @@ -1858,15 +1858,16 @@ async def _run_impl( async def test_multiple_invocations_isolation(request: pytest.FixtureRequest): """Verify that a new invocation ignores events from a previous invocation.""" + run_counts = [] + class CounterNode(BaseNode): name: str = Field(default='counter_node') - run_count: int = Field(default=0) async def _run_impl( self, *, ctx: Context, node_input: Any ) -> AsyncGenerator[Any, None]: - self.run_count += 1 - yield f'Run {self.run_count}' + run_counts.append(1) + yield f'Run {len(run_counts)}' node_a = CounterNode() wf = Workflow(name='wf', edges=[(START, node_a)]) @@ -1883,7 +1884,7 @@ async def _run_impl( ): events1.append(event) - assert node_a.run_count == 1 + assert len(run_counts) == 1 # Invocation 2 (New invocation in SAME session) msg2 = types.Content(parts=[types.Part(text='go 2')], role='user') @@ -1894,7 +1895,7 @@ async def _run_impl( events2.append(event) # If isolation works, CounterNode should run AGAIN! - assert node_a.run_count == 2 + assert len(run_counts) == 2 @pytest.mark.asyncio diff --git a/tests/unittests/workflow/test_workflow_live.py b/tests/unittests/workflow/test_workflow_live.py index 697b5775efb..d5207d666dc 100644 --- a/tests/unittests/workflow/test_workflow_live.py +++ b/tests/unittests/workflow/test_workflow_live.py @@ -27,6 +27,7 @@ from google.adk.workflow._base_node import START from google.adk.workflow._workflow import Workflow from google.genai import types +from pydantic import Field import pytest from . import testing_utils @@ -39,6 +40,7 @@ class _MockNonLiveNode(BaseNode): called: bool = False actual_input: Any = None + shared_state: dict[str, Any] = Field(default_factory=dict) def __init__(self, *, name: str): super().__init__(name=name) @@ -51,6 +53,8 @@ async def _run_impl( ) -> AsyncGenerator[Any, None]: self.called = True self.actual_input = node_input + self.shared_state["called"] = True + self.shared_state["actual_input"] = node_input yield Event(output=f"{self.name}_output") @@ -72,29 +76,6 @@ async def _run_impl( yield Event(output=self.output_value) -class _DynamicLiveSchedulerNode(BaseNode): - """A node that dynamically schedules a child live node using ctx.run_node().""" - - child_node: BaseNode | None = None - child_output: Any = None - - def __init__(self, *, name: str, child_node: BaseNode): - super().__init__(name=name, rerun_on_resume=True) - self.child_node = child_node - - async def _run_impl( - self, - *, - ctx: Context, - node_input: Any, - ) -> AsyncGenerator[Any, None]: - if self.child_node: - self.child_output = await ctx.run_node( - self.child_node, node_input=node_input - ) - yield Event(output=f"{self.name}_output") - - # --- Live Workflow Unit Tests (TDD) --- @@ -468,6 +449,30 @@ async def test_nested_live_node_and_outer_live_node(): @pytest.mark.asyncio async def test_dynamic_node_scheduling_of_live_node(): """CUJ 4: A node in workflow dynamically schedules a live node using ctx.run_node().""" + + class _DynamicLiveSchedulerNode(BaseNode): + """A node that dynamically schedules a child live node using ctx.run_node().""" + + child_node: BaseNode | None = None + child_output: Any = None + shared_state: dict[str, Any] = Field(default_factory=dict) + + def __init__(self, *, name: str, child_node: BaseNode): + super().__init__(name=name, rerun_on_resume=True) + self.child_node = child_node + + async def _run_impl( + self, + *, + ctx: Context, + node_input: Any, + ) -> AsyncGenerator[Any, None]: + if self.child_node: + output = await ctx.run_node(self.child_node, node_input=node_input) + self.child_output = output + self.shared_state["child_output"] = output + yield Event(output=f"{self.name}_output") + mock_model = testing_utils.MockModel.create( responses=[ LlmResponse( @@ -531,7 +536,9 @@ async def test_dynamic_node_scheduling_of_live_node(): {"result": "DynamicLiveNode_output"}, "SchedulerNode_output", ] - assert scheduler_node.child_output == {"result": "DynamicLiveNode_output"} + assert scheduler_node.shared_state.get("child_output") == { + "result": "DynamicLiveNode_output" + } # Assert content events content_texts = [ @@ -663,7 +670,7 @@ async def test_single_turn_agent_runs_as_non_live_in_live_session(): outputs = [e.output for e in events if e.output is not None] assert outputs == ["initial_text_input", "capture_output"] - assert capture.actual_input == "SingleTurn_output" + assert capture.shared_state.get("actual_input") == "SingleTurn_output" # Verify that the model received the initial_text_input (node_input) and NOT the live queue audio assert len(mock_model.requests) == 1 assert ( diff --git a/tests/unittests/workflow/test_workflow_schema.py b/tests/unittests/workflow/test_workflow_schema.py index 5fdfc706ffd..0574aad6d09 100644 --- a/tests/unittests/workflow/test_workflow_schema.py +++ b/tests/unittests/workflow/test_workflow_schema.py @@ -26,6 +26,7 @@ from google.adk.workflow._workflow import Workflow from google.genai import types from pydantic import BaseModel +from pydantic import ValidationError import pytest from .. import testing_utils @@ -630,9 +631,10 @@ async def _run_impl( msg = types.Content(parts=[types.Part(text='hello')], role='user') # We expect it to raise ValidationError - from pydantic import ValidationError - - with pytest.raises(ValidationError): + with pytest.raises( + ValueError, + match=r"Runtime schema validation failed for dynamic node 'node'", + ): async for event in runner.run_async( user_id='u', session_id=session.id, new_message=msg ): From c6dec00a6e6d1370a3c59808d651d9528ebc9708 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Tue, 30 Jun 2026 18:05:56 -0700 Subject: [PATCH 6/6] feat(telemetry): add gen_ai.invoke_agent.{inference,tool}_calls metrics Adds the two per-invocation call-count metrics from semconv #336 (open-telemetry/semantic-conventions-genai#336): * gen_ai.invoke_agent.inference_calls: model calls in one invoke_agent span. * gen_ai.invoke_agent.tool_calls: tool calls in one invoke_agent span. Both are counted on the per-span TelemetryContext at the call sites (record_inference_telemetry, record_tool_execution) and flushed in record_agent_invocation, keyed by gen_ai.agent.name. One point per invoke_agent span, so a multi-agent invocation does not double count. Failed calls still count; only client-side tool calls count. Always emitted (a no-op without a MeterProvider). PiperOrigin-RevId: 940760757 --- src/google/adk/agents/invocation_context.py | 10 ---- src/google/adk/telemetry/_instrumentation.py | 55 +------------------ src/google/adk/telemetry/_metrics.py | 50 ----------------- .../telemetry/functional_node_test_cases.py | 12 ---- .../telemetry/functional_test_cases.py | 12 ---- .../telemetry/functional_test_helpers.py | 10 ---- .../telemetry/test_instrumentation.py | 5 +- 7 files changed, 5 insertions(+), 149 deletions(-) diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 5dac20f6231..6afa0d3573e 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -18,7 +18,6 @@ from typing import Any from typing import cast from typing import Optional -from typing import TYPE_CHECKING from google.adk.platform import uuid as platform_uuid from google.genai import types @@ -48,9 +47,6 @@ from .run_config import RunConfig from .transcription_entry import TranscriptionEntry -if TYPE_CHECKING: - from google.adk.telemetry._instrumentation import TelemetryContext - class LlmCallsLimitExceededError(Exception): """Error thrown when the number of LLM calls exceed the limit.""" @@ -274,12 +270,6 @@ class InvocationContext(BaseModel): of this invocation. """ - _invoke_agent_telemetry_context: Optional[TelemetryContext] = PrivateAttr( - default=None - ) - """TelemetryContext of the active ``invoke_agent`` span, if any. - """ - @property def is_resumable(self) -> bool: """Returns whether the current invocation is resumable.""" diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index e12b7c9b054..2284fb839dc 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -27,11 +27,11 @@ from . import _metrics from . import tracing +from ..events import event as event_lib if TYPE_CHECKING: from ..agents.base_agent import BaseAgent from ..agents.invocation_context import InvocationContext - from ..events import event as event_lib from ..models.llm_request import LlmRequest from ..models.llm_response import LlmResponse from ..tools.base_tool import BaseTool @@ -78,31 +78,11 @@ class TelemetryContext: error_type: str | None = None span: tracing.GenerateContentSpan | trace.Span | None = None _llm_responses: list[LlmResponse] = dataclasses.field(default_factory=list) - _inference_call_count: int = 0 - _tool_call_count: int = 0 @property def llm_responses(self) -> list[LlmResponse]: return self._llm_responses - @property - def inference_call_count(self) -> int: - """Number of model calls counted against this invoke_agent span.""" - return self._inference_call_count - - def increment_inference_calls(self) -> None: - """Counts one model call against this span (including calls that raised).""" - self._inference_call_count += 1 - - def increment_tool_calls(self) -> None: - """Counts one tool call against this invoke_agent span.""" - self._tool_call_count += 1 - - @property - def tool_call_count(self) -> int: - """Number of tool calls counted against this invoke_agent span.""" - return self._tool_call_count - def record_llm_response( self, invocation_context: InvocationContext, response: LlmResponse ) -> None: @@ -130,30 +110,6 @@ def _record_agent_metrics( logger.exception("Failed to record agent metrics for agent %s", agent_name) -def _flush_invoke_agent_metrics( - tel_ctx: TelemetryContext, agent_name: str -) -> None: - """Flushes this span's accumulated inference/tool-call metrics.""" - _metrics.record_invoke_agent_inference_calls( - agent_name, tel_ctx.inference_call_count - ) - _metrics.record_invoke_agent_tool_calls(agent_name, tel_ctx.tool_call_count) - - -def _accumulate_invoke_agent_tool_call(ctx: InvocationContext) -> None: - """Counts one tool call against the active invoke_agent span.""" - span_tel_ctx = ctx._invoke_agent_telemetry_context # pylint: disable=protected-access - if span_tel_ctx is not None: - span_tel_ctx.increment_tool_calls() - - -def _accumulate_invoke_agent_inference_call(ctx: InvocationContext) -> None: - """Counts one model call against the active invoke_agent span.""" - span_tel_ctx = ctx._invoke_agent_telemetry_context # pylint: disable=protected-access - if span_tel_ctx is not None: - span_tel_ctx.increment_inference_calls() - - @contextlib.asynccontextmanager async def record_agent_invocation( ctx: InvocationContext, agent: BaseAgent @@ -163,13 +119,11 @@ async def record_agent_invocation( caught_error: Exception | None = None span: trace.Span | None = None span_name = f"invoke_agent {agent.name}" - tel_ctx = TelemetryContext() try: with tracing.tracer.start_as_current_span(span_name) as s: span = s tracing.trace_agent_invocation(span, agent, ctx) - tel_ctx.otel_context = context_api.get_current() - ctx._invoke_agent_telemetry_context = tel_ctx # pylint: disable=protected-access + tel_ctx = TelemetryContext(otel_context=context_api.get_current()) yield tel_ctx except Exception as e: caught_error = e @@ -182,7 +136,6 @@ async def record_agent_invocation( getattr(getattr(ctx, "session", None), "events", []), caught_error, ) - _flush_invoke_agent_metrics(tel_ctx, agent.name) @contextlib.asynccontextmanager @@ -190,7 +143,7 @@ async def record_tool_execution( tool: BaseTool, agent: BaseAgent, function_args: dict[str, object], - invocation_context: InvocationContext, + invocation_context: InvocationContext | None = None, ) -> AsyncIterator[TelemetryContext]: """Unified context manager for consolidated tool execution telemetry.""" start_time = time.monotonic() @@ -218,7 +171,6 @@ async def record_tool_execution( invocation_context=invocation_context, error_type=tel_ctx.error_type, ) - _accumulate_invoke_agent_tool_call(invocation_context) finally: try: _metrics.record_tool_execution_duration( @@ -253,7 +205,6 @@ async def record_inference_telemetry( yield tel_ctx finally: inference_error = sys.exc_info()[1] - _accumulate_invoke_agent_inference_call(invocation_context) agent = invocation_context.agent elapsed_s = _get_elapsed_s(tel_ctx.span, start_time) try: diff --git a/src/google/adk/telemetry/_metrics.py b/src/google/adk/telemetry/_metrics.py index 7be66bc7a8c..eac65b48ca0 100644 --- a/src/google/adk/telemetry/_metrics.py +++ b/src/google/adk/telemetry/_metrics.py @@ -146,44 +146,6 @@ gen_ai_metrics.create_gen_ai_client_operation_duration(meter) ) _client_token_usage = gen_ai_metrics.create_gen_ai_client_token_usage(meter) -_invoke_agent_inference_calls = meter.create_histogram( - "gen_ai.invoke_agent.inference_calls", - unit="1", - description="Number of inference (model) calls per agent invocation.", - explicit_bucket_boundaries_advisory=[ - 1, - 2, - 3, - 4, - 5, - 6, - 8, - 12, - 16, - 24, - 32, - 64, - ], -) -_invoke_agent_tool_calls = meter.create_histogram( - "gen_ai.invoke_agent.tool_calls", - unit="1", - description="Number of tool calls per agent invocation.", - explicit_bucket_boundaries_advisory=[ - 1, - 2, - 3, - 4, - 5, - 6, - 8, - 12, - 16, - 24, - 32, - 64, - ], -) def record_agent_invocation_duration( @@ -198,18 +160,6 @@ def record_agent_invocation_duration( _agent_invocation_duration.record(elapsed_s, attributes=attrs) -def record_invoke_agent_inference_calls(agent_name: str, count: int): - """Records the number of inference (model) calls in an agent invocation.""" - attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name} - _invoke_agent_inference_calls.record(count, attributes=attrs) - - -def record_invoke_agent_tool_calls(agent_name: str, count: int): - """Records the number of tool calls in an agent invocation.""" - attrs = {gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name} - _invoke_agent_tool_calls.record(count, attributes=attrs) - - def record_agent_request_size( agent_name: str, user_content: types.Content | None ): diff --git a/tests/unittests/telemetry/functional_node_test_cases.py b/tests/unittests/telemetry/functional_node_test_cases.py index 67774235786..56e58a7fa7a 100644 --- a/tests/unittests/telemetry/functional_node_test_cases.py +++ b/tests/unittests/telemetry/functional_node_test_cases.py @@ -2900,12 +2900,6 @@ value=NON_DETERMINISTIC, ), }), - "gen_ai.invoke_agent.inference_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), - }), - "gen_ai.invoke_agent.tool_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), - }), } @@ -2953,12 +2947,6 @@ value=NON_DETERMINISTIC, ), }), - "gen_ai.invoke_agent.inference_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), - }), - "gen_ai.invoke_agent.tool_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), - }), } diff --git a/tests/unittests/telemetry/functional_test_cases.py b/tests/unittests/telemetry/functional_test_cases.py index 0bb7b9b3695..44c352a4c88 100644 --- a/tests/unittests/telemetry/functional_test_cases.py +++ b/tests/unittests/telemetry/functional_test_cases.py @@ -2294,12 +2294,6 @@ value=NON_DETERMINISTIC, ), }), - "gen_ai.invoke_agent.inference_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), - }), - "gen_ai.invoke_agent.tool_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), - }), } @@ -2347,12 +2341,6 @@ value=NON_DETERMINISTIC, ), }), - "gen_ai.invoke_agent.inference_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=2), - }), - "gen_ai.invoke_agent.tool_calls": frozenset({ - MetricPoint(attributes={"gen_ai.agent.name": AGENT_NAME}, value=1), - }), } diff --git a/tests/unittests/telemetry/functional_test_helpers.py b/tests/unittests/telemetry/functional_test_helpers.py index 82f997b6979..e4ae30a3607 100644 --- a/tests/unittests/telemetry/functional_test_helpers.py +++ b/tests/unittests/telemetry/functional_test_helpers.py @@ -339,16 +339,6 @@ class HistogramSpec(NamedTuple): attr="_client_token_usage", metric_name="gen_ai.client.token.usage", ), - HistogramSpec( - module=_metrics, - attr="_invoke_agent_inference_calls", - metric_name="gen_ai.invoke_agent.inference_calls", - ), - HistogramSpec( - module=_metrics, - attr="_invoke_agent_tool_calls", - metric_name="gen_ai.invoke_agent.tool_calls", - ), ) diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index 92fc33b050e..8711aa979ef 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/tests/unittests/telemetry/test_instrumentation.py @@ -15,7 +15,6 @@ # pylint: disable=protected-access import time -import types from unittest import mock from google.adk.telemetry import _instrumentation @@ -95,8 +94,8 @@ async def test_record_agent_invocation_tolerates_minimal_context(): """ agent = mock.MagicMock() agent.name = "test_agent" - # Context-like without `user_content` and without `session`. - bare_ctx = types.SimpleNamespace() + # Bare object without `user_content` and without `session`. + bare_ctx = object() with ( mock.patch.object(