bareagent-cli 0.1.0__py3-none-any.whl

bareagent/mcp/transport/stdio.py

@@ -0,0 +1,202 @@
+"""Stdio transport: subprocess + newline-delimited JSON.
+
+Two background daemon threads consume the subprocess pipes:
+
+- `_reader`: parses stdout as NDJSON, routes responses + notifications.
+- `_stderr_reader`: drains stderr into the project logger so server banners
+  do not block the pipe.
+
+On EOF / subprocess death, every pending future is failed with
+`MCPTransportError` so callers never hang forever.
+"""
+
+from __future__ import annotations
+
+import logging
+import subprocess
+import threading
+from collections.abc import Iterator
+from typing import IO
+
+from ..errors import MCPProtocolError, MCPTransportError
+from ..protocol import Notification, Request, Response, decode_message
+from .base import Transport
+
+_log = logging.getLogger(__name__)
+
+_TERMINATE_TIMEOUT = 2.0
+_KILL_TIMEOUT = 2.0
+
+
+class StdioTransport(Transport):
+    """Spawn an MCP server as a subprocess and talk to it over stdin/stdout."""
+
+    def __init__(
+        self,
+        command: list[str],
+        *,
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+    ) -> None:
+        super().__init__()
+        if not command:
+            raise ValueError("command must be a non-empty list")
+        self._command = list(command)
+        self._env = env
+        self._cwd = cwd
+        self._proc: subprocess.Popen[bytes] | None = None
+        self._reader: threading.Thread | None = None
+        self._stderr_reader: threading.Thread | None = None
+        self._write_lock = threading.Lock()
+        self._closed = False
+        # Set true by ``close()`` so the reader thread can distinguish graceful
+        # shutdown (don't fire disconnect handler) from unexpected EOF.
+        self._closing = False
+
+    def start(self) -> None:
+        if self._proc is not None:
+            raise RuntimeError("StdioTransport already started")
+        try:
+            self._proc = subprocess.Popen(
+                self._command,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                env=self._env,
+                cwd=self._cwd,
+                bufsize=0,
+            )
+        except (OSError, FileNotFoundError) as exc:
+            raise MCPTransportError(f"failed to spawn MCP server: {exc}") from exc
+
+        self._reader = threading.Thread(
+            target=self._read_loop, name="mcp-stdio-reader", daemon=True
+        )
+        self._stderr_reader = threading.Thread(
+            target=self._stderr_loop, name="mcp-stdio-stderr", daemon=True
+        )
+        self._reader.start()
+        self._stderr_reader.start()
+
+    def send(self, message: str) -> None:
+        proc = self._proc
+        if proc is None or proc.stdin is None:
+            raise MCPTransportError("transport not started")
+        if proc.poll() is not None:
+            raise MCPTransportError(f"subprocess exited with code {proc.returncode}")
+        line = message + "\n"
+        try:
+            with self._write_lock:
+                proc.stdin.write(line.encode("utf-8"))
+                proc.stdin.flush()
+        except (BrokenPipeError, OSError) as exc:
+            raise MCPTransportError(
+                f"failed to write to subprocess stdin: {exc}"
+            ) from exc
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._closing = True
+        proc = self._proc
+        if proc is None:
+            return
+        # Graceful shutdown: close stdin → wait → terminate → wait → kill.
+        try:
+            if proc.stdin is not None and not proc.stdin.closed:
+                with self._write_lock:
+                    try:
+                        proc.stdin.close()
+                    except OSError:
+                        pass
+            try:
+                proc.wait(timeout=_TERMINATE_TIMEOUT)
+            except subprocess.TimeoutExpired:
+                proc.terminate()
+                try:
+                    proc.wait(timeout=_KILL_TIMEOUT)
+                except subprocess.TimeoutExpired:
+                    proc.kill()
+                    try:
+                        proc.wait(timeout=_KILL_TIMEOUT)
+                    except subprocess.TimeoutExpired:
+                        pass
+        finally:
+            self._fail_all_pending("stdio transport closed")
+
+    def is_alive(self) -> bool:
+        proc = self._proc
+        return proc is not None and proc.poll() is None
+
+    # --- background threads ---
+
+    def _read_loop(self) -> None:
+        proc = self._proc
+        assert proc is not None and proc.stdout is not None
+        disconnect_reason: str | None = None
+        try:
+            for raw in self._iter_lines(proc.stdout):
+                line = raw.decode("utf-8", errors="replace").strip()
+                if not line:
+                    continue
+                try:
+                    msg = decode_message(line)
+                except MCPProtocolError as exc:
+                    _log.warning(
+                        "MCP stdio: ignoring non-protocol line (%s): %r",
+                        exc,
+                        line[:200],
+                    )
+                    continue
+                self._dispatch(msg)
+        except Exception as exc:  # pragma: no cover — defensive
+            disconnect_reason = f"stdio reader crashed: {exc}"
+            _log.warning("MCP stdio reader crashed: %s", exc)
+        finally:
+            # Distinguish graceful (user called close) from unexpected EOF /
+            # subprocess crash. Only fire disconnect handler in the unexpected
+            # case so a normal shutdown does not spuriously notify the manager.
+            if not self._closing:
+                if disconnect_reason is None:
+                    returncode = proc.poll()
+                    if returncode is not None:
+                        disconnect_reason = (
+                            f"stdout EOF (subprocess exited with code {returncode})"
+                        )
+                    else:
+                        disconnect_reason = "stdout EOF (subprocess exited)"
+                self._invoke_disconnect(disconnect_reason)
+            self._fail_all_pending("stdio reader exited (subprocess EOF or error)")
+
+    def _stderr_loop(self) -> None:
+        proc = self._proc
+        assert proc is not None and proc.stderr is not None
+        try:
+            for raw in self._iter_lines(proc.stderr):
+                text = raw.decode("utf-8", errors="replace").rstrip()
+                if text:
+                    _log.warning("MCP server stderr: %s", text)
+        except Exception:  # pragma: no cover
+            pass
+
+    @staticmethod
+    def _iter_lines(stream: IO[bytes]) -> Iterator[bytes]:
+        # subprocess pipes opened in bytes mode iterate by line (\n-delimited).
+        while True:
+            line = stream.readline()
+            if not line:
+                return
+            yield line
+
+    def _dispatch(self, msg: Request | Response | Notification) -> None:
+        if isinstance(msg, Response):
+            self._route_response(msg)
+        elif isinstance(msg, Notification):
+            self._route_notification(msg)
+        else:
+            # Server-initiated request. PR1 does not support these; in PR2 the
+            # client will handle method dispatch (ping, sampling, elicitation).
+            _log.warning(
+                "MCP stdio: ignoring server-to-client request (method=%r)", msg.method
+            )

