From 61bec007e2eb0d661e286ce772e6d5cc1bf1d9fa Mon Sep 17 00:00:00 2001 From: Lin-Nikaido Date: Mon, 24 Aug 2026 21:10:04 +0900 Subject: [PATCH 1/2] feat --- src/google/adk/flows/llm_flows/functions.py | 154 +++++++ .../test_functions_streaming_results.py | 424 ++++++++++++++++++ 2 files changed, 578 insertions(+) create mode 100644 tests/unittests/flows/llm_flows/test_functions_streaming_results.py diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index d9943bcdb8e..841951324a6 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -20,6 +20,7 @@ import base64 import binascii from concurrent.futures import ThreadPoolExecutor +import contextlib import contextvars import copy import inspect @@ -27,6 +28,7 @@ import logging import threading from typing import Any +from typing import AsyncGenerator from typing import Callable from typing import cast from typing import Dict @@ -637,6 +639,12 @@ async def _run_with_trace() -> Event | None: function_response = await __call_tool_async( tool, args=function_args, tool_context=tool_context ) + function_response = await _drain_tool_result_stream( + function_response, + tool=tool, + tool_context=tool_context, + invocation_context=invocation_context, + ) except Exception as tool_error: error_response = await _run_on_tool_error_callbacks( tool=tool, @@ -1113,6 +1121,152 @@ async def _emit_streaming_tool_event( ) +async def _iter_tool_yields( + generator: object, +) -> AsyncGenerator[object, None]: + """Yields every value a generator tool produces, async or sync. + + Args: + generator: The async or sync generator the tool returned. + + Yields: + Each value the tool yields, in order. + """ + if inspect.isasyncgen(generator): + async with Aclosing(generator) as agen: + async for value in agen: + yield value + return + + with contextlib.closing(cast(Any, generator)) as gen: + for value in gen: + yield value + + +async def _emit_intermediate_tool_result( + function_result: object, + *, + tool: BaseTool, + tool_context: ToolContext, + invocation_context: InvocationContext, +) -> None: + """Streams one intermediate result of a generator tool to the client. + + Args: + function_result: The value the tool yielded. + tool: The tool that yielded it. + tool_context: The context of the call, for its function call id. + invocation_context: The invocation to enqueue on. + """ + remaining_result, function_response_parts = _extract_multimodal_parts( + function_result + ) + content = _build_function_response_content( + tool, + _normalize_tool_result(remaining_result), + tool_context.function_call_id, + function_response_parts, + ) + for part in content.parts or []: + if part.function_response is not None: + part.function_response.will_continue = True + + await invocation_context._enqueue_event( + Event( + content=content, + author=_require_agent_name(invocation_context), + invocation_id=invocation_context.invocation_id, + branch=invocation_context.branch, + partial=True, + actions=EventActions(), + ) + ) + + +async def _drain_tool_result_stream( + function_response: object, + *, + tool: BaseTool, + tool_context: ToolContext, + invocation_context: InvocationContext, +) -> object: + """Streams a generator tool's progress and returns the result for the model. + + A tool that yields instead of returning reports progress. Every value it + yields but the last is streamed to the client as a partial event and is + never shown to the model; the last one is the tool result, and is the only + one the model sees. A yielded ``Event`` is a user-facing message rather than + a result, so it is streamed and then does not compete to be the last value. + + Args: + function_response: The tool's return value. Anything that is not a + generator is returned unchanged. + tool: The tool being called. + tool_context: The context of the call. + invocation_context: The invocation to stream on. + + Returns: + The last non-Event value the tool yielded, or None if it yielded none. + """ + if not inspect.isasyncgen(function_response) and not inspect.isgenerator( + function_response + ): + return function_response + + can_stream = invocation_context._event_queue is not None + if not can_stream: + logger.warning( + 'Tool `%s` yields intermediate results, but this invocation has no' + ' event queue to deliver them on, so only its last result is' + ' reported. Streaming intermediate results requires running the agent' + ' through Runner.', + tool.name, + ) + + result: object = None + has_result = False + try: + async with Aclosing(_iter_tool_yields(function_response)) as stream: + async for value in stream: + if isinstance(value, Event): + if can_stream: + await _emit_streaming_tool_event( + value, + tool=tool, + tool_context=tool_context, + invocation_context=invocation_context, + ) + continue + if has_result and can_stream: + await _emit_intermediate_tool_result( + result, + tool=tool, + tool_context=tool_context, + invocation_context=invocation_context, + ) + result = value + has_result = True + + except Exception: + if has_result and can_stream: + try: + await _emit_intermediate_tool_result( + result, + tool=tool, + tool_context=tool_context, + invocation_context=invocation_context, + ) + except Exception: # pylint: disable=broad-except + # Never masks the tool's own failure, which is what the caller has to + # see, with a failure to report progress about it. + logger.exception( + 'Failed to report the last progress of failing tool `%s`.', + tool.name, + ) + raise + return result + + async def _process_function_live_helper( tool: BaseTool, tool_context: ToolContext, diff --git a/tests/unittests/flows/llm_flows/test_functions_streaming_results.py b/tests/unittests/flows/llm_flows/test_functions_streaming_results.py new file mode 100644 index 00000000000..d73f060feaf --- /dev/null +++ b/tests/unittests/flows/llm_flows/test_functions_streaming_results.py @@ -0,0 +1,424 @@ +# 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. + +"""Tests for tools that yield intermediate results in non-live runs.""" + +import asyncio +import json +from typing import Any +from typing import AsyncGenerator +from typing import Generator + +from google.adk.agents.llm_agent import Agent +from google.adk.events.event import Event +from google.adk.tools.tool_context import ToolContext +from google.genai import types +import pytest + +from ... import testing_utils + + +def _function_responses(event: Event) -> list[types.FunctionResponse]: + """Returns the function responses an event carries.""" + if not event.content or not event.content.parts: + return [] + return [ + part.function_response + for part in event.content.parts + if part.function_response is not None + ] + + +def _tool_progress(events: list[Event]) -> list[tuple[bool, Any]]: + """Summarizes every function response as (is intermediate, payload).""" + return [ + (bool(event.partial), function_response.response) + for event in events + for function_response in _function_responses(event) + ] + + +@pytest.mark.asyncio +async def test_async_generator_tool_streams_every_yield_but_the_last(): + function_call = types.Part.from_function_call( + name='search', args={'q': 'adk'} + ) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search(q: str) -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress', 'message': f'searching {q}'} + yield {'status': 'inProgress', 'message': 'reading pages'} + yield {'status': 'ok', 'result': ['page1', 'page2']} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert _tool_progress(events) == [ + (True, {'status': 'inProgress', 'message': 'searching adk'}), + (True, {'status': 'inProgress', 'message': 'reading pages'}), + (False, {'status': 'ok', 'result': ['page1', 'page2']}), + ] + + +@pytest.mark.asyncio +async def test_only_intermediate_results_are_marked_will_continue(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + will_continue = [ + function_response.will_continue + for event in events + for function_response in _function_responses(event) + ] + assert will_continue == [True, None] + + +@pytest.mark.asyncio +async def test_intermediate_results_reach_an_sse_client_as_will_continue(): + """Pins the shape a client reads off the wire, not just the Event in memory.""" + function_call = types.Part.from_function_call(name='search', args={}) + function_call.function_call.id = 'call-1' + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + intermediate = next(event for event in events if event.partial) + # The same serialization the SSE endpoint applies. + payload = json.loads( + intermediate.model_dump_json(exclude_none=True, by_alias=True) + ) + assert payload['partial'] is True + assert payload['content']['parts'][0]['functionResponse'] == { + 'willContinue': True, + 'id': 'call-1', + 'name': 'search', + 'response': {'status': 'inProgress'}, + } + + +@pytest.mark.asyncio +async def test_intermediate_results_carry_the_call_id_and_tool_name(): + function_call = types.Part.from_function_call(name='search', args={}) + function_call.function_call.id = 'call-1' + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + addressing = [ + (function_response.id, function_response.name) + for event in events + for function_response in _function_responses(event) + ] + assert addressing == [('call-1', 'search'), ('call-1', 'search')] + + +@pytest.mark.asyncio +async def test_intermediate_results_are_not_appended_to_the_session(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + await runner.run_async('test') + + assert _tool_progress(runner.session.events) == [(False, {'status': 'ok'})] + + +@pytest.mark.asyncio +async def test_intermediate_results_are_not_sent_to_the_model(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + await runner.run_async('test') + + assert testing_utils.simplify_contents(mock_model.requests[1].contents) == [ + ('user', 'test'), + ('model', types.Part.from_function_call(name='search', args={})), + ( + 'user', + types.Part.from_function_response( + name='search', response={'status': 'ok'} + ), + ), + ] + + +@pytest.mark.asyncio +async def test_sync_generator_tool_streams_every_yield_but_the_last(): + function_call = types.Part.from_function_call(name='count', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + def count() -> Generator[dict[str, Any], None, None]: + yield {'status': 'inProgress', 'done': 1} + yield {'status': 'inProgress', 'done': 2} + yield {'status': 'ok', 'done': 3} + + agent = Agent(name='root_agent', model=mock_model, tools=[count]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert _tool_progress(events) == [ + (True, {'status': 'inProgress', 'done': 1}), + (True, {'status': 'inProgress', 'done': 2}), + (False, {'status': 'ok', 'done': 3}), + ] + + +@pytest.mark.asyncio +async def test_non_dict_yields_are_wrapped_like_a_returned_value(): + function_call = types.Part.from_function_call(name='count', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def count() -> AsyncGenerator[int, None]: + yield 1 + yield 2 + + agent = Agent(name='root_agent', model=mock_model, tools=[count]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert _tool_progress(events) == [ + (True, {'result': 1}), + (False, {'result': 2}), + ] + + +@pytest.mark.asyncio +async def test_tool_yielding_once_reports_no_intermediate_result(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert _tool_progress(events) == [(False, {'status': 'ok'})] + + +@pytest.mark.asyncio +async def test_tool_yielding_nothing_reports_a_null_result(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + if False: + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert _tool_progress(events) == [(False, {'result': None})] + + +@pytest.mark.asyncio +async def test_yielded_event_is_delivered_as_a_user_message(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[Any, None]: + yield Event(message='looking it up') + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + # The Event is a message for the user, so it is not a function response and + # does not displace the result the model sees. + assert _tool_progress(events) == [(False, {'status': 'ok'})] + messages = [ + event + for event in events + if event.content and event.content.parts and event.content.parts[0].text + ] + assert [ + (event.content.role, event.content.parts[0].text) for event in messages + ] == [('user', 'looking it up'), ('model', 'done')] + assert messages[0].branch.startswith('search@') + + +@pytest.mark.asyncio +async def test_failure_after_the_first_yield_reaches_on_tool_error(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + errors = [] + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + raise ValueError('no results') + + def on_tool_error(*, tool, args, tool_context, error): + errors.append(error) + return {'status': 'error', 'message': str(error)} + + agent = Agent( + name='root_agent', + model=mock_model, + tools=[search], + on_tool_error_callback=on_tool_error, + ) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert [str(error) for error in errors] == ['no results'] + assert _tool_progress(events) == [ + (True, {'status': 'inProgress'}), + (False, {'status': 'error', 'message': 'no results'}), + ] + + +@pytest.mark.asyncio +async def test_unhandled_failure_after_the_first_yield_propagates(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + raise ValueError('no results') + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + with pytest.raises(ValueError, match='no results'): + await runner.run_async('test') + + +@pytest.mark.asyncio +async def test_parallel_generator_tools_stream_independently(): + function_calls = [ + types.Part.from_function_call(name='fast', args={}), + types.Part.from_function_call(name='slow', args={}), + ] + mock_model = testing_utils.MockModel.create( + responses=[function_calls, 'done'] + ) + + async def fast() -> AsyncGenerator[dict[str, Any], None]: + yield {'tool': 'fast', 'status': 'inProgress'} + yield {'tool': 'fast', 'status': 'ok'} + + async def slow() -> AsyncGenerator[dict[str, Any], None]: + yield {'tool': 'slow', 'status': 'inProgress'} + # Yields the event loop so that `fast` finishes first, proving the two + # calls are drained concurrently rather than one after the other. + await asyncio.sleep(0.05) + yield {'tool': 'slow', 'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[fast, slow]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + intermediate = [ + payload for partial, payload in _tool_progress(events) if partial + ] + final = [ + payload for partial, payload in _tool_progress(events) if not partial + ] + assert sorted(payload['tool'] for payload in intermediate) == ['fast', 'slow'] + # Both calls answer in one merged non-partial event, as parallel calls always + # have. + assert len(final) == 2 + assert all(payload['status'] == 'ok' for payload in final) + + +@pytest.mark.asyncio +async def test_tool_context_is_still_injected_into_a_generator_tool(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + seen = [] + + async def search( + tool_context: ToolContext | None = None, + ) -> AsyncGenerator[dict[str, Any], None]: + seen.append(tool_context.function_call_id) + yield {'status': 'inProgress'} + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + await runner.run_async('test') + + assert len(seen) == 1 and seen[0] + + +@pytest.mark.asyncio +async def test_state_a_generator_tool_sets_is_applied_once_at_the_end(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search( + tool_context: ToolContext | None = None, + ) -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + tool_context.state['hits'] = 2 + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + # An intermediate event is never applied, so it must not carry the delta. + assert [event.actions.state_delta for event in events if event.partial] == [ + {} + ] + assert runner.session.state['hits'] == 2 + + +@pytest.mark.asyncio +async def test_returning_a_plain_value_is_unchanged(): + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + + async def search() -> dict[str, Any]: + return {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + events = await runner.run_async('test') + + assert _tool_progress(events) == [(False, {'status': 'ok'})] From 3842ed6ff9d4bfa2acccf1b88882f04e3a060eac Mon Sep 17 00:00:00 2001 From: Lin-Nikaido Date: Fri, 11 Sep 2026 17:14:17 +0900 Subject: [PATCH 2/2] fix: the yields event is late one iteration --- src/google/adk/flows/llm_flows/functions.py | 61 ++++++--------- .../test_functions_streaming_results.py | 76 ++++++++++++++++--- 2 files changed, 90 insertions(+), 47 deletions(-) diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 841951324a6..9b9052e800c 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -1150,7 +1150,11 @@ async def _emit_intermediate_tool_result( tool_context: ToolContext, invocation_context: InvocationContext, ) -> None: - """Streams one intermediate result of a generator tool to the client. + """Streams one value a generator tool yielded to the client. + + Marked ``will_continue``, because a further FunctionResponse for the same + call always follows: either the tool's next yield, or the result event that + answers the call for the model. Args: function_result: The value the tool yielded. @@ -1193,10 +1197,15 @@ async def _drain_tool_result_stream( """Streams a generator tool's progress and returns the result for the model. A tool that yields instead of returning reports progress. Every value it - yields but the last is streamed to the client as a partial event and is - never shown to the model; the last one is the tool result, and is the only - one the model sees. A yielded ``Event`` is a user-facing message rather than - a result, so it is streamed and then does not compete to be the last value. + yields is streamed to the client as a partial event the moment it is + yielded, and none of them is shown to the model; the last one is also the + tool result, and is the only one the model sees. A yielded ``Event`` is a + user-facing message rather than a result, so it is streamed and then does + not compete to be the last value. + + The last value is delivered twice, once as progress and once as the + result; ``will_continue`` tells the two apart, and is true on the progress + copy because a further FunctionResponse for that call does follow. Args: function_response: The tool's return value. Anything that is not a @@ -1224,46 +1233,26 @@ async def _drain_tool_result_stream( ) result: object = None - has_result = False - try: - async with Aclosing(_iter_tool_yields(function_response)) as stream: - async for value in stream: - if isinstance(value, Event): - if can_stream: - await _emit_streaming_tool_event( - value, - tool=tool, - tool_context=tool_context, - invocation_context=invocation_context, - ) - continue - if has_result and can_stream: - await _emit_intermediate_tool_result( - result, + async with Aclosing(_iter_tool_yields(function_response)) as stream: + async for value in stream: + if isinstance(value, Event): + if can_stream: + await _emit_streaming_tool_event( + value, tool=tool, tool_context=tool_context, invocation_context=invocation_context, ) - result = value - has_result = True - - except Exception: - if has_result and can_stream: - try: + continue + result = value + if can_stream: await _emit_intermediate_tool_result( - result, + value, tool=tool, tool_context=tool_context, invocation_context=invocation_context, ) - except Exception: # pylint: disable=broad-except - # Never masks the tool's own failure, which is what the caller has to - # see, with a failure to report progress about it. - logger.exception( - 'Failed to report the last progress of failing tool `%s`.', - tool.name, - ) - raise + return result diff --git a/tests/unittests/flows/llm_flows/test_functions_streaming_results.py b/tests/unittests/flows/llm_flows/test_functions_streaming_results.py index d73f060feaf..fdcf899e1cc 100644 --- a/tests/unittests/flows/llm_flows/test_functions_streaming_results.py +++ b/tests/unittests/flows/llm_flows/test_functions_streaming_results.py @@ -50,7 +50,7 @@ def _tool_progress(events: list[Event]) -> list[tuple[bool, Any]]: @pytest.mark.asyncio -async def test_async_generator_tool_streams_every_yield_but_the_last(): +async def test_async_generator_tool_streams_every_yield(): function_call = types.Part.from_function_call( name='search', args={'q': 'adk'} ) @@ -65,15 +65,18 @@ async def search(q: str) -> AsyncGenerator[dict[str, Any], None]: runner = testing_utils.InMemoryRunner(agent) events = await runner.run_async('test') + # The last yield is both the last report and the result, so it appears + # twice: once streamed as it was yielded, once as the answer to the call. assert _tool_progress(events) == [ (True, {'status': 'inProgress', 'message': 'searching adk'}), (True, {'status': 'inProgress', 'message': 'reading pages'}), + (True, {'status': 'ok', 'result': ['page1', 'page2']}), (False, {'status': 'ok', 'result': ['page1', 'page2']}), ] @pytest.mark.asyncio -async def test_only_intermediate_results_are_marked_will_continue(): +async def test_only_the_answer_to_the_call_is_not_marked_will_continue(): function_call = types.Part.from_function_call(name='search', args={}) mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) @@ -90,7 +93,9 @@ async def search() -> AsyncGenerator[dict[str, Any], None]: for event in events for function_response in _function_responses(event) ] - assert will_continue == [True, None] + # Every streamed value is followed by a further FunctionResponse for the + # same call, so only the one answering the call is not marked. + assert will_continue == [True, True, None] @pytest.mark.asyncio @@ -141,7 +146,7 @@ async def search() -> AsyncGenerator[dict[str, Any], None]: for event in events for function_response in _function_responses(event) ] - assert addressing == [('call-1', 'search'), ('call-1', 'search')] + assert addressing == [('call-1', 'search')] * 3 @pytest.mark.asyncio @@ -186,7 +191,7 @@ async def search() -> AsyncGenerator[dict[str, Any], None]: @pytest.mark.asyncio -async def test_sync_generator_tool_streams_every_yield_but_the_last(): +async def test_sync_generator_tool_streams_every_yield(): function_call = types.Part.from_function_call(name='count', args={}) mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) @@ -202,6 +207,7 @@ def count() -> Generator[dict[str, Any], None, None]: assert _tool_progress(events) == [ (True, {'status': 'inProgress', 'done': 1}), (True, {'status': 'inProgress', 'done': 2}), + (True, {'status': 'ok', 'done': 3}), (False, {'status': 'ok', 'done': 3}), ] @@ -221,12 +227,13 @@ async def count() -> AsyncGenerator[int, None]: assert _tool_progress(events) == [ (True, {'result': 1}), + (True, {'result': 2}), (False, {'result': 2}), ] @pytest.mark.asyncio -async def test_tool_yielding_once_reports_no_intermediate_result(): +async def test_tool_yielding_once_streams_that_yield_and_answers_with_it(): function_call = types.Part.from_function_call(name='search', args={}) mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) @@ -237,7 +244,10 @@ async def search() -> AsyncGenerator[dict[str, Any], None]: runner = testing_utils.InMemoryRunner(agent) events = await runner.run_async('test') - assert _tool_progress(events) == [(False, {'status': 'ok'})] + assert _tool_progress(events) == [ + (True, {'status': 'ok'}), + (False, {'status': 'ok'}), + ] @pytest.mark.asyncio @@ -271,7 +281,10 @@ async def search() -> AsyncGenerator[Any, None]: # The Event is a message for the user, so it is not a function response and # does not displace the result the model sees. - assert _tool_progress(events) == [(False, {'status': 'ok'})] + assert _tool_progress(events) == [ + (True, {'status': 'ok'}), + (False, {'status': 'ok'}), + ] messages = [ event for event in events @@ -283,6 +296,41 @@ async def search() -> AsyncGenerator[Any, None]: assert messages[0].branch.startswith('search@') +@pytest.mark.asyncio +async def test_a_yield_is_reported_before_the_tool_takes_its_next_step(): + """A progress report reaches the client while the work it describes still runs. + + The tool blocks after its first yield until the consumer has seen that yield + reported. Holding a value back until the following one arrives -- which is + the only way to know it was not the last -- deadlocks here instead, and the + `wait_for` turns that regression into a failure rather than a hang. + """ + function_call = types.Part.from_function_call(name='search', args={}) + mock_model = testing_utils.MockModel.create(responses=[function_call, 'done']) + reported = asyncio.Event() + + async def search() -> AsyncGenerator[dict[str, Any], None]: + yield {'status': 'inProgress'} + await asyncio.wait_for(reported.wait(), timeout=10) + yield {'status': 'ok'} + + agent = Agent(name='root_agent', model=mock_model, tools=[search]) + runner = testing_utils.InMemoryRunner(agent) + session = runner.session + + seen = [] + async for event in runner.runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=testing_utils.get_user_content('test'), + ): + seen.extend(_tool_progress([event])) + if seen: + reported.set() + + assert seen[0] == (True, {'status': 'inProgress'}) + + @pytest.mark.asyncio async def test_failure_after_the_first_yield_reaches_on_tool_error(): function_call = types.Part.from_function_call(name='search', args={}) @@ -359,7 +407,12 @@ async def slow() -> AsyncGenerator[dict[str, Any], None]: final = [ payload for partial, payload in _tool_progress(events) if not partial ] - assert sorted(payload['tool'] for payload in intermediate) == ['fast', 'slow'] + assert sorted(payload['tool'] for payload in intermediate) == [ + 'fast', + 'fast', + 'slow', + 'slow', + ] # Both calls answer in one merged non-partial event, as parallel calls always # have. assert len(final) == 2 @@ -402,9 +455,10 @@ async def search( runner = testing_utils.InMemoryRunner(agent) events = await runner.run_async('test') - # An intermediate event is never applied, so it must not carry the delta. + # A streamed event is never applied, so it must not carry the delta. assert [event.actions.state_delta for event in events if event.partial] == [ - {} + {}, + {}, ] assert runner.session.state['hits'] == 2