flowent 0.0.13 → 0.1.1

package/backend/src/flowent/mcp.py

@@ -0,0 +1,484 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import re
+from collections.abc import Callable
+from contextlib import AsyncExitStack
+from importlib import import_module
+from typing import Any, Protocol
+
+from flowent.storage import StateStore, StoredMcpServer, StoredMcpTool
+from flowent.tools import ToolResult
+
+logger = logging.getLogger("flowent.mcp")
+MCP_CONNECT_TIMEOUT_SECONDS = 10
+
+
+class McpTransport(Protocol):
+    async def connect(self, server: StoredMcpServer) -> list[dict[str, object]]: ...
+
+    async def disconnect(self, server_id: str) -> None: ...
+
+    async def call_tool(
+        self,
+        server_id: str,
+        tool_name: str,
+        arguments: dict[str, object],
+    ) -> dict[str, object]: ...
+
+
+def mcp_tool_name(server_id: str, tool_name: str) -> str:
+    return f"mcp__{server_id}__{tool_name}"
+
+
+def parse_mcp_tool_name(name: str) -> tuple[str, str] | None:
+    if not name.startswith("mcp__"):
+        return None
+    parts = name.removeprefix("mcp__").split("__", 1)
+    if len(parts) != 2 or not parts[0] or not parts[1]:
+        return None
+    return parts[0], parts[1]
+
+
+def stable_mcp_server_id(name: str) -> str:
+    normalized = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")
+    return f"mcp-{normalized or 'server'}"
+
+
+def mcp_tools_from_result(raw_tools: list[dict[str, object]]) -> list[StoredMcpTool]:
+    tools: list[StoredMcpTool] = []
+    for raw_tool in raw_tools:
+        name = raw_tool.get("name")
+        if not isinstance(name, str) or not name:
+            continue
+        input_schema = raw_tool.get("inputSchema", raw_tool.get("input_schema", {}))
+        output_schema = raw_tool.get("outputSchema", raw_tool.get("output_schema"))
+        tools.append(
+            StoredMcpTool(
+                description=str(raw_tool.get("description") or ""),
+                input_schema=input_schema if isinstance(input_schema, dict) else {},
+                name=name,
+                output_schema=output_schema
+                if isinstance(output_schema, dict)
+                else None,
+            )
+        )
+    return tools
+
+
+def mcp_result_content(result: dict[str, object]) -> str:
+    content = result.get("content")
+    if isinstance(content, list):
+        parts: list[str] = []
+        for item in content:
+            if isinstance(item, dict) and isinstance(item.get("text"), str):
+                parts.append(item["text"])
+            elif hasattr(item, "text") and isinstance(item.text, str):
+                parts.append(item.text)
+        if parts:
+            return "\n".join(parts)
+    structured_content = result.get("structuredContent")
+    if structured_content is not None:
+        return json.dumps(structured_content, ensure_ascii=False)
+    return json.dumps(result, ensure_ascii=False)
+
+
+def mcp_result_is_error(result: dict[str, object]) -> bool:
+    return bool(result.get("isError") or result.get("is_error"))
+
+
+_template_pattern = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-(.*?))?\}")
+
+
+def expand_mcp_template(value: str, *, env: dict[str, str] | None = None) -> str:
+    lookup = env or os.environ
+
+    def replace(match: re.Match[str]) -> str:
+        name = match.group(1)
+        default = match.group(2)
+        return lookup.get(name, default or "")
+
+    return _template_pattern.sub(replace, value)
+
+
+def expand_mcp_value(value: Any, *, env: dict[str, str] | None = None) -> Any:
+    if isinstance(value, str):
+        return expand_mcp_template(value, env=env)
+    if isinstance(value, list):
+        return [expand_mcp_value(item, env=env) for item in value]
+    if isinstance(value, dict):
+        return {key: expand_mcp_value(item, env=env) for key, item in value.items()}
+    return value
+
+
+def expand_mcp_config(config: dict[str, object]) -> dict[str, object]:
+    expanded = expand_mcp_value(config) if config else {}
+    return expanded if isinstance(expanded, dict) else {}
+
+
+class DefaultMcpTransport:
+    def __init__(self) -> None:
+        self._sessions: dict[str, Any] = {}
+        self._stacks: dict[str, AsyncExitStack] = {}
+
+    async def connect(self, server: StoredMcpServer) -> list[dict[str, object]]:
+        from mcp import ClientSession
+        from mcp.client.stdio import stdio_client
+
+        await self.disconnect(server.id)
+        stack = AsyncExitStack()
+        config = expand_mcp_config(server.config)
+        if server.type == "url":
+            http_module = import_module("mcp.client.streamable_http")
+            http_headers = self._streamable_http_headers(config) or None
+            if hasattr(http_module, "streamablehttp_client"):
+                read_stream, write_stream, _ = await stack.enter_async_context(
+                    http_module.streamablehttp_client(
+                        server.url or str(config.get("url") or ""),
+                        headers=http_headers,
+                    )
+                )
+            else:
+                import httpx
+
+                http_client = await stack.enter_async_context(
+                    httpx.AsyncClient(headers=http_headers)
+                )
+                read_stream, write_stream, _ = await stack.enter_async_context(
+                    http_module.streamable_http_client(
+                        server.url or str(config.get("url") or ""),
+                        http_client=http_client,
+                    )
+                )
+        else:
+            read_stream, write_stream = await stack.enter_async_context(
+                stdio_client(
+                    self._stdio_parameters(server, config),
+                )
+            )
+        session = await stack.enter_async_context(
+            ClientSession(read_stream, write_stream)
+        )
+        await session.initialize()
+        result = await session.list_tools()
+        self._sessions[server.id] = session
+        self._stacks[server.id] = stack
+        return [self._model_dump(tool) for tool in result.tools]
+
+    def _streamable_http_headers(self, config: dict[str, object]) -> dict[str, str]:
+        headers: dict[str, str] = {}
+        for key in ("http_headers", "headers"):
+            raw_headers = config.get(key)
+            if isinstance(raw_headers, dict):
+                headers.update(
+                    {
+                        str(header_name): str(header_value)
+                        for header_name, header_value in raw_headers.items()
+                    }
+                )
+                break
+
+        raw_env_headers = config.get("env_http_headers") or config.get("envHeaders")
+        if isinstance(raw_env_headers, dict):
+            for header_name, env_name in raw_env_headers.items():
+                if isinstance(env_name, str):
+                    env_value = os.environ.get(env_name)
+                    if env_value is not None:
+                        headers[str(header_name)] = env_value
+
+        bearer_token_env_var = config.get("bearer_token_env_var") or config.get(
+            "bearerTokenEnvVar"
+        )
+        if isinstance(bearer_token_env_var, str):
+            env_bearer_token = os.environ.get(bearer_token_env_var)
+            if env_bearer_token and "Authorization" not in headers:
+                headers["Authorization"] = f"Bearer {env_bearer_token}"
+
+        bearer_token: object = config.get("bearer_token") or config.get("bearerToken")
+        if (
+            isinstance(bearer_token, str)
+            and bearer_token
+            and "Authorization" not in headers
+        ):
+            headers["Authorization"] = f"Bearer {bearer_token}"
+        return headers
+
+    def _stdio_parameters(
+        self,
+        server: StoredMcpServer,
+        config: dict[str, object],
+    ) -> Any:
+        from mcp import StdioServerParameters
+
+        env: dict[str, str] = {}
+        raw_env = config.get("env")
+        if isinstance(raw_env, dict):
+            for key, value in raw_env.items():
+                if isinstance(key, str) and isinstance(value, str):
+                    env[key] = value
+        raw_env_vars = config.get("env_vars")
+        if isinstance(raw_env_vars, list):
+            for key in raw_env_vars:
+                if isinstance(key, str) and key not in env:
+                    value = os.environ.get(key)
+                    if value is not None:
+                        env[key] = value
+        cwd = config.get("cwd")
+        raw_args = config.get("args")
+        config_args = raw_args if isinstance(raw_args, list) else []
+        return StdioServerParameters(
+            command=server.command or str(config.get("command") or ""),
+            args=server.args
+            or [str(argument) for argument in config_args if isinstance(argument, str)],
+            cwd=cwd if isinstance(cwd, str) else None,
+            env=env or None,
+        )
+
+    async def disconnect(self, server_id: str) -> None:
+        stack = self._stacks.pop(server_id, None)
+        self._sessions.pop(server_id, None)
+        if stack is not None:
+            await stack.aclose()
+
+    async def call_tool(
+        self,
+        server_id: str,
+        tool_name: str,
+        arguments: dict[str, object],
+    ) -> dict[str, object]:
+        session = self._sessions.get(server_id)
+        if session is None:
+            raise RuntimeError("Server is not connected.")
+        result = await session.call_tool(tool_name, arguments=arguments)
+        return self._model_dump(result)
+
+    def _model_dump(self, value: Any) -> dict[str, object]:
+        if hasattr(value, "model_dump"):
+            dumped = value.model_dump(by_alias=True)
+            return dumped if isinstance(dumped, dict) else {}
+        if isinstance(value, dict):
+            return value
+        return {}
+
+
+class McpManager:
+    def __init__(
+        self,
+        *,
+        store: StateStore,
+        transport: McpTransport | None = None,
+    ) -> None:
+        self.store = store
+        self.transport = transport or DefaultMcpTransport()
+        self._status_by_server: dict[str, str] = {}
+        self._error_by_server: dict[str, str] = {}
+        self._tools_by_server: dict[str, list[StoredMcpTool]] = {}
+        self._server_names: dict[str, str] = {}
+        self._connect_tasks: dict[str, asyncio.Task[None]] = {}
+
+    async def start_enabled(self) -> None:
+        for server in self.store.read_mcp_servers():
+            if server.enabled:
+                self.schedule_connect_server(server)
+            else:
+                await self.disconnect_server(server.id)
+
+    async def stop_all(self) -> None:
+        tasks = list(self._connect_tasks.values())
+        self._connect_tasks.clear()
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        for server_id in list(self._status_by_server):
+            await self.transport.disconnect(server_id)
+        self._status_by_server.clear()
+        self._error_by_server.clear()
+        self._tools_by_server.clear()
+        self._server_names.clear()
+
+    async def sync_server(self, server: StoredMcpServer) -> StoredMcpServer:
+        await self.cancel_connect_server(server.id)
+        if not server.enabled:
+            await self.disconnect_server(server.id)
+            return self.server_with_status(server)
+        await self.connect_server(server)
+        return self.server_with_status(server)
+
+    def schedule_connect_server(self, server: StoredMcpServer) -> None:
+        task = self._connect_tasks.pop(server.id, None)
+        if task is not None and not task.done():
+            task.cancel()
+        self._server_names[server.id] = server.name
+        self._status_by_server[server.id] = "starting"
+        self._error_by_server[server.id] = ""
+        connect_task = asyncio.create_task(self.connect_server(server))
+        self._connect_tasks[server.id] = connect_task
+        connect_task.add_done_callback(self._connect_task_callback(server.id))
+
+    def _connect_task_callback(
+        self,
+        server_id: str,
+    ) -> Callable[[asyncio.Task[None]], None]:
+        def finish(completed_task: asyncio.Task[None]) -> None:
+            self._finish_connect_task(server_id, completed_task)
+
+        return finish
+
+    def _finish_connect_task(
+        self,
+        server_id: str,
+        task: asyncio.Task[None],
+    ) -> None:
+        if self._connect_tasks.get(server_id) is task:
+            self._connect_tasks.pop(server_id, None)
+        if task.cancelled():
+            return
+        try:
+            task.result()
+        except Exception:
+            logger.exception("MCP server background connect failed")
+
+    async def cancel_connect_server(self, server_id: str) -> None:
+        task = self._connect_tasks.pop(server_id, None)
+        if task is None or task.done():
+            return
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+
+    async def reconnect_server(self, server_id: str) -> StoredMcpServer:
+        server = self.find_server(server_id)
+        await self.cancel_connect_server(server_id)
+        await self.transport.disconnect(server_id)
+        if server.enabled:
+            await self.connect_server(server)
+        else:
+            await self.disconnect_server(server_id)
+        return self.server_with_status(server)
+
+    async def delete_server(self, server_id: str) -> None:
+        await self.cancel_connect_server(server_id)
+        self.store.delete_mcp_server(server_id)
+        try:
+            await self.transport.disconnect(server_id)
+        except Exception:
+            logger.exception("MCP server disconnect failed during delete")
+        self._status_by_server.pop(server_id, None)
+        self._error_by_server.pop(server_id, None)
+        self._tools_by_server.pop(server_id, None)
+        self._server_names.pop(server_id, None)
+
+    async def reload(self) -> list[StoredMcpServer]:
+        await self.stop_all()
+        await self.start_enabled()
+        return self.servers_with_status(self.store.read_mcp_servers())
+
+    def find_server(self, server_id: str) -> StoredMcpServer:
+        for server in self.store.read_mcp_servers():
+            if server.id == server_id:
+                return server
+        raise KeyError(server_id)
+
+    async def connect_server(self, server: StoredMcpServer) -> None:
+        self._server_names[server.id] = server.name
+        self._status_by_server[server.id] = "starting"
+        self._error_by_server[server.id] = ""
+        try:
+            raw_tools = await asyncio.wait_for(
+                self.transport.connect(server),
+                timeout=MCP_CONNECT_TIMEOUT_SECONDS,
+            )
+        except TimeoutError:
+            self._status_by_server[server.id] = "error"
+            self._error_by_server[server.id] = "Connection timed out."
+            self._tools_by_server[server.id] = []
+            return
+        except Exception as error:
+            self._status_by_server[server.id] = "error"
+            self._error_by_server[server.id] = str(error)
+            self._tools_by_server[server.id] = []
+            return
+        tools = mcp_tools_from_result(raw_tools)
+        self.store.save_mcp_tools(server.id, tools)
+        self._tools_by_server[server.id] = tools
+        self._status_by_server[server.id] = "ready"
+        self._error_by_server[server.id] = ""
+
+    async def disconnect_server(self, server_id: str) -> None:
+        await self.transport.disconnect(server_id)
+        self._status_by_server[server_id] = "disabled"
+        self._error_by_server[server_id] = ""
+        self._tools_by_server[server_id] = []
+
+    def server_with_status(self, server: StoredMcpServer) -> StoredMcpServer:
+        if not server.enabled:
+            status = "disabled"
+            error = ""
+            tools: list[StoredMcpTool] = []
+        else:
+            status = self._status_by_server.get(server.id, server.status)
+            error = self._error_by_server.get(server.id, server.error)
+            tools = self._tools_by_server.get(server.id, server.tools)
+        self._server_names[server.id] = server.name
+        return server.model_copy(
+            update={
+                "error": error,
+                "status": status,
+                "tools": tools,
+            }
+        )
+
+    def servers_with_status(
+        self, servers: list[StoredMcpServer]
+    ) -> list[StoredMcpServer]:
+        return [self.server_with_status(server) for server in servers]
+
+    def tool_specs(self) -> list[dict[str, object]]:
+        specs: list[dict[str, object]] = []
+        for server in self.servers_with_status(self.store.read_mcp_servers()):
+            if server.status != "ready":
+                continue
+            for tool in server.tools:
+                specs.append(
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": mcp_tool_name(server.id, tool.name),
+                            "description": tool.description
+                            or f"Call {server.name}.{tool.name}.",
+                            "parameters": tool.input_schema or {"type": "object"},
+                        },
+                    }
+                )
+        return specs
+
+    def tool_title(self, name: str) -> str | None:
+        parsed = parse_mcp_tool_name(name)
+        if parsed is None:
+            return None
+        server_id, tool_name = parsed
+        return f"Calling {self._server_names.get(server_id, server_id)}.{tool_name}"
+
+    async def run_tool(
+        self, name: str, arguments: dict[str, object]
+    ) -> ToolResult | None:
+        parsed = parse_mcp_tool_name(name)
+        if parsed is None:
+            return None
+        server_id, tool_name = parsed
+        result = await self.transport.call_tool(server_id, tool_name, arguments)
+        content = mcp_result_content(result)
+        server_name = self._server_names.get(server_id, server_id)
+        return ToolResult(
+            content=content,
+            data={
+                "server": server_name,
+                "tool": tool_name,
+                "result": result,
+            },
+            ok=not mcp_result_is_error(result),
+            title=f"Calling {server_name}.{tool_name}",
+        )

