fix(llm): keep async agents responsive during model requests

ms_agent/agent/llm_agent.py

@@ -20,6 +20,7 @@
 from ms_agent.callbacks import Callback, callbacks_mapping
 from ms_agent.knowledge_search import SirchmunkSearch
 from ms_agent.llm import multimodal
+from ms_agent.llm.io import run_in_llm_executor
 from ms_agent.llm.llm import LLM
 from ms_agent.llm.message_text import (append_text, flatten_message_text,
                                        prepend_text)
@@ -2040,6 +2041,49 @@ def _append_task_notifications(self,
         messages.append(Message(role='user', content=body))
         return messages
 
+    @staticmethod
+    def _close_llm_stream(llm, response=None) -> None:
+        """Best-effort cleanup, including a response returned after cancellation."""
+        for close in (getattr(llm, 'interrupt',
+                              None), getattr(response, 'close', None)):
+            if callable(close):
+                try:
+                    close()
+                except Exception:  # teardown must not mask the original error
+                    pass
+
+    async def _generate_response(self, messages, tools):
+        """Open a synchronous model request without blocking the event loop.
+
+        Cancellation cannot stop a Python worker thread. If the request returns
+        after its caller has gone, the worker closes that response instead of
+        abandoning it. The lock covers only ownership transfer, never I/O.
+        """
+        lock = threading.Lock()
+        llm = self.llm
+        abandoned = False
+        response = None
+
+        def generate():
+            nonlocal response
+            result = llm.generate(messages, tools=tools)
+            with lock:
+                discard = abandoned
+                if not discard:
+                    response = result
+            if discard:
+                self._close_llm_stream(llm, result)
+            return result
+
+        try:
+            return await run_in_llm_executor(generate)
+        except asyncio.CancelledError:
+            with lock:
+                abandoned = True
+                result = response
+            self._close_llm_stream(llm, result)
+            raise
+
     # retry_if: a hard 4xx (bad payload, content filter, auth) is a verdict on
     # the request, not a transient fault — retrying it 5× only adds ~40s of
     # backoff before the same failure surfaces.
@@ -2118,9 +2162,10 @@ async def step(
                 # ui.events.ToolCallComposing).
                 _composing: Dict[int, int] = {}
                 _reported_images = False
-                _gen = self.llm.generate(messages, tools=tools)
-                _loop = asyncio.get_running_loop()
+                _llm = self.llm
+                _gen = await self._generate_response(messages, tools)
                 _NO_MORE = object()
+                _stopped = threading.Event()
 
                 def _next_chunk(_g=_gen):
                     # Step the BLOCKING sync LLM stream off the event loop, so
@@ -2133,10 +2178,15 @@ def _next_chunk(_g=_gen):
                         return next(_g)
                     except StopIteration:
                         return _NO_MORE
+                    finally:
+                        # A cancelled await leaves next() running in its worker.
+                        # Close the iterator there once that read has unwound.
+                        if _stopped.is_set():
+                            self._close_llm_stream(_llm, _g)
 
                 try:
                     while True:
-                        _chunk = await _loop.run_in_executor(None, _next_chunk)
+                        _chunk = await run_in_llm_executor(_next_chunk)
                         if _chunk is _NO_MORE:
                             break
                         _response_message = _chunk
@@ -2189,24 +2239,13 @@ def _next_chunk(_g=_gen):
                         messages[-1] = _response_message
                         yield messages
                 finally:
+                    _stopped.set()
                     if not _reported_images:
                         # A turn that produced nothing still attached images, and
                         # what became of them is still worth saying.
                         self._record_image_deliveries(messages)
                         _reported_images = True
-                    # Turn abandoned mid-stream (client disconnect / stop): ask
-                    # the provider to close the live upstream response so the
-                    # server stops generating, instead of leaving it to run to
-                    # completion into a dropped connection. Only the data-driven
-                    # provider layer implements interrupt(); the legacy LLM does
-                    # not, so this is a no-op there (unchanged). Harmless on a
-                    # normal finish (the stream is already exhausted).
-                    _interrupt = getattr(self.llm, 'interrupt', None)
-                    if callable(_interrupt):
-                        try:
-                            _interrupt()
-                        except Exception:  # noqa: BLE001 - teardown never raises
-                            pass
+                    self._close_llm_stream(_llm, _gen)
                 if self.stream_output:
                     if _printed_reasoning_header and not _printed_reasoning_footer:
                         self._emit_reasoning_end()
@@ -2222,7 +2261,8 @@ def _next_chunk(_g=_gen):
 
                     self._emit_content_end()
             else:
-                _response_message = self.llm.generate(messages, tools=tools)
+                _response_message = await self._generate_response(
+                    messages, tools)
                 if self.show_reasoning:
                     reasoning_text = (
                         getattr(_response_message, 'reasoning_content', '')

ms_agent/llm/anthropic_llm.py

@@ -5,6 +5,7 @@
 from typing import Any, Dict, Generator, Iterator, List, Optional, Union
 
 from ms_agent.llm import LLM
+from ms_agent.llm.io import OpenedStream, interrupt_stream
 from ms_agent.llm.thinking import create_with_thinking_fallback
 from ms_agent.llm.utils import Message, Tool, ToolCall
 from ms_agent.utils import assert_package_exist, get_logger, retry
@@ -167,6 +168,14 @@ def __init__(
 
         self.args: Dict = OmegaConf.to_container(
             getattr(config, 'generation_config', DictConfig({})))
+        self._active_stream = None
+
+    def interrupt(self) -> None:
+        """Close the owned HTTP stream, including before iteration starts."""
+        stream = self._active_stream
+        interrupt_stream(stream)
+        if self._active_stream is stream:
+            self._active_stream = None
 
     def format_tools(self,
                      tools: Optional[List[Tool]]) -> Optional[List[Dict]]:
@@ -284,7 +293,9 @@ def _call_llm(self,
         def _send(**call):
             call.setdefault('model', self.model)
             if stream:
-                return self.client.messages.stream(**call)
+                opened = OpenedStream(self.client.messages.stream(**call))
+                self._active_stream = opened
+                return opened
             return self.client.messages.create(**call)
 
         # This legacy engine owned no repair at all: a model that cannot think
@@ -333,53 +344,58 @@ def _stream_format_output_message(self,
         )
         tool_call_id_map = {}  # index -> tool_call_id (用于去重 yield)
         with stream_manager as stream:
-            full_content = ''
-            full_thinking = ''
-            for event in stream:
-                event_type = getattr(event, 'type')
-                if event_type == 'message_start':
-                    msg = event.message
-                    current_message.id = msg.id
-                    tool_call_id_map = {}
-                    yield current_message
-                elif event_type == 'content_block_delta':
-                    if event.delta.type == 'thinking_delta':
-                        full_thinking += event.delta.thinking
-                        current_message.reasoning_content = full_thinking
-                    elif event.delta.type == 'text_delta':
-                        full_content += event.delta.text
+            self._active_stream = stream
+            try:
+                full_content = ''
+                full_thinking = ''
+                for event in stream:
+                    event_type = getattr(event, 'type')
+                    if event_type == 'message_start':
+                        msg = event.message
+                        current_message.id = msg.id
+                        tool_call_id_map = {}
+                        yield current_message
+                    elif event_type == 'content_block_delta':
+                        if event.delta.type == 'thinking_delta':
+                            full_thinking += event.delta.thinking
+                            current_message.reasoning_content = full_thinking
+                        elif event.delta.type == 'text_delta':
+                            full_content += event.delta.text
+                            current_message.content = full_content
+                        yield current_message
+                    elif event_type == 'message_stop':
+                        final_msg = getattr(event, 'message')
+                        full_content = ''
+                        used_tool_call_ids = set()
+                        for idx, block in enumerate(event.message.content):
+                            if block is None:
+                                continue
+                            if block.type == 'text':
+                                full_content += block.text
+                            elif block.type == 'tool_use':
+                                tool_call_id = tool_call_id_map.get(idx)
+                                tool_call = ToolCall(
+                                    id=tool_call_id,
+                                    index=len(current_message.tool_calls),
+                                    type='function',
+                                    tool_name=block.name,
+                                    arguments=block.input,
+                                )
+                                current_message.tool_calls.append(tool_call)
+                                used_tool_call_ids.add(tool_call_id)
                         current_message.content = full_content
-                    yield current_message
-                elif event_type == 'message_stop':
-                    final_msg = getattr(event, 'message')
-                    full_content = ''
-                    used_tool_call_ids = set()
-                    for idx, block in enumerate(event.message.content):
-                        if block is None:
-                            continue
-                        if block.type == 'text':
-                            full_content += block.text
-                        elif block.type == 'tool_use':
-                            tool_call_id = tool_call_id_map.get(idx)
-                            tool_call = ToolCall(
-                                id=tool_call_id,
-                                index=len(current_message.tool_calls),
-                                type='function',
-                                tool_name=block.name,
-                                arguments=block.input,
-                            )
-                            current_message.tool_calls.append(tool_call)
-                            used_tool_call_ids.add(tool_call_id)
-                    current_message.content = full_content
-                    current_message.partial = False
-                    current_message.completion_tokens = getattr(
-                        final_msg.usage, 'output_tokens',
-                        current_message.completion_tokens)
-                    current_message.prompt_tokens = getattr(
-                        final_msg.usage, 'input_tokens',
-                        current_message.prompt_tokens)
-
-                    yield current_message
+                        current_message.partial = False
+                        current_message.completion_tokens = getattr(
+                            final_msg.usage, 'output_tokens',
+                            current_message.completion_tokens)
+                        current_message.prompt_tokens = getattr(
+                            final_msg.usage, 'input_tokens',
+                            current_message.prompt_tokens)
+
+                        yield current_message
+            finally:
+                if self._active_stream is stream:
+                    self._active_stream = None
 
     @staticmethod
     def _format_output_message(completion) -> Message:

ms_agent/llm/io.py

@@ -0,0 +1,67 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Keep blocking model I/O out of the application's default executor."""
+import asyncio
+import contextvars
+import socket
+from concurrent.futures import ThreadPoolExecutor
+
+# Shared, bounded, and started lazily by ThreadPoolExecutor. Its lifetime is
+# the process's; the executor joins its workers on process exit. Slow models
+# must not occupy the host's executor for session storage and cancellation.
+_executor = ThreadPoolExecutor(thread_name_prefix='ms-agent-llm')
+
+
+async def run_in_llm_executor(func, *args):
+    context = contextvars.copy_context()
+    return await asyncio.get_running_loop().run_in_executor(
+        _executor, context.run, func, *args)
+
+
+def interrupt_stream(stream):
+    """Interrupt an owned HTTP/1 response before closing its SDK stream.
+
+    A socket close in another thread need not wake a blocked recv(). Shutdown
+    does. Never shut down an HTTP/2 connection, which may carry other requests.
+    Custom transports without HTTPX's network extension retain normal close.
+    """
+    try:
+        response = getattr(stream, 'response', None)
+        if (response is not None and not response.is_closed
+                and response.http_version in ('HTTP/1.0', 'HTTP/1.1')):
+            network = response.extensions.get('network_stream')
+            sock = network.get_extra_info('socket') if network else None
+            if sock is not None:
+                sock.shutdown(socket.SHUT_RDWR)
+    except Exception:  # shutdown is best effort; always attempt SDK cleanup
+        pass
+    try:
+        if stream is not None:
+            stream.close()
+    except Exception:
+        pass
+
+
+class OpenedStream:
+    """Open a lazy stream manager while the request still owns cancellation.
+
+    Anthropic sends the HTTP request in __enter__, not messages.stream().
+    Opening it during generate() lets LLMAgent close a late response before
+    reading any body. The consumer retains the manager's context protocol.
+    """
+
+    def __init__(self, manager):
+        self._manager = manager
+        self._stream = manager.__enter__()
+
+    def __enter__(self):
+        return self._stream
+
+    @property
+    def response(self):
+        return getattr(self._stream, 'response', None)
+
+    def __exit__(self, *exc):
+        return self._manager.__exit__(*exc)
+
+    def close(self):
+        self.__exit__(None, None, None)