From 4d09f33cd0a4072ccbb972a5759e995a5cfe0b45 Mon Sep 17 00:00:00 2001 From: Stephen Allen Date: Fri, 31 Jul 2026 16:52:03 -0500 Subject: [PATCH] feat(eval): add support for live workflows --- .../samples/live/live_workflow/README.md | 52 +++ .../samples/live/live_workflow/__init__.py | 15 + .../samples/live/live_workflow/agent.py | 120 ++++++ .../live_workflow/live_workflow.evalset.json | 20 + .../live/live_workflow/test_config.json | 109 +++++ .../adk/evaluation/evaluation_generator.py | 184 ++++++++- .../adk/flows/llm_flows/base_llm_flow.py | 12 + .../evaluation/test_evaluation_generator.py | 374 ++++++++++++++++++ 8 files changed, 866 insertions(+), 20 deletions(-) create mode 100644 contributing/samples/live/live_workflow/README.md create mode 100644 contributing/samples/live/live_workflow/__init__.py create mode 100644 contributing/samples/live/live_workflow/agent.py create mode 100644 contributing/samples/live/live_workflow/live_workflow.evalset.json create mode 100644 contributing/samples/live/live_workflow/test_config.json diff --git a/contributing/samples/live/live_workflow/README.md b/contributing/samples/live/live_workflow/README.md new file mode 100644 index 00000000000..1446e9a8d31 --- /dev/null +++ b/contributing/samples/live/live_workflow/README.md @@ -0,0 +1,52 @@ +# Live Workflow Sample + +## Overview + +This sample composes three short, single-purpose **live (voice) agents** into a graph-based workflow: + +1. `greeter_agent` — greets and confirms the caller's name. +1. `dob_verifier_agent` — captures and validates the caller's date of birth + (using the `validate_date_of_birth` tool). +1. `goals_agent` — once identity is verified, delivers the call goals and wraps + up the conversation. + +Each stage runs in `mode='task'` and hands a typed result to the next +(`GreeterOutput`, `DobOutput`). The stages are wired directly into the +workflow's `edges`, so the framework runs them in order. + +## Running the agent + +```bash +uv run adk web contributing/samples/live/live_workflow +``` + +Open the ADK web interface and start a Live Session with the agent. + +## Evaluating this agent + +`test_config.json` and `live_workflow.evalset.json` evaluate the workflow in +**live mode** with an `llm_audio` user simulator (each user turn is synthesized +to audio and streamed to the live agent). The eval scores response quality, +tool-use quality, and overall trajectory quality against custom rubrics across +the staged conversation. + +This sample uses the in-process, rubric-based LLM-as-judge metrics +(`rubric_based_final_response_quality_v1`, `rubric_based_tool_use_quality_v1`, +and `rubric_based_multi_turn_trajectory_quality_v1`), which support multi-agent +conversations. The first two are scored per turn; the trajectory metric judges +the whole conversation end-to-end. + +The eval case uses a `conversation_scenario`, so an LLM-simulated user adapts to +each stage of the workflow instead of following a fixed script. + +1. Install the eval extra: `uv pip install -e ".[eval]"`. +1. Add a `.env` in this directory with Vertex AI credentials (see + `live_bidi_streaming_single_agent/.env`). The project needs access to both + the Live API and Gemini TTS models. +1. Run the eval: + ```bash + uv run adk eval \ + contributing/samples/live/live_workflow \ + contributing/samples/live/live_workflow/live_workflow.evalset.json \ + --config_file_path contributing/samples/live/live_workflow/test_config.json + ``` diff --git a/contributing/samples/live/live_workflow/__init__.py b/contributing/samples/live/live_workflow/__init__.py new file mode 100644 index 00000000000..4015e47d6e4 --- /dev/null +++ b/contributing/samples/live/live_workflow/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import agent diff --git a/contributing/samples/live/live_workflow/agent.py b/contributing/samples/live/live_workflow/agent.py new file mode 100644 index 00000000000..1fb228035bd --- /dev/null +++ b/contributing/samples/live/live_workflow/agent.py @@ -0,0 +1,120 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""An example of how to build a graph-based live (voice) agent workflow.""" + +from google.adk.agents.llm_agent import Agent +from google.adk.tools.tool_context import ToolContext +from google.adk.workflow import START +from google.adk.workflow import Workflow +from pydantic import BaseModel +from pydantic import Field + +LIVE_MODEL = 'gemini-live-2.5-flash-native-audio' + + +# --- Typed handoffs between stages ----------------------------------------- +class GreeterOutput(BaseModel): + result: str = Field( + default='', description='The confirmed name of the person on the line.' + ) + + +class DobOutput(BaseModel): + result: str = Field( + default='', description='Identity verification result, e.g. "verified".' + ) + + +# --- Tool ------------------------------------------------------------------ +def validate_date_of_birth(dob: str, tool_context: ToolContext) -> dict: + """Validate a confirmed date of birth against records (mocked). + + Args: + dob: The patient's date of birth in YYYY-MM-DD format. + + Returns: + A dict with a ``match`` boolean. + """ + match = dob == '1985-07-12' # Mock record for the demo persona. + tool_context.state['dob_verified'] = match + return {'match': match} + + +# --- Stage 1: Greeting + identity ------------------------------------------ +greeter_agent = Agent( + model=LIVE_MODEL, + name='greeter_agent', + description='Greets on a recorded line and confirms the right person.', + mode='task', + output_schema=GreeterOutput, + instruction=""" + You are Sam, a friendly care-team assistant. + Greet the caller and confirm you are speaking with John Doe before + sharing anything else. Ask one question per turn. Once the name is + confirmed, briefly acknowledge it and complete your task, passing the + confirmed name as 'result'. + """, +) + + +# --- Stage 2: DOB verification --------------------------------------------- +dob_verifier_agent = Agent( + model=LIVE_MODEL, + name='dob_verifier_agent', + description='Captures and validates the date of birth before serving.', + mode='task', + output_schema=DobOutput, + tools=[validate_date_of_birth], + instruction=""" + Verify the caller's identity by date of birth. Ask for their date of + birth, read it back to confirm, then validate it with + `validate_date_of_birth` using YYYY-MM-DD format. Once it matches, let + the caller know their identity is verified and complete your task with + "verified". If it still does not match after two tries, complete your + task with "unverified". Ask one question per turn. + """, +) + + +# --- Stage 3: Conversation goals + ending ---------------------------------- +goals_agent = Agent( + model=LIVE_MODEL, + name='goals_agent', + description='Delivers the call goals once identity is verified.', + mode='task', + instruction=""" + Identity is already verified. As soon as it is your turn, proactively + tell the caller about their upcoming appointment on Tuesday, June 16th at + 3 PM with Dr. Example, and ask if they have any questions for the visit. + Do not wait to be asked. Answer any questions briefly, ask if there is + anything else, then wrap up warmly, end with "Goodbye.", and complete + your task. + """, +) + + +# --- The workflow: agents sequenced directly by edges ---------------------- +root_agent = Workflow( + name='live_workflow', + description=( + 'A Workflow of live voice agents: confirm the caller, verify their' + ' date of birth, then share the call details.' + ), + edges=[ + (START, greeter_agent), + (greeter_agent, dob_verifier_agent), + (dob_verifier_agent, goals_agent), + ], +) diff --git a/contributing/samples/live/live_workflow/live_workflow.evalset.json b/contributing/samples/live/live_workflow/live_workflow.evalset.json new file mode 100644 index 00000000000..02e96186176 --- /dev/null +++ b/contributing/samples/live/live_workflow/live_workflow.evalset.json @@ -0,0 +1,20 @@ +{ + "eval_set_id": "live_workflow", + "name": "live_workflow", + "description": "Live eval cases for the live workflow. Exercises the audio user simulator driving each stage: greeting/identity, DOB verification, and delivering the call goals.", + "eval_cases": [ + { + "eval_id": "verified_patient_scenario", + "conversation_scenario": { + "starting_prompt": "Hello?", + "conversation_plan": "You are John Doe. Confirm your name when greeted. When asked for your date of birth, give July 12th, 1985, and confirm it when read back. Listen to the appointment details, ask what you should bring to the visit, then say you have no other questions and let the call wrap up.", + "user_persona": "NOVICE" + }, + "session_input": { + "app_name": "live_workflow", + "user_id": "test_user_id", + "state": {} + } + } + ] +} diff --git a/contributing/samples/live/live_workflow/test_config.json b/contributing/samples/live/live_workflow/test_config.json new file mode 100644 index 00000000000..152cfd187b0 --- /dev/null +++ b/contributing/samples/live/live_workflow/test_config.json @@ -0,0 +1,109 @@ +{ + "criteria": { + "rubric_based_final_response_quality_v1": { + "threshold": 0.7, + "judge_model_options": { + "judge_model": "gemini-3.5-flash", + "num_samples": 1 + }, + "rubrics": [ + { + "rubric_id": "no_details_before_verification", + "rubric_content": { + "text_property": "If the agent shares appointment details in this turn, the caller's name and date of birth must already have been confirmed earlier in the conversation." + } + }, + { + "rubric_id": "states_appointment_when_asked", + "rubric_content": { + "text_property": "If the caller asks about their appointment in this turn, the agent states the appointment on Tuesday, June 16th at 3 PM with Dr. Example." + } + }, + { + "rubric_id": "offers_further_help", + "rubric_content": { + "text_property": "If the agent gives the caller information they requested in this turn, it also asks whether there is anything else it can help with." + } + }, + { + "rubric_id": "ends_with_goodbye", + "rubric_content": { + "text_property": "If the caller indicates they have no further questions, the agent ends the conversation warmly with a goodbye." + } + } + ] + }, + "rubric_based_tool_use_quality_v1": { + "threshold": 0.7, + "judge_model_options": { + "judge_model": "gemini-3.5-flash", + "num_samples": 1 + }, + "rubrics": [ + { + "rubric_id": "validates_confirmed_dob", + "rubric_content": { + "text_property": "After the caller confirms the date of birth that was read back to them, the agent calls validate_date_of_birth with the confirmed date in YYYY-MM-DD format." + } + } + ] + }, + "rubric_based_multi_turn_trajectory_quality_v1": { + "threshold": 0.7, + "judge_model_options": { + "judge_model": "gemini-3.5-flash", + "num_samples": 1 + }, + "rubrics": [ + { + "rubric_id": "verifies_identity_first", + "rubric_content": { + "text_property": "Across the call, the agent confirms the caller's name and validates their date of birth before disclosing any appointment details." + } + }, + { + "rubric_id": "staged_handoff_order", + "rubric_content": { + "text_property": "The conversation proceeds through the intended stages in order: greeting and name confirmation, then date-of-birth verification, then appointment delivery." + } + }, + { + "rubric_id": "calls_validation_tool", + "rubric_content": { + "text_property": "The agent calls validate_date_of_birth once, only after the caller confirms the read-back date, using YYYY-MM-DD format." + } + }, + { + "rubric_id": "delivers_appointment_and_answers", + "rubric_content": { + "text_property": "After identity is verified, the agent delivers the appointment on Tuesday, June 16th at 3 PM with Dr. Example and answers the caller's follow-up question." + } + }, + { + "rubric_id": "closes_conversation", + "rubric_content": { + "text_property": "The agent ends the call warmly with a goodbye once the caller has no further questions." + } + } + ] + } + }, + "live_model_config": { + "timeout_seconds": 300 + }, + "user_simulator_config": { + "type": "llm_audio", + "model": "gemini-3.5-flash", + "max_allowed_invocations": 10, + "audio_model": "gemini-3.1-flash-tts-preview", + "audio_model_configuration": { + "response_modalities": ["AUDIO"], + "speech_config": { + "voice_config": { + "prebuilt_voice_config": { "voice_name": "Kore" } + }, + "language_code": "en-US" + } + } + } +} diff --git a/src/google/adk/evaluation/evaluation_generator.py b/src/google/adk/evaluation/evaluation_generator.py index c6c397855bf..d7b9fff1eaf 100644 --- a/src/google/adk/evaluation/evaluation_generator.py +++ b/src/google/adk/evaluation/evaluation_generator.py @@ -30,6 +30,7 @@ from websockets.exceptions import ConnectionClosed from websockets.exceptions import ConnectionClosedOK +from ..agents.base_agent import BaseAgent from ..agents.callback_context import CallbackContext from ..agents.invocation_context import InvocationContext from ..agents.live_request_queue import LiveRequestQueue @@ -69,11 +70,22 @@ if TYPE_CHECKING: from types import TracebackType + from ..workflow import BaseNode + logger = logging.getLogger("google_adk." + __name__) _USER_AUTHOR = "user" _DEFAULT_AUTHOR = "agent" +# Function calls that end the agent and hand off instead of continuing the turn +# with a tool response, so their `turn_complete` is real. See +# `_consume_node_events`. +_TURN_ENDING_FUNCTION_CALLS = frozenset({ + "finish_task", + "transfer_to_agent", + "task_completed", +}) + # Chunk size for streaming audio blobs to the Live API. # See https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/live-api#technical-specifications _AUDIO_CHUNK_BYTES = 16000 @@ -205,19 +217,12 @@ async def __aenter__(self) -> _LiveSession: async def _consume_events(self) -> None: """Background task: consume events from run_live.""" try: - run_config = RunConfig( - streaming_mode=StreamingMode.BIDI, - response_modalities=["AUDIO"], - output_audio_transcription=types.AudioTranscriptionConfig(), - input_audio_transcription=types.AudioTranscriptionConfig(), - # Disable server-side voice-activity detection so turn boundaries are - # controlled explicitly via activity markers around the sent audio. - realtime_input_config=types.RealtimeInputConfig( - automatic_activity_detection=types.AutomaticActivityDetection( - disabled=True - ) - ), - ) + # Workflows have no _llm_flow/run_live; drive through Runner.run_live + if not isinstance(self.runner.agent, Agent): + await self._consume_node_events() + return + + run_config = self._live_run_config() invocation_context = self.runner._new_invocation_context_for_live( self.session, @@ -330,6 +335,149 @@ async def _consume_events(self) -> None: self.live_finished.set() self.turn_complete_event.set() # Unblock any waiters + async def _consume_node_events(self) -> None: + """Drives a non-Agent `BaseNode` (e.g. `Workflow`) via `Runner.run_live`.""" + from google.genai import errors + + # TODO: Remove once the live flow fires before/after_model_callback natively. + callback_context_by_author = await self._record_node_app_details() + + # Track tool-call turns so the user simulator isn't prompted before the + # agent has actually finished responding. + in_function_call_loop = False + + try: + async with Aclosing( + self.runner.run_live( + user_id=self.user_id, + session_id=self.session_id, + live_request_queue=self.live_request_queue, + run_config=self._live_run_config(), + ) + ) as agen: + async for event in agen: + event.invocation_id = self.current_invocation_id + callback_context = callback_context_by_author.get(event.author) + if callback_context is not None: + await self.runner.plugin_manager.run_after_model_callback( + callback_context=callback_context, + llm_response=event, + ) + await self.event_queue.put(event) + # Terminal/handoff calls end the agent; the next `turn_complete` is + # real, so they must not arm the guard. + if any( + fc.name not in _TURN_ENDING_FUNCTION_CALLS + for fc in event.get_function_calls() + ): + in_function_call_loop = True + if event.turn_complete and event.author != _USER_AUTHOR: + if not in_function_call_loop: + self.turn_complete_event.set() + else: + in_function_call_loop = False + except (ConnectionClosed, errors.APIError) as e: + # A clean session close ends the stream; keep the transcript so far + # instead of failing the eval case. + if not self._is_normal_closure(e): + raise + logger.info("Ignored WebSocket normal closure exception: %s", e) + + @staticmethod + def _live_run_config() -> RunConfig: + """Builds the live run config shared by the live drivers.""" + return RunConfig( + streaming_mode=StreamingMode.BIDI, + response_modalities=["AUDIO"], + output_audio_transcription=types.AudioTranscriptionConfig(), + input_audio_transcription=types.AudioTranscriptionConfig(), + # Disable server-side voice-activity detection so turn boundaries are + # controlled explicitly via activity markers around the sent audio. + realtime_input_config=types.RealtimeInputConfig( + automatic_activity_detection=types.AutomaticActivityDetection( + disabled=True + ) + ), + ) + + @staticmethod + def _is_normal_closure(exc: BaseException) -> bool: + """Reports whether an exception is a normal Live WebSocket closure (1000).""" + from google.genai import errors + + return isinstance(exc, ConnectionClosedOK) or ( + isinstance(exc, errors.APIError) and exc.code == 1000 + ) + + @staticmethod + async def _record_app_details_for_agent( + invocation_context: InvocationContext, + ) -> CallbackContext: + """Records the agent's live request so the autorater can score it. + + By default, live API calls do not fire before_model_callback, but the + plugins rely on it to capture the agent instructions and tool declarations + that the autorater needs. We run the callback manually here and return the + callback context so callers can replay after_model_callback per event. + + TODO: Remove once the live flow fires before/after_model_callback natively. + """ + agent = invocation_context.agent + llm_request = LlmRequest() + async with Aclosing( + agent._llm_flow._preprocess_async(invocation_context, llm_request) + ) as agen: + async for _ in agen: + pass + + callback_context = CallbackContext(invocation_context) + await invocation_context.plugin_manager.run_before_model_callback( + callback_context=callback_context, + llm_request=llm_request, + ) + return callback_context + + async def _record_node_app_details(self) -> dict[str, CallbackContext]: + """Records live requests for each agent in a node graph root. + + Multi-agent node roots serve several agents over one live stream, so we + record each agent's request up front and return a {author: callback_context} + map the caller uses to fire after_model_callback for each agent's events. A + failure for one agent is logged and skipped so it never aborts the eval run. + + TODO: Remove once the live flow fires before/after_model_callback natively. + """ + callback_context_by_author: dict[str, CallbackContext] = {} + + graph = getattr(self.runner.agent, "graph", None) + if graph is None: + return callback_context_by_author + + base_invocation_context = self.runner._new_invocation_context_for_live( + self.session, + live_request_queue=self.live_request_queue, + run_config=self._live_run_config(), + ) + + for node in graph.nodes: + if not isinstance(node, Agent): + continue + try: + invocation_context = base_invocation_context.model_copy( + update={"agent": node} + ) + callback_context_by_author[node.name] = ( + await self._record_app_details_for_agent(invocation_context) + ) + except Exception: # pylint: disable=broad-except + logger.warning( + "Failed to record app details for agent %s.", + node.name, + exc_info=True, + ) + + return callback_context_by_author + async def __aexit__( self, exc_type: type[BaseException] | None, @@ -355,11 +503,7 @@ async def __aexit__( # connection is closed with code 1000. Some client libraries may raise an # exception rather than handling it silently. We log this as INFO to # avoid false-positive error reports for expected behavior. - is_normal_closure = isinstance(e, ConnectionClosedOK) or ( - isinstance(e, errors.APIError) and e.code == 1000 - ) - - if is_normal_closure: + if self._is_normal_closure(e): logger.info("Ignored WebSocket normal closure exception: %s", e) else: raise @@ -550,7 +694,7 @@ async def _generate_inferences_for_single_user_invocation_live( @staticmethod async def _generate_inferences_from_root_agent_live( - root_agent: Agent, + root_agent: BaseAgent | BaseNode, user_simulator: UserSimulator, reset_func: Optional[Any] = None, initial_session: Optional[SessionInput] = None, @@ -667,7 +811,7 @@ async def _generate_inferences_from_root_agent_live( @staticmethod async def _generate_inferences_from_root_agent( - root_agent: Agent, + root_agent: BaseAgent | BaseNode, user_simulator: UserSimulator, reset_func: Optional[Any] = None, initial_session: Optional[SessionInput] = None, diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 954751a47bd..111bd95ba80 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -749,6 +749,12 @@ async def run_live( 'Connection closed (%s), reconnecting with session handle.', e ) continue + # No resumption handle + normal (1000) close = the model ended the + # session cleanly; end the stream instead of erroring so live nodes + # finish normally. + if isinstance(e, ConnectionClosedOK): + logger.info('Connection closed normally: %s.', e) + return logger.error('Connection closed: %s.', e) raise except errors.APIError as e: @@ -763,6 +769,12 @@ async def run_live( 'Connection lost (%s), reconnecting with session handle.', e ) continue + # No resumption handle + normal (1000) close = the model ended the + # session cleanly; end the stream instead of erroring so live nodes + # finish normally. + if e.code == 1000: + logger.info('Live session closed normally: %s.', e) + return logger.error('APIError in live flow: %s', e) raise diff --git a/tests/unittests/evaluation/test_evaluation_generator.py b/tests/unittests/evaluation/test_evaluation_generator.py index 8f01f767a75..5f7a249f44d 100644 --- a/tests/unittests/evaluation/test_evaluation_generator.py +++ b/tests/unittests/evaluation/test_evaluation_generator.py @@ -1340,6 +1340,380 @@ async def mock_run_live(*args, **kwargs): assert called_after_args.kwargs["llm_response"] == mock_event +class TestLiveSessionNodeRouting: + """Verifies non-Agent BaseNode roots are driven via `Runner.run_live`.""" + + @pytest.mark.asyncio + async def test_workflow_root_uses_runner_run_live(self, mocker): + from google.adk.workflow import Workflow + + # A Workflow has no `_llm_flow`/`run_live`; the session must fall back to + # `Runner.run_live`, which handles node scheduling and function calls. + mock_workflow = mocker.MagicMock(spec=Workflow) + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + + mock_event = Event( + author="agent", + content=types.Content(parts=[types.Part(text="Hi")]), + invocation_id="unused", + turn_complete=True, + ) + + async def mock_run_live(*args, **kwargs): + yield mock_event + + mock_runner.run_live.return_value = mock_run_live() + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="test_user", + session_id="test_session", + ) + + await live_session._consume_events() + + # The Agent-only driver must not be touched for a Workflow root. + mock_runner._new_invocation_context_for_live.assert_not_called() + mock_runner.run_live.assert_called_once() + call_kwargs = mock_runner.run_live.call_args.kwargs + assert call_kwargs["user_id"] == "test_user" + assert call_kwargs["session_id"] == "test_session" + + # The event is re-stamped with the turn's invocation id and enqueued, and + # the agent's turn-complete releases the waiter. + queued_event = await live_session.event_queue.get() + assert queued_event.invocation_id == live_session.current_invocation_id + assert live_session.turn_complete_event.is_set() + assert live_session.live_finished.is_set() + + @pytest.mark.asyncio + async def test_workflow_root_swallows_tool_call_turn_complete(self, mocker): + from google.adk.workflow import Workflow + + # The intermediate `turn_complete` that accompanies a tool call must not + # release the waiter; only the true end-of-turn after the model continues. + mock_workflow = mocker.MagicMock(spec=Workflow) + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + + # A tool call and its `turn_complete` arrive as separate events. + fc_event = Event( + author="dob_verifier_agent", + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name="validate_date_of_birth", + args={"dob": "1985-07-12"}, + ) + ) + ] + ), + invocation_id="unused", + ) + intermediate_turn_complete = Event( + author="dob_verifier_agent", + invocation_id="unused", + turn_complete=True, + ) + # The real end-of-turn after the model continues. + final_event = Event( + author="dob_verifier_agent", + content=types.Content( + parts=[types.Part(text="Your identity is verified.")] + ), + invocation_id="unused", + turn_complete=True, + ) + + async def mock_run_live(*args, **kwargs): + yield fc_event + assert not live_session.turn_complete_event.is_set() + yield intermediate_turn_complete + assert not live_session.turn_complete_event.is_set() + yield final_event + + mock_runner.run_live.return_value = mock_run_live() + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="test_user", + session_id="test_session", + ) + mocker.patch.object( + live_session, + "_record_node_app_details", + new=mocker.AsyncMock(return_value={}), + ) + + await live_session._consume_node_events() + + # The waiter is released only on the true end-of-turn. + assert live_session.turn_complete_event.is_set() + + @pytest.mark.asyncio + async def test_workflow_root_finish_task_does_not_swallow_next_turn( + self, mocker + ): + from google.adk.workflow import Workflow + + # `finish_task` hands off to the next agent, so its following `turn_complete` + # (the next agent's question) is a real end-of-turn and must not be swallowed + # -- otherwise the harness hangs. + mock_workflow = mocker.MagicMock(spec=Workflow) + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + + finish_task_event = Event( + author="greeter_agent", + content=types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name="finish_task", args={"result": "John Doe"} + ) + ) + ] + ), + invocation_id="unused", + ) + # The next agent takes over and asks its first question, ending the turn. + next_agent_turn = Event( + author="dob_verifier_agent", + content=types.Content( + parts=[types.Part(text="What is your date of birth?")] + ), + invocation_id="unused", + turn_complete=True, + ) + + async def mock_run_live(*args, **kwargs): + yield finish_task_event + yield next_agent_turn + + mock_runner.run_live.return_value = mock_run_live() + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="test_user", + session_id="test_session", + ) + mocker.patch.object( + live_session, + "_record_node_app_details", + new=mocker.AsyncMock(return_value={}), + ) + + await live_session._consume_node_events() + + # The next agent's turn_complete must release the waiter. + assert live_session.turn_complete_event.is_set() + + @pytest.mark.asyncio + async def test_workflow_root_tolerates_normal_ws_closure(self, mocker): + from google.adk.workflow import Workflow + from google.genai import errors + + # The node runner unwraps the end-of-session `1000` closure and re-raises it + # from `run_live`; the session must treat it as a normal end of stream + # rather than aborting the eval case. + mock_workflow = mocker.MagicMock(spec=Workflow) + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + + async def mock_run_live(*args, **kwargs): + raise errors.APIError(1000, {}, None) + yield # pragma: no cover - makes this an async generator + + mock_runner.run_live.return_value = mock_run_live() + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="test_user", + session_id="test_session", + ) + + # Must not raise; the normal closure is swallowed. + await live_session._consume_events() + + # The `finally` still flags the stream finished and unblocks waiters. + assert live_session.live_finished.is_set() + assert live_session.turn_complete_event.is_set() + + @pytest.mark.asyncio + async def test_workflow_root_reraises_abnormal_ws_closure(self, mocker): + from google.adk.workflow import Workflow + from google.genai import errors + + # A non-1000 closure is a real failure and must propagate. + mock_workflow = mocker.MagicMock(spec=Workflow) + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + + async def mock_run_live(*args, **kwargs): + raise errors.APIError(1011, {}, None) + yield # pragma: no cover - makes this an async generator + + mock_runner.run_live.return_value = mock_run_live() + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="test_user", + session_id="test_session", + ) + + with pytest.raises(errors.APIError): + await live_session._consume_events() + + @pytest.mark.asyncio + async def test_workflow_fires_after_model_callback_by_author(self, mocker): + from google.adk.workflow import Workflow + + mock_workflow = mocker.MagicMock(spec=Workflow) + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + mock_runner.plugin_manager.run_after_model_callback = mocker.AsyncMock() + + greeter_event = Event(author="greeter", invocation_id="x") + other_event = Event(author="unrecorded_agent", invocation_id="x") + + async def mock_run_live(*args, **kwargs): + yield greeter_event + yield other_event + + mock_runner.run_live.return_value = mock_run_live() + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="u", + session_id="s", + ) + greeter_context = mocker.MagicMock() + mocker.patch.object( + live_session, + "_record_node_app_details", + new=mocker.AsyncMock(return_value={"greeter": greeter_context}), + ) + + await live_session._consume_node_events() + + # Only the recorded author's events replay after_model_callback, and with + # that author's callback context. + mock_runner.plugin_manager.run_after_model_callback.assert_called_once_with( + callback_context=greeter_context, llm_response=greeter_event + ) + + +class TestLiveSessionNodeAppDetails: + """Verifies the stop-gap that records autorater app details for node roots. + + ``_record_app_details_for_agent`` (the per-agent preprocess + callback firing) + is exercised by ``TestLiveSessionCallbacks``; here we focus on the node graph + enumeration/mapping/skip logic in ``_record_node_app_details``. + """ + + def _make_agent_node(self, mocker, name: str): + from google.adk.agents.llm_agent import Agent + + agent = mocker.MagicMock(spec=Agent) + agent.name = name + return agent + + def _make_runner(self, mocker, nodes): + mock_workflow = mocker.MagicMock() + mock_workflow.graph.nodes = nodes + + mock_runner = mocker.MagicMock() + mock_runner.agent = mock_workflow + # `model_copy` carries the target agent onto the per-agent context so + # `_record_app_details_for_agent` sees the right agent. + base_ic = mock_runner._new_invocation_context_for_live.return_value + base_ic.model_copy.side_effect = lambda update: mocker.MagicMock( + agent=update["agent"] + ) + return mock_runner + + @pytest.mark.asyncio + async def test_record_node_app_details_maps_each_agent(self, mocker): + greeter = self._make_agent_node(mocker, "greeter") + verifier = self._make_agent_node(mocker, "verifier") + mock_runner = self._make_runner(mocker, [greeter, verifier]) + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="u", + session_id="s", + ) + + # Return a distinct callback context per recorded agent. + greeter_context = mocker.sentinel.greeter_context + verifier_context = mocker.sentinel.verifier_context + mocker.patch.object( + live_session, + "_record_app_details_for_agent", + new=mocker.AsyncMock(side_effect=[greeter_context, verifier_context]), + ) + + callback_context_by_author = await live_session._record_node_app_details() + + assert callback_context_by_author == { + "greeter": greeter_context, + "verifier": verifier_context, + } + + @pytest.mark.asyncio + async def test_record_node_app_details_no_graph_returns_empty(self, mocker): + mock_runner = mocker.MagicMock() + mock_runner.agent = object() # no `graph` attribute + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="u", + session_id="s", + ) + + assert await live_session._record_node_app_details() == {} + + @pytest.mark.asyncio + async def test_record_node_app_details_skips_failed_agent(self, mocker): + good = self._make_agent_node(mocker, "good") + bad = self._make_agent_node(mocker, "bad") + mock_runner = self._make_runner(mocker, [bad, good]) + + live_session = _LiveSession( + runner=mock_runner, + session=mocker.MagicMock(), + user_id="u", + session_id="s", + ) + + good_context = mocker.sentinel.good_context + + async def record(ic): + if ic.agent.name == "bad": + raise RuntimeError("boom") + return good_context + + mocker.patch.object( + live_session, + "_record_app_details_for_agent", + new=mocker.AsyncMock(side_effect=record), + ) + + # The failing agent is skipped; the healthy one is still recorded. + callback_context_by_author = await live_session._record_node_app_details() + assert callback_context_by_author == {"good": good_context} + + def test_convert_events_preserves_tool_calls_when_skip_summarization(): """Regression test for tool calls dropped from invocation_events.