bareagent/memory/__init__.py

@@ -0,0 +1 @@
+"""Memory modules for BareAgent."""

bareagent/memory/compact.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+import copy
+import logging
+from typing import Any
+
+from bareagent.core.fileutil import collect_tool_names, is_tool_result_message, stringify
+from bareagent.memory.token_counter import estimate_tokens
+from bareagent.memory.transcript import TranscriptManager
+from bareagent.provider.base import BaseLLMProvider
+
+logger = logging.getLogger(__name__)
+
+_TRUNCATED_PREFIX = "[truncated:"
+_SUMMARY_SYSTEM_PROMPT = (
+    "你是 BareAgent 的上下文压缩助手。请用中文总结对话中的目标、已完成工作、"
+    "关键约束、重要文件路径、工具执行结果要点，以及接下来继续工作所需的上下文。"
+    "输出简洁但不能遗漏事实。"
+)
+
+
+def _micro_compact(
+    messages: list[dict[str, Any]],
+    keep_recent: int = 3,
+    tool_name_by_id: dict[str, str] | None = None,
+) -> dict[str, str]:
+    if tool_name_by_id is None:
+        tool_name_by_id = collect_tool_names(messages)
+    tool_result_indices = [
+        index
+        for index, message in enumerate(messages)
+        if is_tool_result_message(message)
+    ]
+    if keep_recent > 0:
+        compact_indices = set(tool_result_indices[:-keep_recent])
+    else:
+        compact_indices = set(tool_result_indices)
+
+    for index in compact_indices:
+        message = messages[index]
+        content = message.get("content")
+        if not isinstance(content, list):
+            continue
+        for block in content:
+            if not isinstance(block, dict) or block.get("type") != "tool_result":
+                continue
+            original_text = stringify(block.get("content", ""))
+            if original_text.startswith(_TRUNCATED_PREFIX):
+                continue
+            tool_use_id = str(block.get("tool_use_id", ""))
+            tool_name = tool_name_by_id.get(tool_use_id, "unknown")
+            block["content"] = (
+                f"[truncated: {tool_name} result, {len(original_text)} chars]"
+            )
+    return tool_name_by_id
+
+
+def _serialize(
+    messages: list[dict[str, Any]],
+    tool_name_by_id: dict[str, str] | None = None,
+) -> str:
+    lines: list[str] = []
+    if tool_name_by_id is None:
+        tool_name_by_id = collect_tool_names(messages)
+    for message in messages:
+        role = str(message.get("role", "unknown"))
+        lines.append(f"[{role}]")
+        lines.append(_serialize_content(message.get("content"), tool_name_by_id))
+        lines.append("")
+    return "\n".join(lines).strip()
+
+
+class Compactor:
+    def __init__(
+        self,
+        provider: BaseLLMProvider,
+        transcript_mgr: TranscriptManager | None,
+        threshold: int = 50000,
+        session_id: str = "default",
+    ) -> None:
+        self._provider = provider
+        self._transcript_mgr = transcript_mgr
+        self._threshold = threshold
+        self._session_id = session_id
+
+    def get_session_id(self) -> str:
+        return self._session_id
+
+    def set_session_id(self, new_session_id: str) -> None:
+        self._session_id = new_session_id
+
+    def __call__(self, messages: list[dict[str, Any]], force: bool = False) -> None:
+        if not force and estimate_tokens(messages) <= self._threshold:
+            return
+
+        _backup = [_clone_message(m) for m in messages]
+
+        tool_name_by_id = _micro_compact(messages, keep_recent=3)
+
+        history_messages, pending_user_message = _split_pending_user_turn(messages)
+        summary_source_messages = [
+            message for message in history_messages if message.get("role") != "system"
+        ]
+        if not summary_source_messages:
+            messages[:] = _backup
+            return
+
+        if self._transcript_mgr is not None:
+            self._transcript_mgr.save(messages, self._session_id)
+        try:
+            summary = self._provider.create(
+                messages=[
+                    {"role": "system", "content": _SUMMARY_SYSTEM_PROMPT},
+                    {
+                        "role": "user",
+                        "content": (
+                            "请简洁总结以下对话的关键信息和已完成的工作，供后续继续工作使用：\n\n"
+                            + _serialize(summary_source_messages, tool_name_by_id)
+                        ),
+                    },
+                ],
+                tools=[],
+                max_tokens=2000,
+            )
+        except Exception:
+            logger.warning("Context compression failed", exc_info=True)
+            messages[:] = _backup
+            return
+
+        system_messages = [
+            _clone_message(message)
+            for message in messages
+            if message.get("role") == "system"
+        ]
+        messages.clear()
+        messages.extend(system_messages)
+        messages.extend(
+            [
+                {"role": "user", "content": f"[Context Compressed]\n{summary.text}"},
+                {
+                    "role": "assistant",
+                    "content": "收到，我已理解之前的上下文，继续工作。",
+                },
+            ]
+        )
+        if pending_user_message is not None:
+            messages.append(pending_user_message)
+
+
+make_compact_fn = Compactor
+
+
+def _serialize_content(content: Any, tool_name_by_id: dict[str, str]) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for block in content:
+            serialized = _serialize_block(block, tool_name_by_id)
+            if serialized:
+                parts.append(serialized)
+        return "\n".join(parts)
+    return stringify(content)
+
+
+def _serialize_block(block: Any, tool_name_by_id: dict[str, str]) -> str:
+    if isinstance(block, str):
+        return block
+    if not isinstance(block, dict):
+        return stringify(block)
+
+    block_type = block.get("type")
+    if block_type == "text":
+        return str(block.get("text", ""))
+    if block_type == "tool_use":
+        return (
+            f"[tool_use:{block.get('name', 'unknown')}] "
+            f"{stringify(block.get('input', {}))}"
+        )
+    if block_type == "tool_result":
+        tool_use_id = str(block.get("tool_use_id", ""))
+        tool_name = tool_name_by_id.get(tool_use_id, "unknown")
+        return f"[tool_result:{tool_name}] {stringify(block.get('content', ''))}"
+
+    if "content" in block:
+        return _serialize_content(block.get("content"), tool_name_by_id)
+    if "text" in block:
+        return str(block.get("text", ""))
+    return stringify(block)
+
+
+def _clone_message(message: dict[str, Any]) -> dict[str, Any]:
+    return copy.deepcopy(message)
+
+
+def _split_pending_user_turn(
+    messages: list[dict[str, Any]],
+) -> tuple[list[dict[str, Any]], dict[str, Any] | None]:
+    if messages and messages[-1].get("role") == "user":
+        return messages[:-1], _clone_message(messages[-1])
+    return messages, None

