windcode 0.1.0__py3-none-any.whl

windcode/extensions/mcp/catalog.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import hashlib
+import re
+from dataclasses import dataclass
+from typing import Any
+
+from mcp.types import (
+    InitializeResult,
+    ListPromptsResult,
+    ListResourcesResult,
+    ListResourceTemplatesResult,
+    ListToolsResult,
+)
+
+from windcode.extensions.models import normalize_id
+
+# Provider function-calling names allow only these characters, capped at 64.
+# MCP stable_ids ("mcp:{server}/tool/{name}") contain ":" and "/", and tool
+# names may contain ".", so they must be sanitized before reaching a provider.
+_WIRE_UNSAFE = re.compile(r"[^A-Za-z0-9_-]")
+_WIRE_MAX = 64
+
+
+def mcp_tool_wire_name(server_id: str, tool_name: str) -> str:
+    """Derive a provider-safe function name for an MCP tool.
+
+    Keeps the ``{server}__{tool}`` shape when it fits the provider charset
+    and length limit, falling back to a deterministic hash suffix otherwise so
+    distinct tools never collide on the wire.
+    """
+
+    base = _WIRE_UNSAFE.sub("_", f"{server_id}__{tool_name}")
+    if len(base) <= _WIRE_MAX:
+        return base
+    digest = hashlib.sha256(f"{server_id}\0{tool_name}".encode()).hexdigest()[:8]
+    return f"{base[: _WIRE_MAX - len(digest) - 1]}_{digest}"
+
+
+def _suffix(value: str) -> str:
+    normalized = value.strip().lower()
+    try:
+        return normalize_id(normalized)
+    except ValueError:
+        return hashlib.sha256(value.encode()).hexdigest()[:16]
+
+
+@dataclass(frozen=True, slots=True)
+class McpToolDefinition:
+    stable_id: str
+    server_id: str
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+    annotations: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class McpResourceDefinition:
+    stable_id: str
+    server_id: str
+    uri: str
+    name: str
+    description: str
+
+
+@dataclass(frozen=True, slots=True)
+class McpPromptDefinition:
+    stable_id: str
+    server_id: str
+    name: str
+    description: str
+
+
+@dataclass(frozen=True, slots=True)
+class McpCatalog:
+    server_id: str
+    instructions: str | None
+    tools: tuple[McpToolDefinition, ...]
+    resources: tuple[McpResourceDefinition, ...]
+    prompts: tuple[McpPromptDefinition, ...]
+
+
+def build_catalog(
+    server_id: str,
+    initialize: InitializeResult,
+    tools: ListToolsResult,
+    resources: ListResourcesResult,
+    prompts: ListPromptsResult,
+    resource_templates: ListResourceTemplatesResult | None = None,
+    *,
+    max_metadata_chars: int = 65_536,
+) -> McpCatalog:
+    server = normalize_id(server_id)
+    instructions = initialize.instructions
+    if instructions is not None and len(instructions) > max_metadata_chars:
+        raise ValueError("MCP instructions exceed metadata limit")
+    tool_items: list[McpToolDefinition] = []
+    seen_tools: set[str] = set()
+    for tool in tools.tools:
+        name = normalize_id(tool.name)
+        if name in seen_tools:
+            raise ValueError(f"duplicate MCP tool: {name}")
+        seen_tools.add(name)
+        schema = tool.inputSchema
+        annotations = {} if tool.annotations is None else tool.annotations.model_dump(mode="json")
+        tool_items.append(
+            McpToolDefinition(
+                f"mcp:{server}/tool/{name}",
+                server,
+                tool.name,
+                tool.description or "",
+                schema,
+                annotations,
+            )
+        )
+    resource_items = [
+        McpResourceDefinition(
+            f"mcp:{server}/resource/{_suffix(str(resource.uri))}",
+            server,
+            str(resource.uri),
+            resource.name,
+            resource.description or "",
+        )
+        for resource in resources.resources
+    ]
+    if resource_templates is not None:
+        resource_items.extend(
+            McpResourceDefinition(
+                f"mcp:{server}/resource/{_suffix(template.uriTemplate)}",
+                server,
+                template.uriTemplate,
+                template.name,
+                template.description or "",
+            )
+            for template in resource_templates.resourceTemplates
+        )
+    prompt_items = tuple(
+        McpPromptDefinition(
+            f"mcp:{server}/prompt/{normalize_id(prompt.name)}",
+            server,
+            prompt.name,
+            prompt.description or "",
+        )
+        for prompt in prompts.prompts
+    )
+    return McpCatalog(
+        server,
+        instructions,
+        tuple(sorted(tool_items, key=lambda item: item.stable_id)),
+        tuple(sorted(resource_items, key=lambda item: item.stable_id)),
+        tuple(sorted(prompt_items, key=lambda item: item.stable_id)),
+    )

