comate-cli 0.2.0__tar.gz → 0.2.2__tar.gz

{comate_cli-0.2.0 → comate_cli-0.2.2}/.gitignore

@@ -8,6 +8,7 @@ set_env.sh
 # C extensions
 *.so
 
+.claude/
 # Distribution / packaging
 .Python
 build/
@@ -218,4 +219,7 @@ __marimo__/
 .streamlit/secrets.toml
 
 .env
-./tmp/
+./tmp/
+
+# OS generated files
+.DS_Store

{comate_cli-0.2.0 → comate_cli-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comate-cli
-Version: 0.2.0
+Version: 0.2.2
 Summary: Comate terminal CLI built on comate-agent-sdk
 Project-URL: Homepage, https://github.com/AndyLee1024/agent-sdk
 Project-URL: Repository, https://github.com/AndyLee1024/agent-sdk

{comate_cli-0.2.0 → comate_cli-0.2.2}/comate_cli/mcp_cli.py

@@ -23,8 +23,13 @@ class McpCliError(ValueError):
         self.exit_code = int(exit_code)
 
 
+class _McpArgumentParser(argparse.ArgumentParser):
+    def error(self, message: str) -> None:
+        raise McpCliError(message)
+
+
 def _build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
+    parser = _McpArgumentParser(
         prog="comate mcp",
         description="Configure and manage MCP servers",
         formatter_class=argparse.RawDescriptionHelpFormatter,
@@ -143,11 +148,11 @@ def _parse_env_entries(entries: list[str]) -> dict[str, str]:
     return env
 
 
-def _clean_remainder_args(raw_args: list[str]) -> list[str]:
-    args = list(raw_args)
-    if args and args[0] == "--":
-        return args[1:]
-    return args
+def _split_trailing_command_args(argv: list[str]) -> tuple[list[str], list[str]]:
+    if "--" not in argv:
+        return list(argv), []
+    separator_index = argv.index("--")
+    return list(argv[:separator_index]), list(argv[separator_index + 1 :])
 
 
 def _build_add_server_config(args: argparse.Namespace) -> McpServerConfig:
@@ -155,9 +160,7 @@ def _build_add_server_config(args: argparse.Namespace) -> McpServerConfig:
     header_entries = [str(item) for item in list(args.header or [])]
     env_entries = [str(item) for item in list(args.env or [])]
     command_or_url = str(args.command_or_url or "").strip()
-    remainder = _clean_remainder_args(
-        [str(item) for item in list(getattr(args, "extra_args", []) or [])]
-    )
+    remainder = [str(item) for item in list(getattr(args, "extra_args", []) or [])]
 
     if transport == "http":
         if env_entries:
@@ -177,6 +180,9 @@ def _build_add_server_config(args: argparse.Namespace) -> McpServerConfig:
 
     if header_entries:
         raise McpCliError("--header is only supported for http transport")
+    stdio_args = list(remainder)
+    if not command_or_url and stdio_args:
+        command_or_url = stdio_args.pop(0)
     if not command_or_url:
         raise McpCliError("stdio transport requires command after <name>")
 
@@ -184,8 +190,8 @@ def _build_add_server_config(args: argparse.Namespace) -> McpServerConfig:
     cfg = {
         "command": command_or_url,
     }
-    if remainder:
-        cfg["args"] = remainder
+    if stdio_args:
+        cfg["args"] = stdio_args
     if env:
         cfg["env"] = env
     return cfg  # type: ignore[return-value]
@@ -365,16 +371,17 @@ def _cmd_get(args: argparse.Namespace, *, project_root: PathInput | None) -> Non
 
 def run_mcp_command(argv: list[str], *, project_root: PathInput | None = None) -> None:
     parser = _build_parser()
-    parsed, extra_args = parser.parse_known_args(argv)
+    parser_argv, trailing_args = _split_trailing_command_args(argv)
+    parsed = parser.parse_args(parser_argv)
     command = str(getattr(parsed, "command", "") or "").strip()
     if not command:
         parser.print_help(sys.stdout)
         return
-    if command != "add" and extra_args:
-        joined = " ".join(extra_args)
+    if command != "add" and trailing_args:
+        joined = " ".join(["--", *trailing_args])
         raise McpCliError(f"Unrecognized arguments: {joined}")
     if command == "add":
-        setattr(parsed, "extra_args", list(extra_args))
+        setattr(parsed, "extra_args", list(trailing_args))
 
     handlers = {
         "add": _cmd_add,

{comate_cli-0.2.0 → comate_cli-0.2.2}/comate_cli/terminal_agent/app.py

@@ -15,8 +15,7 @@ from pathlib import Path
 
 from rich.console import Console
 
-from comate_agent_sdk import Agent
-from comate_agent_sdk.agent import AgentConfig, ChatSession
+from comate_agent_sdk.agent import Agent, AgentConfig, ChatSession
 from comate_agent_sdk.context import EnvOptions
 from comate_agent_sdk.tools import tool
 
@@ -230,16 +229,20 @@ def _format_resume_hint(session_id: str | None) -> str | None:
 
 async def _preload_mcp_in_tui(session: ChatSession) -> None:
     """在 TUI 内异步加载 MCP，初始化阶段不输出 scrollback 文案。"""
-    if not bool(session._agent.options.mcp_enabled):
+    runtime = session.runtime
+    if not bool(runtime.options.mcp_enabled):
         return
 
     try:
-        await session._agent.ensure_mcp_tools_loaded()
+        preload_task = runtime.start_mcp_preload()
+        if preload_task is None:
+            return
+        await preload_task
     except Exception as e:
         logger.warning(f"MCP init failed: {e}", exc_info=True)
         return
 
-    mgr = session._agent._mcp_manager
+    mgr = runtime._mcp_manager
     if mgr is None:
         return
 
@@ -311,7 +314,7 @@ async def run(
 
     # 配置 TUI logging handler（将 SDK 日志输出到 TUI）
     from comate_cli.terminal_agent.logging_adapter import setup_tui_logging
-    setup_tui_logging(renderer)
+    logging_session = setup_tui_logging(renderer)
 
     tui = TerminalAgentTUI(session, status_bar, renderer)
     tui.add_resume_history(mode)
@@ -333,6 +336,7 @@ async def run(
                 exc_info=True,
             )
     finally:
+        logging_session.close()
         if active_session is session:
             await _graceful_shutdown(active_session)
         else:

{comate_cli-0.2.0 → comate_cli-0.2.2}/comate_cli/terminal_agent/logging_adapter.py

@@ -2,6 +2,8 @@
 from __future__ import annotations
 
 import logging
+import os
+import sys
 from typing import TYPE_CHECKING
 
 from prompt_toolkit.application import run_in_terminal
@@ -106,42 +108,119 @@ class TUILoggingHandler(logging.Handler):
                 # DummyApplication，直接调用
                 _append()
             else:
-                # 真实 app，使用 run_in_terminal
-                run_in_terminal(_append, in_executor=False)
+                app.create_background_task(
+                    run_in_terminal(_append, in_executor=False)
+                )
         except Exception:
             # 没有 app 或导入失败，直接调用
             _append()
 
 
-def setup_tui_logging(renderer: EventRenderer) -> None:
-    """统一日志初始化：文件 + TUI 双通道。
-
-    - 所有日志（含 traceback）写入 ~/.comate/logs/agent.log（RotatingFileHandler）
-    - WARNING/ERROR 以用户友好格式显示在 TUI scrollback
-    - 清除 root logger 默认的 stderr handler，阻止 traceback 泄漏到终端
-    """
-    import os
+class _DropTerminalStreamFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        del record
+        return False
+
+
+class TUILoggingSession:
+    def __init__(
+        self,
+        *,
+        root_logger: logging.Logger,
+        added_handlers: list[logging.Handler],
+        muted_handlers: list[tuple[logging.Handler, logging.Filter]],
+        previous_level: int,
+    ) -> None:
+        self._root_logger = root_logger
+        self._added_handlers = list(added_handlers)
+        self._muted_handlers = list(muted_handlers)
+        self._previous_level = previous_level
+        self._closed = False
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+
+        for handler, log_filter in self._muted_handlers:
+            try:
+                handler.removeFilter(log_filter)
+            except Exception:
+                continue
+
+        for handler in self._added_handlers:
+            try:
+                self._root_logger.removeHandler(handler)
+            except Exception:
+                continue
+            try:
+                handler.close()
+            except Exception:
+                continue
+
+        self._root_logger.setLevel(self._previous_level)
+
+
+def _build_file_handler() -> logging.Handler:
     from logging.handlers import RotatingFileHandler
 
-    # 1. 日志文件 handler（完整调试信息，含 traceback）
     log_dir = os.path.join(os.path.expanduser("~"), ".comate", "logs")
     os.makedirs(log_dir, exist_ok=True)
     log_path = os.path.join(log_dir, "agent.log")
     file_handler = RotatingFileHandler(
-        log_path, maxBytes=10 * 1024 * 1024, backupCount=3, encoding="utf-8",
+        log_path,
+        maxBytes=10 * 1024 * 1024,
+        backupCount=3,
+        encoding="utf-8",
     )
     file_handler.setLevel(logging.DEBUG)
     file_handler.setFormatter(
         logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
     )
+    return file_handler
+
+
+def _should_mute_terminal_stream(handler: logging.Handler) -> bool:
+    if not isinstance(handler, logging.StreamHandler):
+        return False
+    stream = getattr(handler, "stream", None)
+    return stream in {sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__}
+
 
-    # 2. 清除 root logger 的所有 handler，阻止 stderr 泄漏
+def setup_tui_logging(renderer: EventRenderer) -> TUILoggingSession:
+    """统一日志初始化：文件 + TUI 双通道。
+
+    - 所有日志（含 traceback）写入 ~/.comate/logs/agent.log（RotatingFileHandler）
+    - WARNING/ERROR 以用户友好格式显示在 TUI scrollback
+    - 临时静默 root logger 上写往 stdout/stderr 的 stream handler，避免污染 TUI
+    """
     root = logging.getLogger()
-    root.handlers.clear()
+    previous_level = root.level
+
+    # 1. 日志文件 handler（完整调试信息，含 traceback）
+    file_handler = _build_file_handler()
     root.addHandler(file_handler)
     root.setLevel(logging.DEBUG)
 
-    # 3. TUI handler 挂到 root（覆盖所有命名空间的 WARNING/ERROR）
+    # 2. TUI handler 挂到 root（覆盖所有命名空间的 WARNING/ERROR）
     tui_handler = TUILoggingHandler(renderer)
     tui_handler.setLevel(logging.WARNING)
     root.addHandler(tui_handler)
+
+    # 3. 临时静默直写终端的 handler，避免污染 prompt_toolkit UI
+    muted_handlers: list[tuple[logging.Handler, logging.Filter]] = []
+    for handler in list(root.handlers):
+        if handler in {file_handler, tui_handler}:
+            continue
+        if not _should_mute_terminal_stream(handler):
+            continue
+        log_filter = _DropTerminalStreamFilter()
+        handler.addFilter(log_filter)
+        muted_handlers.append((handler, log_filter))
+
+    return TUILoggingSession(
+        root_logger=root,
+        added_handlers=[file_handler, tui_handler],
+        muted_handlers=muted_handlers,
+        previous_level=previous_level,
+    )

{comate_cli-0.2.0 → comate_cli-0.2.2}/comate_cli/terminal_agent/rpc_protocol.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 from dataclasses import asdict, is_dataclass
-from typing import Any
+from typing import Any, TypeGuard
 
 
 class ErrorCodes:
@@ -30,6 +30,13 @@ class JSONRPCProtocolError(Exception):
         self.data = data
 
 
+JSONRPCRequestId = str | int
+
+
+def is_jsonrpc_request_id(raw_value: Any) -> TypeGuard[JSONRPCRequestId]:
+    return isinstance(raw_value, (str, int)) and not isinstance(raw_value, bool)
+
+
 def parse_jsonrpc_message(raw_line: str) -> dict[str, Any]:
     try:
         parsed = json.loads(raw_line)
@@ -52,13 +59,13 @@ def parse_jsonrpc_message(raw_line: str) -> dict[str, Any]:
     return parsed
 
 
-def build_success_response(request_id: str | int, result: Any) -> dict[str, Any]:
+def build_success_response(request_id: JSONRPCRequestId, result: Any) -> dict[str, Any]:
     return {"jsonrpc": "2.0", "id": request_id, "result": _to_jsonable(result)}
 
 
 def build_error_response(
     *,
-    request_id: str | int | None,
+    request_id: JSONRPCRequestId | None,
     code: int,
     message: str,
     data: Any = None,
@@ -82,8 +89,8 @@ def build_event_notification(event: Any) -> dict[str, Any]:
     return {"jsonrpc": "2.0", "method": "event", "params": params}
 
 
-def _coerce_request_id(raw_value: Any) -> str | int | None:
-    if isinstance(raw_value, str | int):
+def _coerce_request_id(raw_value: Any) -> JSONRPCRequestId | None:
+    if is_jsonrpc_request_id(raw_value):
         return raw_value
     return None
 

{comate_cli-0.2.0 → comate_cli-0.2.2}/comate_cli/terminal_agent/rpc_stdio.py

@@ -14,6 +14,7 @@ from comate_cli.terminal_agent.rpc_protocol import (
     build_error_response,
     build_event_notification,
     build_success_response,
+    is_jsonrpc_request_id,
     parse_jsonrpc_message,
 )
 
@@ -33,8 +34,13 @@ class StdioRPCBridge:
         self._init_error: Exception | None = None
         self._active_prompt_request_id: str | int | None = None
         self._active_prompt_result: asyncio.Future[dict[str, Any]] | None = None
+        self._finalize_prompt_task: asyncio.Task[None] | None = None
+        self._stdin_queue: asyncio.Queue[str | None] | None = None
+        self._stdin_reader_fd: int | None = None
+        self._stdin_error: Exception | None = None
 
     async def run(self) -> None:
+        self._install_stdin_reader()
         self._event_pump_task = asyncio.create_task(
             self._run_event_pump(),
             name="rpc-event-pump",
@@ -42,10 +48,29 @@ class StdioRPCBridge:
         try:
             while not self._closing:
                 pump_task = self._event_pump_task
-                if pump_task is not None and pump_task.done():
+                input_task = asyncio.create_task(
+                    self._read_next_input_line(),
+                    name="rpc-stdin-read",
+                )
+                wait_tasks: set[asyncio.Task[Any]] = {input_task}
+                if pump_task is not None:
+                    wait_tasks.add(pump_task)
+
+                done, _pending = await asyncio.wait(
+                    wait_tasks,
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+
+                if pump_task is not None and pump_task in done:
+                    input_task.cancel()
+                    try:
+                        await input_task
+                    except asyncio.CancelledError:
+                        pass
                     await pump_task
+                    break
 
-                raw_line = await asyncio.to_thread(sys.stdin.readline)
+                raw_line = input_task.result()
                 if raw_line == "":
                     break
 
@@ -56,12 +81,14 @@ class StdioRPCBridge:
                 await self._handle_incoming_line(line)
         finally:
             self._closing = True
+            self._remove_stdin_reader()
             await self._cancel_active_prompt()
             active_future = self._active_prompt_result
             if active_future is not None and not active_future.done():
                 active_future.set_result(
                     {"status": "cancelled", "stop_reason": "cancelled"}
                 )
+            await self._await_finalize_prompt_task()
             pump_task = self._event_pump_task
             if pump_task is not None:
                 pump_task.cancel()
@@ -92,7 +119,7 @@ class StdioRPCBridge:
         if method is None:
             await self._send(
                 build_error_response(
-                    request_id=request_id if isinstance(request_id, str | int) else None,
+                    request_id=request_id if is_jsonrpc_request_id(request_id) else None,
                     code=ErrorCodes.INVALID_REQUEST,
                     message="response payload is not supported on this endpoint",
                 )
@@ -114,14 +141,14 @@ class StdioRPCBridge:
 
         await self._send(
             build_error_response(
-                request_id=request_id if isinstance(request_id, str | int) else None,
+                request_id=request_id if is_jsonrpc_request_id(request_id) else None,
                 code=ErrorCodes.METHOD_NOT_FOUND,
                 message=f"method not found: {method}",
             )
         )
 
     async def _handle_initialize(self, request_id: Any) -> None:
-        if not isinstance(request_id, str | int):
+        if not is_jsonrpc_request_id(request_id):
             await self._send(
                 build_error_response(
                     request_id=None,
@@ -155,7 +182,7 @@ class StdioRPCBridge:
         )
 
     async def _handle_prompt(self, request_id: Any, params: Any) -> None:
-        if not isinstance(request_id, str | int):
+        if not is_jsonrpc_request_id(request_id):
             await self._send(
                 build_error_response(
                     request_id=None,
@@ -213,13 +240,13 @@ class StdioRPCBridge:
             )
             return
 
-        asyncio.create_task(
+        self._finalize_prompt_task = asyncio.create_task(
             self._finalize_prompt_result(result_future=result_future, request_id=request_id),
             name=f"rpc-prompt-finalize-{request_id}",
         )
 
     async def _handle_cancel(self, request_id: Any) -> None:
-        if not isinstance(request_id, str | int):
+        if not is_jsonrpc_request_id(request_id):
             await self._send(
                 build_error_response(
                     request_id=None,
@@ -233,7 +260,7 @@ class StdioRPCBridge:
         await self._send(build_success_response(request_id, {"cancelled": cancelled}))
 
     async def _handle_replay(self, request_id: Any) -> None:
-        if not isinstance(request_id, str | int):
+        if not is_jsonrpc_request_id(request_id):
             await self._send(
                 build_error_response(
                     request_id=None,
@@ -313,6 +340,8 @@ class StdioRPCBridge:
             )
         finally:
             self._clear_active_prompt(result_future)
+            if self._finalize_prompt_task is asyncio.current_task():
+                self._finalize_prompt_task = None
 
     async def _cancel_active_prompt(self) -> bool:
         if not self._has_active_prompt():
@@ -347,6 +376,72 @@ class StdioRPCBridge:
         self._active_prompt_result = None
         self._active_prompt_request_id = None
 
+    def _install_stdin_reader(self) -> None:
+        if self._stdin_queue is not None:
+            return
+
+        loop = asyncio.get_running_loop()
+        self._stdin_queue = asyncio.Queue()
+        try:
+            stdin_fd = sys.stdin.fileno()
+            loop.add_reader(stdin_fd, self._on_stdin_ready)
+        except (AttributeError, NotImplementedError, OSError, ValueError) as exc:
+            self._stdin_queue = None
+            raise RuntimeError(f"failed to install stdin reader: {exc}") from exc
+
+        self._stdin_reader_fd = stdin_fd
+        self._stdin_error = None
+
+    def _remove_stdin_reader(self) -> None:
+        stdin_fd = self._stdin_reader_fd
+        if stdin_fd is None:
+            return
+
+        try:
+            asyncio.get_running_loop().remove_reader(stdin_fd)
+        except RuntimeError:
+            pass
+        self._stdin_reader_fd = None
+        self._stdin_queue = None
+
+    def _on_stdin_ready(self) -> None:
+        queue = self._stdin_queue
+        if queue is None:
+            return
+
+        try:
+            raw_line = sys.stdin.readline()
+        except Exception as exc:
+            logger.exception("rpc stdin read failed")
+            self._stdin_error = RuntimeError(f"stdin read failed: {exc}")
+            self._remove_stdin_reader()
+            queue.put_nowait(None)
+            return
+
+        if raw_line == "":
+            self._remove_stdin_reader()
+            queue.put_nowait(None)
+            return
+        queue.put_nowait(raw_line)
+
+    async def _read_next_input_line(self) -> str:
+        queue = self._stdin_queue
+        if queue is None:
+            raise RuntimeError("stdin reader is not installed")
+
+        raw_line = await queue.get()
+        if raw_line is None:
+            if self._stdin_error is not None:
+                raise self._stdin_error
+            return ""
+        return raw_line
+
+    async def _await_finalize_prompt_task(self) -> None:
+        finalize_task = self._finalize_prompt_task
+        if finalize_task is None:
+            return
+        await finalize_task
+
     async def _send(self, payload: dict[str, Any]) -> None:
         encoded = f"{self._dump_json(payload)}\n"
         async with self._write_lock:

{comate_cli-0.2.0 → comate_cli-0.2.2}/comate_cli/terminal_agent/tui.py

@@ -741,6 +741,15 @@ class TerminalAgentTUI(
                     }
                 if isinstance(event, StopEvent):
                     stop_reason = str(event.reason or "")
+                    if event.reason == "shutdown_request":
+                        shutdown_recipients = event.metadata.get("shutdown_recipients", [])
+                        if shutdown_recipients:
+                            from comate_agent_sdk.agent.runner_engine.query_stream import (
+                                _send_team_shutdown_response,
+                            )
+                            _send_team_shutdown_response(
+                                self._session.runtime, recipients=shutdown_recipients
+                            )
                 is_waiting, new_questions = self._renderer.handle_event(event)
                 self._maybe_flash_tool_result(event)
                 if is_waiting:

{comate_cli-0.2.0 → comate_cli-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "comate-cli"
-version = "0.2.0"
+version = "0.2.2"
 description = "Comate terminal CLI built on comate-agent-sdk"
 readme = "README.md"
 requires-python = ">=3.11"

comate_cli-0.2.2/tests/test_app_mcp_preload.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import asyncio
+import types
+import unittest
+
+from comate_cli.terminal_agent.app import _preload_mcp_in_tui
+
+
+class _FakeRuntime:
+    def __init__(self) -> None:
+        self.options = types.SimpleNamespace(mcp_enabled=True)
+        self._mcp_manager = types.SimpleNamespace(
+            failed_servers=[],
+            tool_infos=[types.SimpleNamespace(server_alias="ctx7")],
+        )
+        self.start_calls = 0
+        self.awaited = asyncio.Event()
+
+    def start_mcp_preload(self) -> asyncio.Task[None]:
+        self.start_calls += 1
+
+        async def _worker() -> None:
+            await asyncio.sleep(0)
+            self.awaited.set()
+
+        return asyncio.create_task(_worker())
+
+
+class _FakeSession:
+    def __init__(self, runtime: _FakeRuntime) -> None:
+        self.runtime = runtime
+        self._agent = None
+
+
+class TestAppMcpPreload(unittest.IsolatedAsyncioTestCase):
+    async def test_preload_mcp_in_tui_uses_runtime_public_entry(self) -> None:
+        runtime = _FakeRuntime()
+        session = _FakeSession(runtime)
+
+        await _preload_mcp_in_tui(session)  # type: ignore[arg-type]
+
+        self.assertEqual(runtime.start_calls, 1)
+        self.assertTrue(runtime.awaited.is_set())
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)