bareagent/memory/conversation_io.py

@@ -0,0 +1,226 @@
+"""Conversation import/export serialization (REPL-independent, unit-testable).
+
+Pure helpers backing the ``/export`` and ``/import`` REPL commands:
+
+- :func:`render_markdown` — human-readable Markdown (mirrors the traversal
+  structure of ``main._replay_stdio_transcript``).
+- :func:`to_export_json` — self-contained, faithful JSON wrapper.
+- :func:`parse_import` — auto-detecting loader with shape validation.
+
+No dependency on ``bareagent.main`` / UI / loop, so this module can be exercised in
+isolation.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+EXPORT_VERSION = 1
+_DEFAULT_MAX_TOOL_CHARS = 2000
+_TOOL_INPUT_PREVIEW_CHARS = 200
+_TRUNCATION_MARKER = "… (truncated)"
+
+
+def _truncate(text: str, limit: int) -> str:
+    if limit < 0 or len(text) <= limit:
+        return text
+    return text[:limit] + _TRUNCATION_MARKER
+
+
+def _stringify_content(content: Any) -> str:
+    """Coerce arbitrary tool_result content into display text."""
+    if isinstance(content, str):
+        return content
+    if content is None:
+        return ""
+    return json.dumps(content, ensure_ascii=False, default=str)
+
+
+def _tool_use_summary(block: dict[str, Any]) -> str:
+    name = str(block.get("name", "unknown"))
+    raw_input = block.get("input", {})
+    preview = json.dumps(raw_input, ensure_ascii=False, default=str)
+    preview = _truncate(preview, _TOOL_INPUT_PREVIEW_CHARS)
+    return f"- **Tool call** `{name}`: `{preview}`"
+
+
+def render_markdown(
+    messages: list[dict[str, Any]],
+    *,
+    include_thinking: bool = False,
+    max_tool_chars: int = _DEFAULT_MAX_TOOL_CHARS,
+    title: str | None = None,
+) -> str:
+    """Render *messages* as human-readable Markdown.
+
+    Skips ``system`` messages, renders user/assistant text, turns ``tool_use``
+    blocks into single-line summaries and ``tool_result`` blocks into truncated
+    code fences. ``thinking`` blocks are omitted unless *include_thinking* is
+    set. Mirrors the traversal of ``main._replay_stdio_transcript`` (including
+    the ``tool_use_id`` → tool name association). *title*, when provided, is
+    emitted as a top-level heading.
+    """
+    tool_name_by_id: dict[str, str] = {}
+    lines: list[str] = []
+
+    if title:
+        lines.append(f"# {title}")
+        lines.append("")
+
+    for message in messages:
+        role = message.get("role")
+        content = message.get("content")
+
+        if role == "system":
+            continue
+
+        if role == "user":
+            if isinstance(content, str):
+                lines.append("## User")
+                lines.append("")
+                lines.append(content)
+                lines.append("")
+                continue
+            if isinstance(content, list):
+                for block in content:
+                    if not isinstance(block, dict):
+                        continue
+                    if block.get("type") == "text":
+                        text_value = str(block.get("text", ""))
+                        if text_value:
+                            lines.append("## User")
+                            lines.append("")
+                            lines.append(text_value)
+                            lines.append("")
+                        continue
+                    if block.get("type") != "tool_result":
+                        continue
+                    tool_name = tool_name_by_id.get(
+                        str(block.get("tool_use_id", "")),
+                        "unknown",
+                    )
+                    result_text = _truncate(
+                        _stringify_content(block.get("content", "")),
+                        max_tool_chars,
+                    )
+                    error_marker = " (error)" if block.get("is_error") else ""
+                    lines.append(f"### Tool result: {tool_name}{error_marker}")
+                    lines.append("")
+                    lines.append("```")
+                    lines.append(result_text)
+                    lines.append("```")
+                    lines.append("")
+            continue
+
+        if role != "assistant":
+            continue
+
+        if isinstance(content, str):
+            lines.append("## Assistant")
+            lines.append("")
+            lines.append(content)
+            lines.append("")
+            continue
+        if not isinstance(content, list):
+            continue
+
+        for block in content:
+            if not isinstance(block, dict):
+                continue
+            block_type = block.get("type")
+            if block_type == "text":
+                text_value = str(block.get("text", ""))
+                if text_value:
+                    lines.append("## Assistant")
+                    lines.append("")
+                    lines.append(text_value)
+                    lines.append("")
+                continue
+            if block_type == "thinking":
+                if not include_thinking:
+                    continue
+                thinking_value = str(block.get("thinking", ""))
+                if thinking_value:
+                    lines.append("### Thinking")
+                    lines.append("")
+                    lines.append("```")
+                    lines.append(thinking_value)
+                    lines.append("```")
+                    lines.append("")
+                continue
+            if block_type != "tool_use":
+                continue
+
+            tool_id = str(block.get("id", ""))
+            if tool_id:
+                tool_name_by_id[tool_id] = str(block.get("name", "unknown"))
+            lines.append(_tool_use_summary(block))
+            lines.append("")
+
+    return "\n".join(lines).rstrip("\n") + "\n"
+
+
+def to_export_json(
+    messages: list[dict[str, Any]],
+    *,
+    session_id: str,
+    exported_at: str,
+) -> str:
+    """Serialize *messages* into a self-contained, faithful JSON wrapper.
+
+    The wrapper preserves every message verbatim (including ``system`` /
+    ``thinking`` / tool blocks) so a round-trip through :func:`parse_import`
+    recovers an equivalent conversation.
+    """
+    payload = {
+        "version": EXPORT_VERSION,
+        "session_id": session_id,
+        "exported_at": exported_at,
+        "messages": messages,
+    }
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+
+
+def _validate_messages(value: Any) -> list[dict[str, Any]]:
+    if not isinstance(value, list):
+        raise ValueError("conversation must be a list of messages")
+    for index, item in enumerate(value):
+        if not isinstance(item, dict):
+            raise ValueError(f"message at index {index} is not an object")
+        if "role" not in item:
+            raise ValueError(f"message at index {index} is missing a 'role' key")
+    return value
+
+
+def parse_import(text: str) -> list[dict[str, Any]]:
+    """Parse imported conversation *text* into validated messages.
+
+    Auto-detects the format: first try parsing the whole document as JSON — a
+    dict with a ``messages`` key yields those messages, a bare list is used
+    directly. If whole-document JSON parsing fails, fall back to JSONL
+    (one JSON object per non-blank line).
+
+    Validation: the result must be a list where every element is a dict
+    containing a ``role`` key; otherwise :class:`ValueError` is raised with a
+    human-readable reason. No other rewriting is performed (faithful load).
+    """
+    try:
+        document = json.loads(text)
+    except json.JSONDecodeError:
+        messages: list[Any] = []
+        for line_no, line in enumerate(text.splitlines(), 1):
+            stripped = line.strip()
+            if not stripped:
+                continue
+            try:
+                messages.append(json.loads(stripped))
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"invalid JSON on line {line_no}: {exc}") from exc
+        return _validate_messages(messages)
+
+    if isinstance(document, dict):
+        if "messages" not in document:
+            raise ValueError("JSON object is missing a 'messages' key")
+        return _validate_messages(document["messages"])
+    return _validate_messages(document)