diff --git a/omnigent/inner/acp_executor.py b/omnigent/inner/acp_executor.py index e296b7ba17..d8de83edf7 100644 --- a/omnigent/inner/acp_executor.py +++ b/omnigent/inner/acp_executor.py @@ -1108,57 +1108,66 @@ async def run_turn( deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS accumulated_text: list[str] = [] - while True: - remaining = deadline - loop.time() - if remaining <= 0: - yield ExecutorError(message="Timeout waiting for ACP response", retryable=True) - return - - # Complete only once the future is resolved AND the queue is drained, - # so trailing chunks aren't truncated. - if fut.done() and self._queue.empty(): - try: - response = fut.result() - except Exception as exc: # noqa: BLE001 - self._session_id = None - self._system_prompt_sent = False - yield ExecutorError(message=f"ACP process error: {exc}", retryable=True) + # The stdout reader pops ``req_id`` only when it matches a RESPONSE + # (see ``_read_stdout``). On the timeout path below, or when the reader + # resolves ``fut`` with an exception (EOF / reader error) without a + # match, nothing removes ``req_id`` — so ``self._pending`` would grow one + # stale entry per silent/failed turn on a long-lived ACP session. Drop + # it in ``finally`` on every exit (mirrors ``_rpc``'s timeout cleanup). + try: + while True: + remaining = deadline - loop.time() + if remaining <= 0: + yield ExecutorError(message="Timeout waiting for ACP response", retryable=True) return - if "error" in response: - error_msg = response["error"].get("message", "Unknown ACP error") - if "Session not found" in error_msg: + + # Complete only once the future is resolved AND the queue is drained, + # so trailing chunks aren't truncated. + if fut.done() and self._queue.empty(): + try: + response = fut.result() + except Exception as exc: # noqa: BLE001 self._session_id = None self._system_prompt_sent = False - yield ExecutorError(message=error_msg, retryable=True) + yield ExecutorError(message=f"ACP process error: {exc}", retryable=True) + return + if "error" in response: + error_msg = response["error"].get("message", "Unknown ACP error") + if "Session not found" in error_msg: + self._session_id = None + self._system_prompt_sent = False + yield ExecutorError(message=error_msg, retryable=True) + return + result = response.get("result", {}) if isinstance(response, dict) else {} + usage = self._usage_from_result(result) if isinstance(result, dict) else None + yield TurnComplete(response="".join(accumulated_text), usage=usage) return - result = response.get("result", {}) if isinstance(response, dict) else {} - usage = self._usage_from_result(result) if isinstance(result, dict) else None - yield TurnComplete(response="".join(accumulated_text), usage=usage) - return - try: - notification = await asyncio.wait_for( - self._queue.get(), timeout=min(remaining, 2.0) - ) - except asyncio.TimeoutError: - continue + try: + notification = await asyncio.wait_for( + self._queue.get(), timeout=min(remaining, 2.0) + ) + except asyncio.TimeoutError: + continue - method = notification.get("method", "") - params = notification.get("params", {}) - - if method == _CLIENT_NOTIFICATION_SESSION_UPDATE: - update = params.get("update", {}) - for event in self._handle_session_update(update): - if isinstance(event, TextChunk): - accumulated_text.append(event.text) - yield event - elif notification.get("id") is not None and notification.get("method"): - # Server-initiated request (session/request_permission / fs/*): - # routes through policy + elicitation. Blocks while the human decides. - await self._respond_to_agent_request(notification) - - # Inbound message = progress; reset the idle deadline. - deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS + method = notification.get("method", "") + params = notification.get("params", {}) + + if method == _CLIENT_NOTIFICATION_SESSION_UPDATE: + update = params.get("update", {}) + for event in self._handle_session_update(update): + if isinstance(event, TextChunk): + accumulated_text.append(event.text) + yield event + elif notification.get("id") is not None and notification.get("method"): + # Server-initiated request (session/request_permission / fs/*): + # routes through policy + elicitation. Blocks while the human decides. + await self._respond_to_agent_request(notification) + + # Inbound message = progress; reset the idle deadline. + deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS + finally: + self._pending.pop(req_id, None) async def interrupt_session(self, session_key: str) -> bool: # noqa: ARG002 — one ACP session per process """Abort the running turn via the ACP ``session/cancel`` notification. diff --git a/omnigent/inner/goose_executor.py b/omnigent/inner/goose_executor.py index 9fd3e987a2..db234398d1 100644 --- a/omnigent/inner/goose_executor.py +++ b/omnigent/inner/goose_executor.py @@ -1184,75 +1184,85 @@ async def run_turn( deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS accumulated_text: list[str] = [] - while True: - remaining = deadline - loop.time() - if remaining <= 0: - yield ExecutorError(message="Timeout waiting for goose response", retryable=True) - return - - # Complete only once the future is resolved AND the queue is drained, - # so trailing chunks aren't truncated. - if fut.done() and self._queue.empty(): - try: - response = fut.result() - except Exception as exc: # noqa: BLE001 - self._session_id = None - self._system_prompt_sent = False - yield ExecutorError(message=f"goose process error: {exc}", retryable=True) + # The stdout reader pops ``req_id`` only when it matches a RESPONSE. + # On timeout, or when the reader resolves ``fut`` with an exception + # (EOF / reader error) without a match, nothing removes it, so + # ``self._pending`` would grow one stale entry per silent/failed turn + # on a long-lived session. Drop it on every exit (mirrors ``_rpc``). + try: + while True: + remaining = deadline - loop.time() + if remaining <= 0: + yield ExecutorError( + message="Timeout waiting for goose response", retryable=True + ) return - if "error" in response: - error_msg = response["error"].get("message", "Unknown ACP error") - if "Session not found" in error_msg: + + # Complete only once the future is resolved AND the queue is drained, + # so trailing chunks aren't truncated. + if fut.done() and self._queue.empty(): + try: + response = fut.result() + except Exception as exc: # noqa: BLE001 self._session_id = None self._system_prompt_sent = False - yield ExecutorError(message=error_msg, retryable=True) + yield ExecutorError(message=f"goose process error: {exc}", retryable=True) + return + if "error" in response: + error_msg = response["error"].get("message", "Unknown ACP error") + if "Session not found" in error_msg: + self._session_id = None + self._system_prompt_sent = False + yield ExecutorError(message=error_msg, retryable=True) + return + result = response.get("result", {}) if isinstance(response, dict) else {} + usage = self._usage_from_result(result) if isinstance(result, dict) else None + yield TurnComplete(response="".join(accumulated_text), usage=usage) return - result = response.get("result", {}) if isinstance(response, dict) else {} - usage = self._usage_from_result(result) if isinstance(result, dict) else None - yield TurnComplete(response="".join(accumulated_text), usage=usage) - return - try: - notification = await asyncio.wait_for( - self._queue.get(), timeout=min(remaining, 2.0) - ) - except asyncio.TimeoutError: - continue + try: + notification = await asyncio.wait_for( + self._queue.get(), timeout=min(remaining, 2.0) + ) + except asyncio.TimeoutError: + continue - method = notification.get("method", "") - params = notification.get("params", {}) - - if method == _CLIENT_NOTIFICATION_SESSION_UPDATE: - update = params.get("update", {}) - update_type = update.get("sessionUpdate", "") - - if update_type == _UPDATE_AGENT_MESSAGE_CHUNK: - content = update.get("content", {}) - text = content.get("text", "") if isinstance(content, dict) else "" - if text: - accumulated_text.append(text) - yield TextChunk(text=text) - elif update_type == _UPDATE_USAGE: - size = update.get("size") - if isinstance(size, int) and size > 0: - self._context_window = size - elif update_type == _UPDATE_TOOL_CALL: - logger.debug("goose tool_call: %s", update.get("title", "tool_call")) - elif update_type == _UPDATE_TOOL_CALL_UPDATE: - pass - - elif notification.get("id") is not None and notification.get("method"): - # Server-initiated request (session/request_permission): routes - # through policy + elicitation. Blocks while the human decides. - await self._respond_to_agent_request(notification) - # Surface any fs ToolCall events the handler buffered so the - # I/O shows in history. - while self._fs_events: - yield self._fs_events.pop(0) - - # Inbound message = progress; reset the idle deadline (after the - # approval block so a slow approval doesn't time out). - deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS + method = notification.get("method", "") + params = notification.get("params", {}) + + if method == _CLIENT_NOTIFICATION_SESSION_UPDATE: + update = params.get("update", {}) + update_type = update.get("sessionUpdate", "") + + if update_type == _UPDATE_AGENT_MESSAGE_CHUNK: + content = update.get("content", {}) + text = content.get("text", "") if isinstance(content, dict) else "" + if text: + accumulated_text.append(text) + yield TextChunk(text=text) + elif update_type == _UPDATE_USAGE: + size = update.get("size") + if isinstance(size, int) and size > 0: + self._context_window = size + elif update_type == _UPDATE_TOOL_CALL: + logger.debug("goose tool_call: %s", update.get("title", "tool_call")) + elif update_type == _UPDATE_TOOL_CALL_UPDATE: + pass + + elif notification.get("id") is not None and notification.get("method"): + # Server-initiated request (session/request_permission): routes + # through policy + elicitation. Blocks while the human decides. + await self._respond_to_agent_request(notification) + # Surface any fs ToolCall events the handler buffered so the + # I/O shows in history. + while self._fs_events: + yield self._fs_events.pop(0) + + # Inbound message = progress; reset the idle deadline (after the + # approval block so a slow approval doesn't time out). + deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS + finally: + self._pending.pop(req_id, None) async def close_session(self, session_key: str) -> None: """Close a named session (no-op; the ACP session is per-process).""" diff --git a/omnigent/inner/qwen_executor.py b/omnigent/inner/qwen_executor.py index 12ca130286..b51b7d3ed0 100644 --- a/omnigent/inner/qwen_executor.py +++ b/omnigent/inner/qwen_executor.py @@ -1340,97 +1340,107 @@ async def run_turn( # (see _accumulate_usage). Stays empty when qwen reports none. turn_usage: dict[str, int] = {} - while True: - remaining = deadline - loop.time() - if remaining <= 0: - yield ExecutorError(message="Timeout waiting for qwen response", retryable=True) - return - - # Complete only once the future is resolved AND the queue is drained. - # The reader resolves the future directly but enqueues chunks, and - # the response trails the chunk stream — a bare fut.done() check - # could return with chunks still buffered, truncating the response. - if fut.done() and self._queue.empty(): - try: - response = fut.result() - except Exception as exc: # noqa: BLE001 - # The stdout reader sets an exception on the future when the - # subprocess dies. Surface it as a clean retryable error - # rather than letting it raise out of the generator. - self._session_id = None - self._system_prompt_sent = False - yield ExecutorError(message=f"qwen process error: {exc}", retryable=True) + # The stdout reader pops ``req_id`` only when it matches a RESPONSE. + # On timeout, or when the reader resolves ``fut`` with an exception + # (EOF / reader error) without a match, nothing removes it, so + # ``self._pending`` would grow one stale entry per silent/failed turn + # on a long-lived session. Drop it on every exit (mirrors ``_rpc``). + try: + while True: + remaining = deadline - loop.time() + if remaining <= 0: + yield ExecutorError( + message="Timeout waiting for qwen response", retryable=True + ) return - if "error" in response: - error_msg = response["error"].get("message", "Unknown ACP error") - # If the session was lost, reset so next turn creates a new - # one — and re-send the system prompt into that fresh session. - if "Session not found" in error_msg: + + # Complete only once the future is resolved AND the queue is drained. + # The reader resolves the future directly but enqueues chunks, and + # the response trails the chunk stream — a bare fut.done() check + # could return with chunks still buffered, truncating the response. + if fut.done() and self._queue.empty(): + try: + response = fut.result() + except Exception as exc: # noqa: BLE001 + # The stdout reader sets an exception on the future when the + # subprocess dies. Surface it as a clean retryable error + # rather than letting it raise out of the generator. self._session_id = None self._system_prompt_sent = False - yield ExecutorError(message=error_msg, retryable=True) + yield ExecutorError(message=f"qwen process error: {exc}", retryable=True) + return + if "error" in response: + error_msg = response["error"].get("message", "Unknown ACP error") + # If the session was lost, reset so next turn creates a new + # one — and re-send the system prompt into that fresh session. + if "Session not found" in error_msg: + self._session_id = None + self._system_prompt_sent = False + yield ExecutorError(message=error_msg, retryable=True) + return + # Successful completion. Attach the per-turn token usage qwen + # reported over the stream (None when it reported none) and feed + # the cost observer, mirroring the codex executor. + usage = turn_usage or None + if usage is not None: + _notify_usage_from_dict(model=self._model, usage=usage) + final_text = "".join(accumulated_text) + yield TurnComplete(response=final_text if final_text else "", usage=usage) return - # Successful completion. Attach the per-turn token usage qwen - # reported over the stream (None when it reported none) and feed - # the cost observer, mirroring the codex executor. - usage = turn_usage or None - if usage is not None: - _notify_usage_from_dict(model=self._model, usage=usage) - final_text = "".join(accumulated_text) - yield TurnComplete(response=final_text if final_text else "", usage=usage) - return - - # Otherwise consume queued notifications. - try: - notification = await asyncio.wait_for( - self._queue.get(), timeout=min(remaining, 2.0) - ) - except asyncio.TimeoutError: - continue - method = notification.get("method", "") - params = notification.get("params", {}) - - if method == _CLIENT_NOTIFICATION_SESSION_UPDATE: - update = params.get("update", {}) - update_type = update.get("sessionUpdate", "") - - if update_type == _UPDATE_AGENT_MESSAGE_CHUNK: - # qwen rides per-call token usage on an agent_message_chunk - # with empty text + a populated _meta.usage, so fold usage - # before the text check (the usage-bearing chunk has none). - self._accumulate_usage(turn_usage, update) - content = update.get("content", {}) - if isinstance(content, dict): - text = content.get("text", "") - else: - text = "" - if text: - accumulated_text.append(text) - yield TextChunk(text=text) - - elif update_type == _UPDATE_TOOL_CALL: - # Qwen is executing a built-in tool — surface it as info. - tool_title = update.get("title", "tool_call") - logger.debug("qwen tool_call: %s", tool_title) - - elif update_type == _UPDATE_TOOL_CALL_UPDATE: - # Status update on an in-progress tool call — skip. - pass - - elif notification.get("id") is not None and notification.get("method"): - # Server-initiated request (e.g. session/request_permission): - # permission goes through policy + elicitation; anything else - # gets method-not-found. Blocks while the human decides. - await self._respond_to_agent_request(notification) - # Surface any fs ToolCall events the handler buffered so the - # I/O shows in history. - while self._fs_events: - yield self._fs_events.pop(0) - - # Inbound message = progress; reset the idle deadline. Runs after the - # human-approval block above so a slow approval doesn't time out. - deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS + # Otherwise consume queued notifications. + try: + notification = await asyncio.wait_for( + self._queue.get(), timeout=min(remaining, 2.0) + ) + except asyncio.TimeoutError: + continue + + method = notification.get("method", "") + params = notification.get("params", {}) + + if method == _CLIENT_NOTIFICATION_SESSION_UPDATE: + update = params.get("update", {}) + update_type = update.get("sessionUpdate", "") + + if update_type == _UPDATE_AGENT_MESSAGE_CHUNK: + # qwen rides per-call token usage on an agent_message_chunk + # with empty text + a populated _meta.usage, so fold usage + # before the text check (the usage-bearing chunk has none). + self._accumulate_usage(turn_usage, update) + content = update.get("content", {}) + if isinstance(content, dict): + text = content.get("text", "") + else: + text = "" + if text: + accumulated_text.append(text) + yield TextChunk(text=text) + + elif update_type == _UPDATE_TOOL_CALL: + # Qwen is executing a built-in tool — surface it as info. + tool_title = update.get("title", "tool_call") + logger.debug("qwen tool_call: %s", tool_title) + + elif update_type == _UPDATE_TOOL_CALL_UPDATE: + # Status update on an in-progress tool call — skip. + pass + + elif notification.get("id") is not None and notification.get("method"): + # Server-initiated request (e.g. session/request_permission): + # permission goes through policy + elicitation; anything else + # gets method-not-found. Blocks while the human decides. + await self._respond_to_agent_request(notification) + # Surface any fs ToolCall events the handler buffered so the + # I/O shows in history. + while self._fs_events: + yield self._fs_events.pop(0) + + # Inbound message = progress; reset the idle deadline. Runs after the + # human-approval block above so a slow approval doesn't time out. + deadline = loop.time() + _PROMPT_TIMEOUT_SECONDS + finally: + self._pending.pop(req_id, None) async def close_session(self, session_key: str) -> None: """Close a named session (no-op; sessions are per-process).""" diff --git a/tests/inner/test_acp_executor.py b/tests/inner/test_acp_executor.py index 23c9cc53d0..d5c34cb07d 100644 --- a/tests/inner/test_acp_executor.py +++ b/tests/inner/test_acp_executor.py @@ -25,6 +25,7 @@ from omnigent.inner._acp_omnigent_mcp import OmnigentAcpMcp, _to_acp_mcp_servers from omnigent.inner.acp_executor import AcpAgentConfig, AcpExecutor from omnigent.inner.executor import ( + ExecutorError, ReasoningChunk, TextChunk, ToolCallComplete, @@ -569,3 +570,51 @@ async def deny(tool_name: str, tool_input: dict) -> bool: # Turn still completes even though the tool was rejected. assert any(isinstance(e, TurnComplete) for e in events) + + +# --------------------------------------------------------------------------- +# run_turn cleans up its pending prompt future on every exit path +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_turn_timeout_does_not_leak_pending_future( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A turn that times out must not leave its request future in ``_pending``. + + The stdout reader pops ``req_id`` only when it matches a response. On the + timeout path (and the EOF/reader-error path) nothing removes it, so before + the ``finally`` cleanup ``_pending`` grew one stale entry per silent/failed + turn on a long-lived ACP session. Drive a turn whose response never arrives + and assert the map is empty afterward. + """ + from types import SimpleNamespace + + import omnigent.inner.acp_executor as acp_mod + + # Time out the prompt loop near-instantly instead of after the 300s default. + monkeypatch.setattr(acp_mod, "_PROMPT_TIMEOUT_SECONDS", 0.05) + + ex = AcpExecutor(AcpAgentConfig(command="x")) + + async def _noop(*_a: object, **_k: object) -> None: + return None + + async def _sess() -> str: + return "sess-1" + + # No real subprocess / handshake; nothing ever resolves the prompt future, + # so run_turn falls through to the timeout branch. + monkeypatch.setattr(ex, "_start_process", _noop) + monkeypatch.setattr(ex, "_ensure_initialized", _noop) + monkeypatch.setattr(ex, "_ensure_session", _sess) + monkeypatch.setattr(ex, "_send", _noop) + ex._proc = SimpleNamespace(returncode=None) # type: ignore[assignment] + + events = [event async for event in ex.run_turn([{"role": "user", "content": "hi"}], [], "sys")] + + # It surfaced the timeout... + assert any(isinstance(e, ExecutorError) for e in events) + # ...and, the point of the fix, left no stale future behind. + assert ex._pending == {}, "run_turn must drop its prompt future from _pending on timeout" diff --git a/tests/inner/test_goose_executor.py b/tests/inner/test_goose_executor.py index 63444788fa..10eae81020 100644 --- a/tests/inner/test_goose_executor.py +++ b/tests/inner/test_goose_executor.py @@ -1484,3 +1484,39 @@ def test_create_app_returns_fastapi() -> None: from omnigent.inner import goose_harness assert isinstance(goose_harness.create_app(), FastAPI) + + +@pytest.mark.asyncio +async def test_run_turn_timeout_does_not_leak_pending_future( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A turn that times out must not leave its request future in ``_pending``. + + The stdout reader pops ``req_id`` only on a matched response; the timeout / + EOF paths did not, leaking one stale future per silent turn on a long-lived + goose session. Drive a turn whose response never arrives and assert cleanup. + """ + from types import SimpleNamespace + + import omnigent.inner.goose_executor as goose_mod + + monkeypatch.setattr(goose_mod, "_PROMPT_TIMEOUT_SECONDS", 0.05) + ex = GooseExecutor(goose_path="goose") + + async def _noop(*_a: object, **_k: object) -> None: + return None + + async def _sess() -> str: + return "sess-1" + + monkeypatch.setattr(ex, "_start_process", _noop) + monkeypatch.setattr(ex, "_ensure_initialized", _noop) + monkeypatch.setattr(ex, "_ensure_session", _sess) + monkeypatch.setattr(ex, "_send", _noop) + ex._proc = SimpleNamespace(returncode=None) # type: ignore[assignment] + + events = [ + event async for event in ex.run_turn([{"role": "user", "content": "hi"}], [], "sys") + ] + assert any(isinstance(e, ExecutorError) for e in events) + assert ex._pending == {}, "run_turn must drop its prompt future from _pending on timeout" diff --git a/tests/inner/test_qwen_executor.py b/tests/inner/test_qwen_executor.py index d49b678254..81a69b7c28 100644 --- a/tests/inner/test_qwen_executor.py +++ b/tests/inner/test_qwen_executor.py @@ -2130,3 +2130,39 @@ async def test_ensure_initialized_image_capability_defaults_false() -> None: await executor._ensure_initialized() assert executor._initialized is True assert executor._image_supported is False + + +@pytest.mark.asyncio +async def test_run_turn_timeout_does_not_leak_pending_future( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A turn that times out must not leave its request future in ``_pending``. + + The stdout reader pops ``req_id`` only on a matched response; the timeout / + EOF paths did not, leaking one stale future per silent turn on a long-lived + qwen session. Drive a turn whose response never arrives and assert cleanup. + """ + from types import SimpleNamespace + + import omnigent.inner.qwen_executor as qwen_mod + + monkeypatch.setattr(qwen_mod, "_PROMPT_TIMEOUT_SECONDS", 0.05) + ex = QwenExecutor(qwen_path="qwen") + + async def _noop(*_a: object, **_k: object) -> None: + return None + + async def _sess() -> str: + return "sess-1" + + monkeypatch.setattr(ex, "_start_process", _noop) + monkeypatch.setattr(ex, "_ensure_initialized", _noop) + monkeypatch.setattr(ex, "_ensure_session", _sess) + monkeypatch.setattr(ex, "_send", _noop) + ex._proc = SimpleNamespace(returncode=None) # type: ignore[assignment] + + events = [ + event async for event in ex.run_turn([{"role": "user", "content": "hi"}], [], "sys") + ] + assert any(isinstance(e, ExecutorError) for e in events) + assert ex._pending == {}, "run_turn must drop its prompt future from _pending on timeout"