integrations/tau/bridge.py
Rama main · réplica de lectura del repositorio original.
"""A persistent MCP connection with one task owning the stdio context lifetime."""
from __future__ import annotations
import asyncio
import os
from collections.abc import AsyncIterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from pathlib import Path
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.types import CallToolResult, PaginatedRequestParams, Tool
from pydantic import Field, JsonValue, StrictBool, model_validator
from .knowledge import CodingProfile, GeneralProfile, SessionProfile
class Settings(GeneralProfile):
repositories: list[CodingProfile] = Field(default_factory=list)
command: str = Field(default="bm", min_length=1)
args: list[str] = Field(default_factory=lambda: ["mcp", "--transport", "stdio"])
auto_recall: StrictBool = True
checkpoint_on_compact: StrictBool = True
capture_knowledge: StrictBool = True
summarize_on_shutdown: StrictBool = True
summary_timeout_seconds: float = Field(default=60, gt=0, le=300)
summary_chunk_chars: int = Field(default=16000, ge=1000, le=50000)
capture_transcript: StrictBool = False
capture_folder: str = "tau/transcripts"
timeout_seconds: float = Field(default=30, gt=0, le=300)
recall_chars: int = Field(default=12000, ge=100, le=50000)
@model_validator(mode="after")
def unique_repository_roots(self) -> Settings:
roots = [profile.root for profile in self.repositories]
if len(set(roots)) != len(roots):
raise ValueError("repository roots must be unique")
return self
def profile_for(self, cwd: Path | None) -> SessionProfile:
if cwd is None:
return self
directory = cwd.resolve()
matches = [p for p in self.repositories if directory.is_relative_to(p.root)]
# Explicit user profiles are scoped by checkout; the closest wins for nested roots.
return max(matches, key=lambda p: len(p.root.parts)) if matches else self
def load_settings() -> Settings:
# Only explicitly selected/user-owned configuration can launch a subprocess.
# Ambient repository files must not change the executable or capture destination.
path = Path(os.environ.get("TAU_BASIC_MEMORY_CONFIG", "~/.tau/basic-memory.json")).expanduser()
return Settings.model_validate_json(path.read_text()) if path.exists() else Settings()
@asynccontextmanager
async def connect(settings: Settings) -> AsyncIterator[ClientSession]:
# Preserve BM's configured environment/routing, not Tau's Python import paths.
env = {k: v for k, v in os.environ.items() if k not in {"PYTHONPATH", "PYTHONHOME"}}
params = StdioServerParameters(command=settings.command, args=settings.args, env=env)
# MCP requires a synchronous file descriptor; opening the null device does no disk I/O.
with open(os.devnull, "w") as errlog: # noqa: ASYNC230
async with stdio_client(params, errlog=errlog) as (read, write):
async with ClientSession(read, write) as session:
await asyncio.wait_for(session.initialize(), settings.timeout_seconds)
yield session
async def list_tools(session: ClientSession) -> list[Tool]:
tools: list[Tool] = []
cursors: set[str] = set()
names: set[str] = set()
cursor: str | None = None
while True:
page = await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
for tool in page.tools:
if tool.name in names:
raise ValueError(f"MCP server advertised duplicate tool: {tool.name}")
names.add(tool.name)
tools.append(tool)
cursor = page.next_cursor
if cursor is None:
return tools
if cursor in cursors:
raise ValueError("MCP tools/list repeated a pagination cursor")
cursors.add(cursor)
async def discover(settings: Settings) -> list[Tool]:
async with connect(settings) as session:
return await asyncio.wait_for(list_tools(session), settings.timeout_seconds)
def discover_sync(settings: Settings) -> list[Tool]:
# Tau composes tools immediately after synchronous setup, before session_start.
# A temporary loop in a joined thread performs discovery and closes its child
# before setup returns. No long-lived resource can leak if staged loading fails.
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="bm-discovery") as executor:
return executor.submit(lambda: asyncio.run(discover(settings))).result()
class McpConnection:
"""Own a persistent session separately from the tasks making tool calls."""
def __init__(self, settings: Settings) -> None:
self.settings = settings
self.session: ClientSession | None = None
self.owner: asyncio.Task[None] | None = None
self.stop_requested = asyncio.Event()
self.active_calls: set[asyncio.Task[CallToolResult]] = set()
async def start(self) -> None:
if self.owner is not None:
raise RuntimeError("Basic Memory connection already started")
self.stop_requested = asyncio.Event()
ready = asyncio.Future[None]()
self.owner = asyncio.create_task(self._own_session(ready), name="basic-memory-mcp")
try:
await ready
except BaseException:
await self.close()
raise
async def _own_session(self, ready: asyncio.Future[None]) -> None:
try:
async with connect(self.settings) as session:
self.session = session
if not ready.done():
ready.set_result(None)
await self.stop_requested.wait()
except Exception as exc:
if not ready.done():
ready.set_exception(
RuntimeError(f"Basic Memory startup failed ({type(exc).__name__})")
)
else:
raise
finally:
self.session = None
async def close(self) -> None:
self.stop_requested.set()
calls = tuple(self.active_calls)
for call in calls:
call.cancel()
await asyncio.gather(*calls, return_exceptions=True)
if self.owner is not None:
if self.session is None and not self.owner.done():
self.owner.cancel()
try:
await self.owner
except asyncio.CancelledError:
# Startup cancellation intentionally retires the owner before it is ready.
caller = asyncio.current_task()
if caller is not None and caller.cancelling():
raise
finally:
self.owner = None
async def call(self, name: str, arguments: dict[str, JsonValue]) -> CallToolResult:
if self.session is None or self.stop_requested.is_set():
raise RuntimeError("Basic Memory is disconnected; check configuration and /reload")
call = asyncio.create_task(self.session.call_tool(name, arguments))
self.active_calls.add(call)
try:
return await asyncio.wait_for(call, self.settings.timeout_seconds)
finally:
self.active_calls.discard(call)