Merge branch 'main' into fix/polly-review-one-shot-mode

_drain_extra_turns now waits for synthesis after dispatch. The three tests
that dispatch sub-agents (distinct-models, list-then-dispatch, canonical-id)
only configured Polly's dispatch turn — the process would hang waiting for
a synthesis response that never came.

Sub-agents (openai-agents, OPENAI_BASE_URL → mock server) fail fast when
no response is queued for their model key, triggering the inbox wake notice.
Polly's synthesis turn then needs a mock response — add one to each affected
test. Also raise _RUN_TIMEOUT_SEC 120 → 300 to give the extra turn room.

test_polly_rejects_cross_family_model_dispatch is unaffected: the dispatch
fails validation before creating any child, so _subagent_work_by_parent is
empty → runner emits 'idle' → fast-exit as before.

Co-authored-by: Tomu Hirata

omnigent/chat.py

@@ -2258,39 +2258,63 @@ async def _query_sessions_once(
     # sub-agents and are auto-woken by inbox completions across multiple
     # turns.
     #
-    # Fast-exit: refresh() at the TOP of each iteration catches the common
-    # case (single-turn agent, session already idle) with one HTTP round-
-    # trip (~100 ms) instead of waiting up to _PER_TURN_TIMEOUT_S for a
-    # stream subscription to time out.
+    # Why not _collect_query for the fast-exit signal: the runtime emits
+    # ``session.status: waiting`` AFTER ``response.completed`` (the runner
+    # finishes dispatching tools, then enters the async drain). _collect_query
+    # exits at CompletedEvent and never sees the subsequent "waiting".
     #
-    # Race window: a turn MAY complete in the gap between the top-of-loop
-    # refresh() showing "waiting" and await_turn() opening its subscription.
-    # The window is O(ms) in practice (subagents take seconds). If it fires,
-    # await_turn() times out, the bottom refresh() shows "idle", and we exit
-    # — the only cost is one _PER_TURN_TIMEOUT_S wait and possibly missing
-    # that turn's text.
+    # Why not refresh() for the fast-exit signal: the snapshot API collapses
+    # the ``"waiting"`` relay status to ``"idle"`` once the turn loop exits,
+    # even while sub-agents are still running.
     #
-    # Timeouts: 120 s per turn bounds the race-window penalty. A global
-    # 1800 s wall-clock budget caps the loop regardless of turn count.
+    # Probe approach: subscribe to the live stream for a short window after
+    # the first turn. The "waiting" event arrives O(ms–s) after CompletedEvent
+    # (runner dispatches tools, spawns sub-agents, then parks). The probe
+    # catches it before sub-agents have a chance to complete.
+    #
+    # ``await_turn`` resets ``last_turn_saw_waiting`` to False on
+    # ``session.status: running`` (synthesis starting), so the flag cleanly
+    # reflects only the most recent dispatch state after each call.
+    #
+    # Single-turn agents: no "waiting" event ever → probe times out in
+    # _STATUS_PROBE_TIMEOUT_S (~30 s) and the loop exits.
     _MAX_EXTRA_TURNS = 30
-    _PER_TURN_TIMEOUT_S = 120.0
-    _LOOP_TIMEOUT_S = 1800.0
+    # The runner emits session.status:waiting (not idle) when a turn ends with
+    # running sub-agents. The relay cache holds "waiting", which the snapshot
+    # collapses to "running". refresh() is therefore the authoritative signal:
+    # "running" → async orchestrator still waiting for inbox; "idle" → done.
+    #
+    # A short probe await_turn runs first: it catches synthesis text or the
+    # status event if the subscription opens before the event arrives. Both
+    # "waiting" and "idle" break the probe immediately so the generator closes
+    # cleanly without hitting the timeout.
+    #
+    # refresh() is called after every await_turn (probe + loop) — it is correct
+    # even when await_turn times out (sub-agents still running), unlike the
+    # last_turn_saw_waiting flag which would incorrectly exit on timeout.
+    _STATUS_PROBE_TIMEOUT_S = 5.0  # brief window; status events arrive fast
+    _PER_TURN_TIMEOUT_S = 120.0  # race-window guard per synthesis turn
+    _LOOP_TIMEOUT_S = 1800.0  # 30 min total
 
     async def _drain_extra_turns() -> None:
+        # Probe: collect synthesis text or status events that arrive quickly.
+        probe = await chat.await_turn(timeout=_STATUS_PROBE_TIMEOUT_S)
+        if probe.text:
+            all_text_parts.append(probe.text)
+        # refresh() is the authoritative check: "running" means the runner's
+        # relay cache holds "waiting" (sub-agents still running); "idle" means
+        # truly done (single-turn agent, or synthesis completed in the probe).
+        await chat.refresh()
+        if chat.status not in ("running", "launching"):
+            return
+        # Async orchestrator confirmed. Loop, refreshing after each turn.
         for _ in range(_MAX_EXTRA_TURNS):
-            # Fast-exit for single-turn agents.
-            await chat.refresh()
-            if chat.status not in ("waiting", "running", "launching"):
-                return
-            # Session still active; subscribe before the next check to
-            # reduce (not eliminate) the race where a turn completes
-            # between refresh and subscribe.
             extra = await chat.await_turn(timeout=_PER_TURN_TIMEOUT_S)
             if extra.text:
                 all_text_parts.append(extra.text)
             await chat.refresh()
-            if chat.status not in ("waiting", "running", "launching"):
-                return
+            if chat.status not in ("running", "launching"):
+                return  # Idle: synthesis done or all sub-agents complete.
         logger.warning(
             "headless -p hit the %d-turn guard for session %s; "
             "the orchestrator may still be running",

omnigent/repl/_repl.py

@@ -1775,7 +1775,11 @@ class _SessionsChatReplAdapter:
                         flush=True,
                     )
                 async for event in self._client.sessions.stream(self._session_id):
-                    if isinstance(event, _StatusEv) and event.status in ("idle", "failed"):
+                    if isinstance(event, _StatusEv) and event.status in (
+                        "idle",
+                        "waiting",
+                        "failed",
+                    ):
                         turn_done = getattr(self, "_turn_done", None)
                         if turn_done is not None:
                             turn_done.set()
@@ -2886,7 +2890,7 @@ async def run_repl(
                 # a stream-pump reconnect gap or before this REPL
                 # attached is lost — so also re-sync at each turn start.
                 _spawn_metadata_refresh()
-            elif event.status in ("idle", "failed"):
+            elif event.status in ("idle", "waiting", "failed"):
                 from omnigent_client import TextDone
 
                 # A SETUP-phase failure (spec resolution, spawn-env

omnigent/runner/app.py

@@ -8001,7 +8001,20 @@ def create_runner_app(
             _publish_turn_status(conv_id, "failed", error=_normalize_turn_error(error))
         else:
             if not has_buffered:
-                _publish_turn_status(conv_id, "idle")
+                # Emit ``waiting`` instead of ``idle`` when the turn ended
+                # cleanly but sub-agents are still running. This lets the
+                # headless ``-p`` multi-turn loop (``_drain_extra_turns`` in
+                # ``chat.py``) distinguish an async orchestrator that parked
+                # on the inbox drain from a truly finished single-turn agent —
+                # both would otherwise emit ``idle`` here, making them
+                # indistinguishable without a "waiting" signal.
+                children = _subagent_work_by_parent.get(conv_id, set())
+                has_running_children = any(
+                    (e := _subagent_work_by_child.get(c)) is not None
+                    and e.status in ("launching", "running", "waiting")
+                    for c in children
+                )
+                _publish_turn_status(conv_id, "waiting" if has_running_children else "idle")
         if was_interrupted:
             _mark_subagent_terminal_and_wake(
                 conv_id,

omnigent/server/routes/sessions.py

@@ -18225,30 +18225,31 @@ async def _get_session_snapshot(
     if runner_client is None:
         runner_client = get_runner_client()
 
-    status = _session_status_cache.get(session_id)
-    if status is None:
-        # Cache miss: either the server restarted, or the relay
-        # has not yet published the first ``"running"`` event
-        # for a freshly bound session (the relay's GET /stream
-        # is still in its tunnel handshake). Ask the runner for
-        # live status so we don't synthesize a stale ``"idle"``
-        # while a turn is actually in flight.
-        if runner_client is not None:
+    status = _session_status_from_cache(session_id)
+    if status == "idle":
+        # Cache miss (or truly idle): either the server restarted, or the
+        # relay has not yet published the first ``"running"`` event for a
+        # freshly bound session (the relay's GET /stream is still in its
+        # tunnel handshake). Ask the runner for live status so we don't
+        # synthesize a stale ``"idle"`` while a turn is actually in flight.
+        # ``_session_status_from_cache`` already collapses the fine-grained
+        # relay values (``"waiting"`` → ``"running"``), so the raw cache value
+        # is only needed here when it is actually missing (None).
+        if _session_status_cache.get(session_id) is None and runner_client is not None:
             try:
                 resp = await runner_client.get(
                     f"/v1/sessions/{session_id}",
                     timeout=5.0,
                 )
                 if resp.status_code == 200:
-                    status = resp.json().get("status", "idle")
-                    _session_status_cache[session_id] = status
+                    raw = resp.json().get("status", "idle")
+                    _session_status_cache[session_id] = raw
+                    status = _session_status_from_cache(session_id)
             except httpx.HTTPError:
                 _logger.debug(
                     "Runner status query failed for %s",
                     session_id,
                 )
-        if status is None:
-            status = "idle"
     # last_total_tokens and last_task_error come from the context-tokens
     # label written by the forwarder (tasks table has been removed).
     last_total_tokens: int | None = None

sdks/python-client/omnigent_client/_sessions_chat.py

@@ -1105,6 +1105,14 @@ class SessionsChat:
                     if not text_parts:
                         text_parts.extend(_assistant_text_from_response(event.response.output))
                     break
+                elif isinstance(event, SessionStatusEvent) and event.status in (
+                    "waiting",
+                    "idle",
+                ):
+                    # Break on terminal status events so the async generator
+                    # closes cleanly (avoids "aclose(): already running" when
+                    # asyncio.timeout fires mid-stream).
+                    break
                 elif isinstance(event, _TURN_TERMINAL_EVENT_TYPES):
                     break
 

tests/cli/test_chat.py

@@ -3270,33 +3270,31 @@ def _fake_sessions_chat_cls(
     :param extra_turns: Optional list of text strings to return from
         successive ``await_turn()`` calls, simulating async orchestrator
         auto-wakes. When exhausted ``await_turn`` returns empty text and
-        ``status`` returns ``"idle"``.
+        ``last_turn_saw_waiting`` returns ``False``.
     :returns: A class usable as a drop-in for ``SessionsChat``.
     """
     _extra = list(extra_turns or [])
 
     class _FakeSessionsChat:
         def __init__(self, **_kwargs: object) -> None:
-            self._status = "waiting" if _extra else "idle"
             self._pending = list(_extra)
 
         @property
         def status(self) -> str:
-            return self._status
+            # Mirrors the real snapshot: "running" while sub-agents are pending
+            # (the runner emits "waiting" → relay collapses to "running"),
+            # "idle" when done.
+            return "running" if self._pending else "idle"
 
         async def refresh(self) -> None:
-            # After await_turn drains a turn, mark idle when nothing left.
-            if not self._pending:
-                self._status = "idle"
+            pass  # status is derived from _pending; no fetch needed.
 
-        async def query(self, prompt: str) -> object:
-            return await query_impl(prompt)
+        async def query(self, prompt: str) -> QueryResult:
+            return await query_impl(prompt)  # type: ignore[return-value]
 
         async def await_turn(self, *, timeout: float | None = None) -> QueryResult:
             if self._pending:
                 text = self._pending.pop(0)
-                if not self._pending:
-                    self._status = "idle"
                 return QueryResult(text=text, files=[])
             return QueryResult(text="", files=[])
 

tests/e2e/test_polly_subagent_model_e2e.py

@@ -57,7 +57,7 @@ from tests.e2e.test_polly_e2e import (
 # tests/e2e/test_polly_subagent_model_e2e.py -> repo root is 2 parents up.
 _POLLY = _REPO / "examples" / "polly"
 # Mock runs are fast (no real model inference) so a short timeout is enough.
-_RUN_TIMEOUT_SEC = 120
+_RUN_TIMEOUT_SEC = 300
 
 # Models dispatched to each worker in the multi-dispatch test.
 # Under mock (no Databricks creds), the dispatch gate localizes models for
@@ -286,6 +286,9 @@ def test_polly_dispatches_distinct_models_per_worker(
             },
             # Second response: after tool results arrive, end the turn.
             {"text": "Dispatched all three workers. Waiting for inbox notices."},
+            # Third response: synthesis after sub-agents complete (or fail fast on
+            # the mock server with no queued responses).
+            {"text": "All three workers done. Model overrides verified."},
         ],
         key=_MOCK_BRAIN_MODEL,
     )
@@ -476,6 +479,8 @@ def test_polly_lists_models_then_dispatches_pi_from_list(
             },
             # Step 3: end the turn after dispatch.
             {"text": "Dispatched pi on a Claude model from the catalog."},
+            # Step 4: synthesis after pi completes (or fails fast on mock).
+            {"text": "Pi done. Model override verified."},
         ],
         key=_MOCK_BRAIN_MODEL,
     )
@@ -593,6 +598,8 @@ def test_polly_canonical_id_localized_for_gateway_child(
             },
             # End the turn after dispatch.
             {"text": "Dispatched pi on claude-opus-4-8. Waiting for inbox."},
+            # Synthesis after pi completes (or fails fast on mock).
+            {"text": "Pi done. Canonical id localized and persisted."},
         ],
         key=_MOCK_BRAIN_MODEL,
     )