fix(cli): stream_reasoning_engine raises StopIteration RuntimeError on sync generators

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

## Link to Issue or Description of Change
Closes : #6093
  **Problem:**
  On Agent Engine deployments served by the ADK API server, every call to the
  `/api/stream_reasoning_engine` route with a synchronous streaming `class_method`
  (e.g. `stream_query`) ends with `RuntimeError: coroutine raised StopIteration`
  after the last chunk is streamed.

  The cause is the sync-to-async adapter `_aiter_from_iter` in
  `src/google/adk/cli/fast_api.py` (lines 916–922 in v2.2.0):

      async def _aiter_from_iter(iterator):
        while True:
          try:
            chunk = await run_in_threadpool(next, iterator)
            yield chunk
          except StopIteration:
            break

  The `except StopIteration` is unreachable. When the iterator is exhausted,
  `next()` raises `StopIteration` inside the worker thread, anyio sets it on a
  future, and it propagates out of the `run_in_threadpool` coroutine frame.
  Python (PEP 479) forbids `StopIteration` escaping a coroutine and converts it
  to `RuntimeError("coroutine raised StopIteration")` before the `except` clause
  ever sees it.

  **Affected versions:** Regression introduced in v2.2.0 — the route and the
  buggy adapter were added in the same commit. Not present in the v1.x line
  (verified absent at v1.35.0).

  **Solution:**
  Stop relying on `StopIteration` crossing the await boundary; use a sentinel
  default so iterator exhaustion never raises across it:

      _SENTINEL = object()

      async def _aiter_from_iter(iterator):
        while True:
          chunk = await run_in_threadpool(next, iterator, _SENTINEL)
          if chunk is _SENTINEL:
            break
          yield chunk

  This is the minimal, idiomatic fix; the stream now terminates cleanly when the
  sync generator is exhausted.

  ## Testing Plan

  **Unit Tests:**

  - [x] I have added or updated unit tests for my change.
  - [x] All unit tests pass locally.

  Added `test_gemini_stream_reasoning_engine_sync_generator` plus a
  `test_app_with_gemini_enterprise_sync_stream` fixture in
  `tests/unittests/cli/test_fast_api.py`. The pre-existing stream test used an
  *async* generator (the `isasyncgenfunction` branch) and never exercised the
  buggy sync-generator path. The new test fails on the unpatched code with
  `RuntimeError` and passes with the fix.

  pytest summary:

      $ pytest tests/unittests/cli/test_fast_api.py -k stream_reasoning_engine -q
      3 passed, 79 deselected

      $ pytest tests/unittests/cli/test_fast_api.py -q
      82 passed

  **Manual End-to-End (E2E) Tests:**

  The failure and the fix reproduce standalone in ~15 lines, independent of any
  model or deployment:

      import asyncio
      from starlette.concurrency import run_in_threadpool

      async def _aiter_from_iter(iterator):  # old, buggy version
          while True:
              try:
                  chunk = await run_in_threadpool(next, iterator)
                  yield chunk
              except StopIteration:
                  break

      async def main():
          def gen():
              yield 1
              yield 2
          async for c in _aiter_from_iter(gen()):
              print("chunk:", c)

      asyncio.run(main())
      # chunk: 1
      # chunk: 2
      # RuntimeError: coroutine raised StopIteration   <-- before the fix

  With the sentinel version above, the same script prints the two chunks and
  exits cleanly with no exception. Originally observed on a live Vertex AI Agent
  Engine deployment (google-adk==2.2.0, Python 3.11) where every `stream_query`
  call logged the RuntimeError after the final chunk.

  ## Checklist

  - [x] I have read the CONTRIBUTING.md document.
  - [x] I have performed a self-review of my own code.
  - [x] I have commented my code, particularly in hard-to-understand areas.
  - [x] I have added tests that prove my fix is effective or that my feature works.
  - [x] New and existing unit tests pass locally with my changes.
  - [x] I have manually tested my changes end-to-end.
  - [ ] Any dependent changes have been merged and published in downstream modules.

  ## Additional context

  Original server traceback:

      ERROR:    Exception in ASGI application
      Traceback (most recent call last):
        File ".../starlette/responses.py", line 250, in stream_response
          async for chunk in self.body_iterator:
        File ".../google/adk/cli/fast_api.py", line 797, in json_generator
          async for chunk in output:
        File ".../google/adk/cli/fast_api.py", line 919, in _aiter_from_iter
          chunk = await run_in_threadpool(next, iterator)
        File ".../starlette/concurrency.py", line 32, in run_in_threadpool
          return await anyio.to_thread.run_sync(func)
        File ".../anyio/to_thread.py", line 63, in run_sync
          return await get_async_backend().run_sync_in_worker_thread(
        File ".../anyio/_backends/_asyncio.py", line 2518, in run_sync_in_worker_thread
          return await future
      RuntimeError: coroutine raised StopIteration

  Occurs 100% of the time on every sync streaming request once the generator is
  exhausted. The bug is model-agnostic (purely in the FastAPI streaming adapter).

Co-authored-by: George Weale <gweale@google.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/6114 from surajit-1306:fix/stream-reasoning-engine-stopiteration e5ee866074fefc56418ec03441e3706617f9d755
PiperOrigin-RevId: 962875380
This commit is contained in:
Surajit Nandi
2026-08-11 10:53:56 -07:00
committed by Copybara-Service
parent f4fd7d5db9
commit aa9c187f46
2 changed files with 89 additions and 4 deletions
+11 -4
View File
@@ -604,14 +604,21 @@ def get_fast_api_app(
output = await _invoke_callable_or_raise(method, parsed.input or {})
if inspect.isgenerator(output):
# Sentinel-based exhaustion check. We cannot rely on catching
# StopIteration here: when ``next(iterator)`` is called inside the
# threadpool worker, the StopIteration propagates out of the
# ``run_in_threadpool`` coroutine frame, and Python (PEP 479) converts
# it to ``RuntimeError("coroutine raised StopIteration")`` before the
# ``except StopIteration`` clause can ever see it. Passing a default to
# ``next`` avoids raising at the boundary entirely.
_SENTINEL = object()
async def _aiter_from_iter(iterator):
while True:
try:
chunk = await run_in_threadpool(next, iterator)
yield chunk
except StopIteration:
chunk = await run_in_threadpool(next, iterator, _SENTINEL)
if chunk is _SENTINEL:
break
yield chunk
content_iter = _aiter_from_iter(output)
else:
+78
View File
@@ -1046,6 +1046,63 @@ def test_app_with_gemini_enterprise(
yield client
@pytest.fixture
def test_app_with_gemini_enterprise_sync_stream(
mock_session_service,
mock_artifact_service,
mock_memory_service,
mock_agent_loader,
mock_eval_sets_manager,
mock_eval_set_results_manager,
monkeypatch,
):
"""Like test_app_with_gemini_enterprise but stream_query is a sync generator.
This exercises the inspect.isgenerator() branch in stream_reasoning_engine,
where the sync iterator is adapted to an async iterator via a threadpool.
"""
monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project")
mock_agent_loader.list_agents = MagicMock(
return_value=["test_app", "gemini_app"]
)
mock_adk_app_instance = MagicMock()
mock_adk_app_instance._tmpl_attrs = {}
def stream_query_impl(**kwargs):
yield {"chunk": 1, "kwargs": kwargs}
yield {"chunk": 2, "kwargs": kwargs}
mock_adk_app_instance.stream_query = stream_query_impl
with (
patch("google.auth.default", return_value=(MagicMock(), "test-project")),
patch("vertexai.init", new_callable=MagicMock),
patch(
"vertexai.agent_engines.AdkApp", return_value=mock_adk_app_instance
),
patch("google.adk.agents.Agent", new_callable=MagicMock),
patch(
"google.adk.telemetry._agent_engine.TopSpanProcessor",
new_callable=MagicMock,
),
patch(
"google.adk.telemetry._agent_engine.get_propagated_context",
new_callable=MagicMock,
),
):
client = _create_test_client(
mock_session_service,
mock_artifact_service,
mock_memory_service,
mock_agent_loader,
mock_eval_sets_manager,
mock_eval_set_results_manager,
gemini_enterprise_app_name="gemini_app",
)
yield client
#################################################
# Test Cases
#################################################
@@ -3581,6 +3638,27 @@ def test_gemini_stream_reasoning_engine_missing_class_method(
assert response.status_code == 400
def test_gemini_stream_reasoning_engine_sync_generator(
test_app_with_gemini_enterprise_sync_stream,
):
"""Regression test: a synchronous streaming class_method must not raise.
A sync generator is adapted to an async iterator via run_in_threadpool. The
adapter must not rely on catching StopIteration across the await boundary,
since Python (PEP 479) converts an escaping StopIteration into
RuntimeError("coroutine raised StopIteration") after the final chunk.
"""
response = test_app_with_gemini_enterprise_sync_stream.post(
"/api/stream_reasoning_engine",
json={"class_method": "stream_query", "input": {"arg1": 1}},
)
assert response.status_code == 200
lines = response.text.strip().split("\n")
assert len(lines) == 2
assert json.loads(lines[0]) == {"chunk": 1, "kwargs": {"arg1": 1}}
assert json.loads(lines[1]) == {"chunk": 2, "kwargs": {"arg1": 1}}
def test_run_eval_request_live_fields_default():
"""RunEvalRequest defaults to non-live mode."""
from google.adk.cli.dev_server import RunEvalRequest