windcode/extensions/mcp/client.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import tempfile
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from datetime import timedelta
+from pathlib import Path
+from types import TracebackType
+from typing import Self, TextIO
+
+import httpx
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from mcp.client.streamable_http import streamable_http_client
+from mcp.types import (
+    CallToolResult,
+    GetPromptResult,
+    InitializeResult,
+    ListPromptsResult,
+    ListResourcesResult,
+    ListResourceTemplatesResult,
+    ListToolsResult,
+    PaginatedRequestParams,
+    ReadResourceResult,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedStdioServer:
+    command: str
+    args: tuple[str, ...] = ()
+    cwd: Path | None = None
+    env: dict[str, str] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedHttpServer:
+    url: str
+    headers: dict[str, str] | None = None
+
+
+class McpClient:
+    def __init__(
+        self,
+        definition: ResolvedStdioServer | ResolvedHttpServer,
+        *,
+        connect_timeout: float = 10.0,
+        call_timeout: float = 60.0,
+        stderr_limit: int = 16_384,
+    ) -> None:
+        self.definition = definition
+        self.connect_timeout = connect_timeout
+        self.call_timeout = call_timeout
+        self._stack: AsyncExitStack | None = None
+        self._session: ClientSession | None = None
+        self._stderr_limit = stderr_limit
+        self._stderr_file: TextIO | None = None
+        self._stderr_value = ""
+        self.initialize_result: InitializeResult | None = None
+        self._owner_task: asyncio.Task[None] | None = None
+        self._ready: asyncio.Future[InitializeResult] | None = None
+        self._close_requested: asyncio.Event | None = None
+        self._close_error: BaseException | None = None
+
+    @property
+    def stderr(self) -> str:
+        stream = self._stderr_file
+        if stream is None:
+            return self._stderr_value
+        stream.flush()
+        position = stream.tell()
+        stream.seek(0)
+        value = stream.read(self._stderr_limit)
+        stream.seek(position)
+        return value
+
+    @property
+    def connected(self) -> bool:
+        return self._session is not None
+
+    @property
+    def close_error(self) -> BaseException | None:
+        return self._close_error
+
+    async def __aenter__(self) -> Self:
+        await self.connect()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        del exc_type, exc, traceback
+        await self.aclose()
+
+    async def connect(self) -> InitializeResult:
+        if self.initialize_result is not None:
+            return self.initialize_result
+        if self._owner_task is None:
+            loop = asyncio.get_running_loop()
+            self._ready = loop.create_future()
+            self._close_requested = asyncio.Event()
+            self._owner_task = asyncio.create_task(self._own_connection())
+        ready = self._ready
+        if ready is None:
+            raise RuntimeError("MCP client connection owner was not initialized")
+        async with asyncio.timeout(self.connect_timeout):
+            return await asyncio.shield(ready)
+
+    async def _own_connection(self) -> None:
+        stack = AsyncExitStack()
+        try:
+            async with asyncio.timeout(self.connect_timeout):
+                if isinstance(self.definition, ResolvedStdioServer):
+                    self._stderr_file = tempfile.TemporaryFile(
+                        mode="w+", encoding="utf-8", errors="replace"
+                    )
+                    minimum_env = {
+                        key: value
+                        for key in ("PATH", "HOME", "USER", "LANG", "LC_ALL", "TMPDIR")
+                        if (value := os.environ.get(key)) is not None
+                    }
+                    minimum_env.update(self.definition.env or {})
+                    streams = await stack.enter_async_context(
+                        stdio_client(
+                            StdioServerParameters(
+                                command=self.definition.command,
+                                args=list(self.definition.args),
+                                cwd=self.definition.cwd,
+                                env=minimum_env,
+                            ),
+                            errlog=self._stderr_file,
+                        )
+                    )
+                    read_stream, write_stream = streams
+                else:
+                    http_client = await stack.enter_async_context(
+                        httpx.AsyncClient(
+                            headers=self.definition.headers,
+                            timeout=httpx.Timeout(self.call_timeout, connect=self.connect_timeout),
+                        )
+                    )
+                    streams = await stack.enter_async_context(
+                        streamable_http_client(
+                            self.definition.url,
+                            http_client=http_client,
+                            terminate_on_close=False,
+                        )
+                    )
+                    read_stream, write_stream, _ = streams
+                session = await stack.enter_async_context(
+                    ClientSession(
+                        read_stream,
+                        write_stream,
+                        read_timeout_seconds=timedelta(seconds=self.call_timeout),
+                    )
+                )
+                result = await session.initialize()
+            self._stack = stack
+            self._session = session
+            self.initialize_result = result
+            ready = self._ready
+            if ready is not None and not ready.done():
+                ready.set_result(result)
+            close_requested = self._close_requested
+            if close_requested is not None:
+                await close_requested.wait()
+        except BaseException as exc:
+            ready = self._ready
+            if ready is not None and not ready.done():
+                if isinstance(exc, asyncio.CancelledError):
+                    ready.cancel()
+                else:
+                    ready.set_exception(exc)
+            elif not isinstance(exc, asyncio.CancelledError):
+                raise
+        finally:
+            self._session = None
+            self.initialize_result = None
+            await stack.aclose()
+            self._stack = None
+            self._close_stderr()
+
+    def _close_stderr(self) -> None:
+        if self._stderr_file is None:
+            return
+        self._stderr_value = self.stderr
+        self._stderr_file.close()
+        self._stderr_file = None
+
+    def _require_session(self) -> ClientSession:
+        if self._session is None:
+            raise RuntimeError("MCP client is not connected")
+        return self._session
+
+    async def list_tools(self, cursor: str | None = None) -> ListToolsResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().list_tools(
+                params=PaginatedRequestParams(cursor=cursor)
+            )
+
+    async def list_resources(self, cursor: str | None = None) -> ListResourcesResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().list_resources(
+                params=PaginatedRequestParams(cursor=cursor)
+            )
+
+    async def list_resource_templates(
+        self, cursor: str | None = None
+    ) -> ListResourceTemplatesResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().list_resource_templates(
+                params=PaginatedRequestParams(cursor=cursor)
+            )
+
+    async def list_prompts(self, cursor: str | None = None) -> ListPromptsResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().list_prompts(
+                params=PaginatedRequestParams(cursor=cursor)
+            )
+
+    async def call_tool(self, name: str, arguments: dict[str, object]) -> CallToolResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().call_tool(name, arguments)
+
+    async def read_resource(self, uri: str) -> ReadResourceResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().read_resource(uri)  # pyright: ignore[reportArgumentType]
+
+    async def get_prompt(
+        self, name: str, arguments: dict[str, str] | None = None
+    ) -> GetPromptResult:
+        async with asyncio.timeout(self.call_timeout):
+            return await self._require_session().get_prompt(name, arguments)
+
+    async def aclose(self) -> None:
+        task = self._owner_task
+        if task is None:
+            self._settle_ready()
+            self._close_stderr()
+            return
+        close_requested = self._close_requested
+        if close_requested is not None:
+            close_requested.set()
+        try:
+            async with asyncio.timeout(self.connect_timeout):
+                result = await asyncio.gather(task, return_exceptions=True)
+                if result and isinstance(result[0], BaseException):
+                    self._close_error = result[0]
+        except TimeoutError:
+            task.cancel()
+            result = await asyncio.gather(task, return_exceptions=True)
+            self._close_error = TimeoutError("MCP connection close timed out")
+            if result and isinstance(result[0], BaseException):
+                self._close_error = result[0]
+        finally:
+            self._settle_ready()
+            self._owner_task = None
+            self._ready = None
+            self._close_requested = None
+
+    def _settle_ready(self) -> None:
+        ready = self._ready
+        if ready is None:
+            return
+        if not ready.done():
+            ready.cancel()
+        elif not ready.cancelled():
+            ready.exception()

windcode/extensions/mcp/runtime.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from enum import StrEnum
+
+from windcode.extensions.mcp.client import McpClient
+
+
+class McpServerState(StrEnum):
+    DISCOVERED = "discovered"
+    CONNECTING = "connecting"
+    READY = "ready"
+    FAILED = "failed"
+    CLOSING = "closing"
+    CLOSED = "closed"
+
+
+ClientFactory = Callable[[], McpClient]
+McpObserver = Callable[[str, str, str], Awaitable[None]]
+
+
+@dataclass(slots=True)
+class _ServerSlot:
+    factory: ClientFactory
+    required: bool
+    state: McpServerState = McpServerState.DISCOVERED
+    client: McpClient | None = None
+    error: BaseException | None = None
+    reconnects: int = 0
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+
+class McpRuntime:
+    def __init__(
+        self,
+        servers: dict[str, tuple[ClientFactory, bool]],
+        observer: McpObserver | None = None,
+    ) -> None:
+        self._servers = {
+            server_id: _ServerSlot(factory, required)
+            for server_id, (factory, required) in sorted(servers.items())
+        }
+        self._closed = False
+        self.observer = observer
+
+    async def _observe(self, action: str, server_id: str, status: str) -> None:
+        if self.observer is not None:
+            await self.observer(action, server_id, status)
+
+    def state(self, server_id: str) -> McpServerState:
+        return self._servers[server_id].state
+
+    @property
+    def server_ids(self) -> tuple[str, ...]:
+        return tuple(self._servers)
+
+    @property
+    def required_server_ids(self) -> tuple[str, ...]:
+        return tuple(server_id for server_id, slot in self._servers.items() if slot.required)
+
+    async def activate(self, server_id: str) -> McpClient:
+        slot = self._servers[server_id]
+        async with slot.lock:
+            if self._closed:
+                raise RuntimeError("MCP runtime is closed")
+            if slot.state is McpServerState.READY and slot.client is not None:
+                return slot.client
+            slot.state = McpServerState.CONNECTING
+            await self._observe("mcp_connecting", server_id, "connecting")
+            client: McpClient | None = None
+            try:
+                client = slot.factory()
+                await client.connect()
+            except BaseException as exc:
+                slot.state = McpServerState.FAILED
+                slot.error = exc
+                await self._observe("diagnostic", server_id, "failed")
+                if client is not None:
+                    await client.aclose()
+                raise
+            slot.client = client
+            slot.error = None
+            slot.state = McpServerState.READY
+            await self._observe("mcp_connected", server_id, "ready")
+            return client
+
+    async def call(
+        self, server_id: str, operation: Callable[[McpClient], Awaitable[object]]
+    ) -> object:
+        client = await self.activate(server_id)
+        try:
+            result = await operation(client)
+            await self._observe("mcp_called", server_id, "success")
+            return result
+        except (ConnectionError, EOFError, BrokenPipeError):
+            slot = self._servers[server_id]
+            async with slot.lock:
+                if slot.reconnects >= 1:
+                    raise
+                slot.reconnects += 1
+                if slot.client is not None:
+                    await slot.client.aclose()
+                slot.client = None
+                slot.state = McpServerState.DISCOVERED
+            client = await self.activate(server_id)
+            result = await operation(client)
+            await self._observe("mcp_called", server_id, "success")
+            return result
+        except BaseException:
+            await self._observe("diagnostic", server_id, "call_failed")
+            raise
+
+    async def activate_required(self, *, concurrency: int = 4) -> None:
+        semaphore = asyncio.Semaphore(concurrency)
+
+        async def activate_one(server_id: str) -> None:
+            async with semaphore:
+                await self.activate(server_id)
+
+        await asyncio.gather(
+            *(activate_one(server_id) for server_id, slot in self._servers.items() if slot.required)
+        )
+
+    async def aclose(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        for server_id, slot in reversed(tuple(self._servers.items())):
+            async with slot.lock:
+                slot.state = McpServerState.CLOSING
+                client = slot.client
+                try:
+                    if client is not None:
+                        await client.aclose()
+                        if client.close_error is not None:
+                            await self._observe("diagnostic", server_id, "close_failed")
+                except Exception:
+                    await self._observe("diagnostic", server_id, "close_failed")
+                finally:
+                    slot.client = None
+                    slot.state = McpServerState.CLOSED
+                    await self._observe("mcp_closed", server_id, "closed")