package/backend/src/flowent/mcp_import.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import json
+import tomllib
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from flowent.mcp import stable_mcp_server_id
+from flowent.storage import StoredMcpServer
+
+McpImportSource = Literal["claude_code", "codex"]
+
+
+class McpImportSourceResult(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    error: str = ""
+    path: str
+    servers: list[StoredMcpServer] = Field(default_factory=list)
+    source: McpImportSource
+
+
+class McpImportDiscovery(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    servers: list[StoredMcpServer] = Field(default_factory=list)
+    sources: list[McpImportSourceResult] = Field(default_factory=list)
+
+
+def discover_imported_mcp_servers(
+    cwd: Path | None = None,
+    home: Path | None = None,
+    source: McpImportSource | None = None,
+) -> McpImportDiscovery:
+    workspace = (cwd or Path.cwd()).resolve(strict=False)
+    user_home = (home or Path.home()).resolve(strict=False)
+    sources: list[McpImportSourceResult] = []
+
+    for path, config_source in candidate_mcp_config_files(
+        workspace,
+        user_home,
+        source,
+    ):
+        if not path.is_file():
+            continue
+        try:
+            servers = parse_mcp_config_file(path, config_source, workspace)
+            sources.append(
+                McpImportSourceResult(
+                    path=str(path.resolve(strict=False)),
+                    servers=servers,
+                    source=config_source,
+                )
+            )
+        except Exception as error:
+            sources.append(
+                McpImportSourceResult(
+                    error=str(error),
+                    path=str(path.resolve(strict=False)),
+                    source=config_source,
+                )
+            )
+
+    return McpImportDiscovery(
+        servers=dedupe_mcp_servers(
+            server for source_result in sources for server in source_result.servers
+        ),
+        sources=sources,
+    )
+
+
+def candidate_mcp_config_files(
+    cwd: Path,
+    home: Path,
+    source: McpImportSource | None = None,
+) -> list[tuple[Path, McpImportSource]]:
+    candidates: list[tuple[Path, McpImportSource]] = []
+    if source in (None, "claude_code"):
+        candidates.extend(
+            [
+                (cwd / ".mcp.json", "claude_code"),
+                (cwd / ".claude" / "settings.local.json", "claude_code"),
+                (cwd / ".claude" / "settings.json", "claude_code"),
+                (home / ".claude.json", "claude_code"),
+                (home / ".claude" / "settings.json", "claude_code"),
+            ]
+        )
+    if source in (None, "codex"):
+        candidates.extend(
+            [
+                (cwd / ".codex" / "config.toml", "codex"),
+                (home / ".codex" / "config.toml", "codex"),
+            ]
+        )
+    seen: set[tuple[Path, McpImportSource]] = set()
+    unique_candidates: list[tuple[Path, McpImportSource]] = []
+    for path, source in candidates:
+        key = (path.resolve(strict=False), source)
+        if key in seen:
+            continue
+        seen.add(key)
+        unique_candidates.append((path, source))
+    return unique_candidates
+
+
+def parse_mcp_config_file(
+    path: Path,
+    source: McpImportSource,
+    cwd: Path,
+) -> list[StoredMcpServer]:
+    if source == "codex":
+        payload = tomllib.loads(path.read_text(encoding="utf-8"))
+        return parse_codex_mcp_servers(payload)
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    return parse_claude_code_mcp_servers(payload, cwd)
+
+
+def parse_codex_mcp_servers(payload: object) -> list[StoredMcpServer]:
+    if not isinstance(payload, dict):
+        return []
+    raw_servers = payload.get("mcp_servers")
+    if not isinstance(raw_servers, dict):
+        return []
+    return servers_from_map(raw_servers)
+
+
+def parse_claude_code_mcp_servers(payload: object, cwd: Path) -> list[StoredMcpServer]:
+    if not isinstance(payload, dict):
+        return []
+    server_maps: list[dict[object, object]] = []
+    projects = payload.get("projects")
+    if isinstance(projects, dict):
+        workspace_keys = [str(cwd.resolve(strict=False)), str(cwd)]
+        for workspace_key in dict.fromkeys(workspace_keys):
+            project_config = projects.get(workspace_key)
+            if isinstance(project_config, dict):
+                project_servers = project_config.get("mcpServers")
+                if isinstance(project_servers, dict):
+                    server_maps.append(project_servers)
+    top_level_servers = payload.get("mcpServers")
+    if isinstance(top_level_servers, dict):
+        server_maps.append(top_level_servers)
+
+    return dedupe_mcp_servers(
+        server for server_map in server_maps for server in servers_from_map(server_map)
+    )
+
+
+def servers_from_map(raw_servers: dict[object, object]) -> list[StoredMcpServer]:
+    servers: list[StoredMcpServer] = []
+    for raw_name, raw_config in raw_servers.items():
+        if not isinstance(raw_config, dict):
+            continue
+        server = server_from_config(str(raw_name), raw_config)
+        if server is not None:
+            servers.append(server)
+    return dedupe_mcp_servers(servers)
+
+
+def server_from_config(
+    name: str,
+    raw_config: dict[object, object],
+) -> StoredMcpServer | None:
+    config = {str(key): value for key, value in raw_config.items()}
+    url = string_config(config, "url")
+    command = string_config(config, "command")
+    args = string_list_config(config, "args")
+    server_type = "url" if url else "command"
+    if server_type == "command" and not command:
+        return None
+    enabled = enabled_config(config)
+    return StoredMcpServer(
+        args=args if server_type == "command" else [],
+        command=command if server_type == "command" else "",
+        config=config,
+        enabled=enabled,
+        id=stable_mcp_server_id(name),
+        name=name,
+        type=server_type,
+        url=url if server_type == "url" else "",
+    )
+
+
+def string_config(config: dict[str, object], key: str) -> str:
+    value = config.get(key)
+    return value if isinstance(value, str) else ""
+
+
+def string_list_config(config: dict[str, object], key: str) -> list[str]:
+    value = config.get(key)
+    if not isinstance(value, list):
+        return []
+    return [item for item in value if isinstance(item, str)]
+
+
+def enabled_config(config: dict[str, object]) -> bool:
+    enabled = config.get("enabled")
+    if isinstance(enabled, bool):
+        return enabled
+    disabled = config.get("disabled")
+    if isinstance(disabled, bool):
+        return not disabled
+    return True
+
+
+def dedupe_mcp_servers(servers: Iterable[StoredMcpServer]) -> list[StoredMcpServer]:
+    unique_servers: list[StoredMcpServer] = []
+    seen_ids: set[str] = set()
+    for server in servers:
+        if server.id in seen_ids:
+            continue
+        seen_ids.add(server.id)
+        unique_servers.append(server)
+    return unique_servers