diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 34098cdc5e21..37113b61acf0 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -2473,7 +2473,11 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= request_client_holder = {"client": None, "diag": None, "owner_tid": None} # Transport kind of the registered request client — see the non-streaming # variant. Routes _close_request_client_once to anthropic vs openai abort/ - # close helpers (#67142). + # close helpers (#67142). ``kind="stream"`` registers a per-request + # *stream handle* instead of a client — used under the MoA facade, whose + # singleton client has no per-request sockets to abort + # (_abort_request_openai_client is a no-op on it), so interrupts must + # close the stream object itself (#57354). request_client_kind = {"value": "openai"} request_client_lock = threading.Lock() # Request-local cancellation flag — see interruptible_api_call for the full @@ -2493,6 +2497,44 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= request_client_holder["owner_tid"] = threading.get_ident() return client + def _stream_close_callable(stream): + close = getattr(stream, "close", None) + if callable(close): + return close + response = getattr(stream, "response", None) + close = getattr(response, "close", None) + if callable(close): + return close + return None + + def _set_request_stream_handle(stream): + # Register the per-request *stream* under kind="stream" so an + # interrupt closes the stream handle itself. Under the MoA facade the + # registered "client" is the shared facade singleton whose + # per-request abort helpers are no-ops, leaving the underlying HTTP + # stream open until the provider drained it (#57354). + if _stream_close_callable(stream) is None: + return stream + with request_client_lock: + request_client_holder["client"] = stream + request_client_kind["value"] = "stream" + request_client_holder["owner_tid"] = threading.get_ident() + return stream + + def _close_request_stream_handle(stream, reason: str) -> None: + close = _stream_close_callable(stream) + if close is None: + return + try: + close() + logger.info("Streaming response handle closed (%s)", reason) + except Exception as exc: + logger.debug( + "Streaming response handle close failed (%s): %s", + reason, + exc, + ) + def _close_request_client_once(reason: str) -> None: # See #29507 explanation in the non-streaming variant above. A # stranger thread (the interrupt-check / stale-stream detector loop) @@ -2500,9 +2542,15 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # so the worker thread retains ownership of the FD release. with request_client_lock: request_client = request_client_holder.get("client") + request_kind = request_client_kind.get("value", "openai") owner_tid = request_client_holder.get("owner_tid") + # A registered stream handle (kind="stream", MoA facade path) is + # safe to close from any thread — closing IS the abort — so the + # stranger-thread ownership carve-out only applies to real + # per-request clients (#57354). stranger_thread = ( - request_client is not None + request_kind != "stream" + and request_client is not None and owner_tid is not None and owner_tid != threading.get_ident() ) @@ -2511,8 +2559,9 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= request_client_holder["owner_tid"] = None if request_client is None: return - kind = request_client_kind.get("value", "openai") - if kind == "anthropic_messages": + if request_kind == "stream": + _close_request_stream_handle(request_client, reason) + elif request_kind == "anthropic_messages": if stranger_thread: agent._abort_request_anthropic_client(request_client, reason=reason) else: @@ -2691,6 +2740,11 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= _diag = agent._stream_diag_init() request_client_holder["diag"] = _diag stream = request_client.chat.completions.create(**stream_kwargs) + if agent.provider == "moa": + # The MoA facade is a shared singleton — abort/close of the + # registered client is a no-op, so register the stream handle + # itself for interrupt teardown (#57354). + stream = _set_request_stream_handle(stream) # Claim the delta sink for THIS attempt (#65991). If a prior attempt's # stream is somehow still alive (a stale-stream reconnect whose socket # abort raced), this claim supersedes it so its late chunks are fenced diff --git a/tests/run_agent/test_streaming.py b/tests/run_agent/test_streaming.py index 55a0b9c008e7..c73f441141ec 100644 --- a/tests/run_agent/test_streaming.py +++ b/tests/run_agent/test_streaming.py @@ -3,6 +3,7 @@ Tests the unified streaming API call, delta callbacks, tool-call suppression, provider fallback, and CLI streaming display. """ +import threading from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -718,6 +719,68 @@ class TestStreamingFallback: assert agent._disable_streaming is True assert deltas == [] + @patch("run_agent.AIAgent._abort_request_openai_client") + @patch("run_agent.AIAgent._close_request_openai_client") + @patch("run_agent.AIAgent._create_request_openai_client") + def test_moa_interrupt_closes_stream_handle( + self, mock_create, mock_close_openai, mock_abort_openai + ): + """MoA interrupts must close the per-request stream, not the facade client.""" + from run_agent import AIAgent + + class _BlockingClosableStream: + def __init__(self): + self.entered = threading.Event() + self.closed = threading.Event() + self.close_calls = 0 + + def __iter__(self): + return self + + def __next__(self): + self.entered.set() + if not self.closed.wait(timeout=5): + raise TimeoutError("MoA test stream was not closed") + raise RuntimeError("stream closed") + + def close(self): + self.close_calls += 1 + self.closed.set() + + stream = _BlockingClosableStream() + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = stream + mock_create.return_value = mock_client + + agent = AIAgent( + model="default", + provider="moa", + api_key="test-key", + base_url="moa://local", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + agent.api_mode = "chat_completions" + agent._interrupt_requested = False + agent.client = mock_client + + def _request_interrupt(): + assert stream.entered.wait(timeout=2) + agent._interrupt_requested = True + + interrupter = threading.Thread(target=_request_interrupt, daemon=True) + interrupter.start() + + with pytest.raises(InterruptedError): + agent._interruptible_streaming_api_call({"model": "default", "messages": []}) + + assert stream.closed.wait(timeout=2) + assert stream.close_calls == 1 + mock_create.assert_called_once() + mock_close_openai.assert_not_called() + mock_abort_openai.assert_not_called() + @patch("run_agent.AIAgent._create_request_openai_client") @patch("run_agent.AIAgent._close_request_openai_client") def test_stream_error_propagates_original(self, mock_close, mock_create):