fix async queue

Signed-off-by: Richard Chien <[email protected]>

Author: stdrc

src/kimi_cli/app.py

@@ -22,6 +22,7 @@
 from kimi_cli.soul.agent import Runtime, load_agent
 from kimi_cli.soul.context import Context
 from kimi_cli.soul.kimisoul import KimiSoul
+from kimi_cli.utils.aioqueue import QueueShutDown
 from kimi_cli.utils.logging import StreamToLogger, logger
 from kimi_cli.utils.path import shorten_home
 from kimi_cli.wire import Wire, WireUISide
@@ -201,7 +202,7 @@ async def _ui_loop_fn(wire: Wire) -> None:
                 while True:
                     msg = await wire_ui.receive()
                     yield msg
-            except asyncio.QueueShutDown:
+            except QueueShutDown:
                 pass
             finally:
                 # stop consuming Wire messages

src/kimi_cli/soul/__init__.py

@@ -10,6 +10,7 @@
 
 from kosong.message import ContentPart
 
+from kimi_cli.utils.aioqueue import QueueShutDown
 from kimi_cli.utils.logging import logger
 from kimi_cli.wire import Wire
 from kimi_cli.wire.message import WireMessage
@@ -164,7 +165,7 @@ async def run_soul(
         wire.shutdown()
         try:
             await asyncio.wait_for(ui_task, timeout=0.5)
-        except asyncio.QueueShutDown:
+        except QueueShutDown:
             logger.debug("UI loop shut down")
             pass
         except TimeoutError:

src/kimi_cli/soul/approval.py

@@ -6,6 +6,7 @@
 from typing import Literal
 
 from kimi_cli.soul.toolset import get_current_tool_call_or_none
+from kimi_cli.utils.aioqueue import Queue
 from kimi_cli.utils.logging import logger
 from kimi_cli.wire.display import DisplayBlock
 
@@ -25,7 +26,7 @@ class Request:
 
 class Approval:
     def __init__(self, yolo: bool = False):
-        self._request_queue = asyncio.Queue[Request]()
+        self._request_queue = Queue[Request]()
         self._requests: dict[str, tuple[Request, asyncio.Future[bool]]] = {}
         self._yolo = yolo
         self._auto_approve_actions: set[str] = set()  # TODO: persist across sessions

src/kimi_cli/ui/acp/__init__.py

@@ -17,6 +17,7 @@
 from kimi_cli.soul import Soul, run_soul
 from kimi_cli.soul.kimisoul import KimiSoul
 from kimi_cli.soul.toolset import KimiToolset
+from kimi_cli.utils.aioqueue import QueueShutDown
 from kimi_cli.utils.logging import logger
 from kimi_cli.wire import Wire, WireUISide
 from kimi_cli.wire.message import WireMessage
@@ -92,7 +93,7 @@ async def _ui_loop_fn(wire: Wire) -> None:
                 while True:
                     msg = await wire_ui.receive()
                     yield msg
-            except asyncio.QueueShutDown:
+            except QueueShutDown:
                 pass
             finally:
                 # stop consuming Wire messages

src/kimi_cli/ui/print/visualize.py

@@ -1,4 +1,3 @@
-import asyncio
 from dataclasses import dataclass
 from typing import Protocol
 
@@ -8,6 +7,7 @@
 
 from kimi_cli.cli import OutputFormat
 from kimi_cli.soul.message import tool_result_to_message
+from kimi_cli.utils.aioqueue import QueueShutDown
 from kimi_cli.wire import Wire
 from kimi_cli.wire.message import StepBegin, StepInterrupted, WireMessage
 
@@ -162,7 +162,7 @@ async def visualize(output_format: OutputFormat, final_only: bool, wire: Wire) -
     while True:
         try:
             msg = await wire_ui.receive()
-        except asyncio.QueueShutDown:
+        except QueueShutDown:
             handler.flush()
             break
 

src/kimi_cli/ui/shell/keyboard.py

@@ -7,6 +7,8 @@
 from collections.abc import AsyncGenerator, Callable
 from enum import Enum, auto
 
+from kimi_cli.utils.aioqueue import Queue
+
 
 class KeyEvent(Enum):
     UP = auto()
@@ -20,7 +22,7 @@ class KeyEvent(Enum):
 
 async def listen_for_keyboard() -> AsyncGenerator[KeyEvent]:
     loop = asyncio.get_running_loop()
-    queue = asyncio.Queue[KeyEvent]()
+    queue = Queue[KeyEvent]()
     cancel_event = threading.Event()
 
     def emit(event: KeyEvent) -> None:

src/kimi_cli/ui/shell/replay.py

@@ -16,6 +16,7 @@
 from kimi_cli.ui.shell.console import console
 from kimi_cli.ui.shell.prompt import PROMPT_SYMBOL
 from kimi_cli.ui.shell.visualize import visualize
+from kimi_cli.utils.aioqueue import QueueShutDown
 from kimi_cli.utils.logging import logger
 from kimi_cli.utils.message import message_stringify
 from kimi_cli.wire import Wire
@@ -66,7 +67,7 @@ async def replay_recent_history(
             wire.soul_side.send(event)
             await asyncio.sleep(0)  # yield to UI loop
         wire.shutdown()
-        with contextlib.suppress(asyncio.QueueShutDown):
+        with contextlib.suppress(QueueShutDown):
             await ui_task
 
 

src/kimi_cli/ui/shell/visualize.py

@@ -19,6 +19,7 @@
 from kimi_cli.tools import extract_key_argument
 from kimi_cli.ui.shell.console import console
 from kimi_cli.ui.shell.keyboard import KeyEvent, listen_for_keyboard
+from kimi_cli.utils.aioqueue import QueueShutDown
 from kimi_cli.utils.rich.columns import BulletColumns
 from kimi_cli.utils.rich.markdown import Markdown
 from kimi_cli.wire import WireUISide
@@ -350,7 +351,7 @@ def keyboard_handler(event: KeyEvent) -> None:
                 while True:
                     try:
                         msg = await wire.receive()
-                    except asyncio.QueueShutDown:
+                    except QueueShutDown:
                         self.cleanup(is_interrupt=False)
                         live.update(self.compose())
                         break

src/kimi_cli/ui/wire/__init__.py

@@ -9,6 +9,7 @@
 
 from kimi_cli.soul import LLMNotSet, LLMNotSupported, MaxStepsReached, RunCancelled, Soul, run_soul
 from kimi_cli.soul.kimisoul import KimiSoul
+from kimi_cli.utils.aioqueue import Queue, QueueShutDown
 from kimi_cli.utils.logging import logger
 from kimi_cli.wire import Wire
 from kimi_cli.wire.message import ApprovalRequest, Request
@@ -37,7 +38,7 @@ def __init__(self, soul: Soul):
 
         # outward
         self._write_task: asyncio.Task[None] | None = None
-        self._write_queue: asyncio.Queue[JSONRPCOutMessage] = asyncio.Queue()
+        self._write_queue: Queue[JSONRPCOutMessage] = Queue()
 
         # inward
         self._dispatch_tasks: set[asyncio.Task[None]] = set()
@@ -65,7 +66,7 @@ async def _write_loop(self) -> None:
             while True:
                 try:
                     msg = await self._write_queue.get()
-                except asyncio.QueueShutDown:
+                except QueueShutDown:
                     logger.debug("Send queue shut down, stopping Wire server write loop")
                     break
                 self._writer.write(msg.model_dump_json().encode("utf-8") + b"\n")
@@ -141,7 +142,7 @@ async def _dispatch_msg(self, msg: JSONRPCInMessage) -> None:
     async def _send_msg(self, msg: JSONRPCOutMessage) -> None:
         try:
             await self._write_queue.put(msg)
-        except asyncio.QueueShutDown:
+        except QueueShutDown:
             logger.error("Send queue shut down; dropping message: {msg}", msg=msg)
 
     @property

src/kimi_cli/utils/aioqueue.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import asyncio
+import sys
+
+if sys.version_info >= (3, 13):
+    QueueShutDown = asyncio.QueueShutDown  # type: ignore[assignment]
+
+    class Queue[T](asyncio.Queue[T]):
+        """Asyncio Queue with shutdown support."""
+
+else:
+
+    class QueueShutDown(Exception):
+        """Raised when operating on a shut down queue."""
+
+    class _Shutdown:
+        """Sentinel for queue shutdown."""
+
+    _SHUTDOWN = _Shutdown()
+
+    class Queue[T](asyncio.Queue[T | _Shutdown]):
+        """Asyncio Queue with shutdown support for Python < 3.13."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self._shutdown = False
+
+        def shutdown(self, immediate: bool = False) -> None:
+            if self._shutdown:
+                return
+            self._shutdown = True
+            if immediate:
+                self._queue.clear()
+
+            getters = list(getattr(self, "_getters", []))
+            count = max(1, len(getters))
+            self._enqueue_shutdown(count)
+
+        def _enqueue_shutdown(self, count: int) -> None:
+            for _ in range(count):
+                try:
+                    super().put_nowait(_SHUTDOWN)
+                except asyncio.QueueFull:
+                    self._queue.clear()
+                    super().put_nowait(_SHUTDOWN)
+
+        async def get(self) -> T:
+            if self._shutdown and self.empty():
+                raise QueueShutDown
+            item = await super().get()
+            if isinstance(item, _Shutdown):
+                raise QueueShutDown
+            return item
+
+        def get_nowait(self) -> T:
+            if self._shutdown and self.empty():
+                raise QueueShutDown
+            item = super().get_nowait()
+            if isinstance(item, _Shutdown):
+                raise QueueShutDown
+            return item
+
+        async def put(self, item: T) -> None:
+            if self._shutdown:
+                raise QueueShutDown
+            await super().put(item)
+
+        def put_nowait(self, item: T) -> None:
+            if self._shutdown:
+                raise QueueShutDown
+            super().put_nowait(item)