Artemis-Eapil/artemis/plugins/session.py

"""Task-owned external MCP transport. A timed-out write is never replayed."""

from __future__ import annotations

import asyncio
from contextlib import AsyncExitStack
import os
from pathlib import Path
import re
from typing import Any

import httpx
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client

from artemis.config.settings import settings
from artemis.plugins.models import PluginManifest, PluginProfile, ServerConfig


def resolve_server(
    manifest: PluginManifest, profile: PluginProfile, directory: Path
) -> ServerConfig:
    if manifest.server is None:
        raise ValueError("此插件不包含 MCP 服务")
    values = {name: field.default for name, field in manifest.config.items() if not field.secret}
    values.update(profile.values)

    def substitute(value: Any) -> Any:
        if isinstance(value, dict):
            return {k: substitute(v) for k, v in value.items()}
        if isinstance(value, list):
            return [substitute(v) for v in value]
        if not isinstance(value, str):
            return value

        def replace(match: re.Match) -> str:
            key = match.group(1)
            if key == "plugin_dir":
                return str(directory)
            if key.startswith("config.") and key[7:] in values:
                return values[key[7:]]
            if key.startswith("secret.") and key[7:] in profile.secrets:
                return profile.secrets[key[7:]].get_secret_value()
            raise ValueError(f"插件配置引用无法解析：{key}")

        return re.sub(r"\$\{([^}]+)\}", replace, value)

    return ServerConfig.model_validate(substitute(manifest.server.model_dump()))


class PluginSession:
    """One owner task enters/exits anyio transport scopes, even during cancellation."""

    def __init__(self, server: ServerConfig, directory: Path):
        self.server = server
        self.directory = directory
        self.tools: list[Any] = []
        self._queue: asyncio.Queue = asyncio.Queue()
        self._ready = asyncio.Event()
        self._owner: asyncio.Task | None = None
        self._error: BaseException | None = None
        self.stop_uncertain = False

    async def start(self) -> PluginSession:
        self._owner = asyncio.create_task(self._run(), name="artemis-plugin-session")
        try:
            await asyncio.wait_for(
                self._ready.wait(), settings.ARTEMIS_PLUGIN_START_TIMEOUT_SECONDS
            )
            if self._error:
                raise RuntimeError(
                    "插件 MCP 启动失败，请检查命令、连接地址及必填配置"
                ) from self._error
            return self
        except BaseException:
            await self.aclose()
            raise

    async def _run(self):
        active = None
        try:
            async with AsyncExitStack() as stack:
                if self.server.transport == "stdio":
                    # A plugin can print credentials to stderr; never forward it to task reports.
                    errlog = stack.enter_context(open(os.devnull, "w"))
                    streams = await stack.enter_async_context(
                        stdio_client(
                            StdioServerParameters(
                                command=self.server.command,
                                args=self.server.args,
                                env=self.server.env,
                                cwd=str(self.directory),
                            ),
                            errlog=errlog,
                        )
                    )
                else:
                    client = await stack.enter_async_context(
                        httpx.AsyncClient(
                            headers=self.server.headers,
                            timeout=settings.ARTEMIS_PLUGIN_TOOL_TIMEOUT_SECONDS,
                            follow_redirects=False,
                        )
                    )
                    streams = await stack.enter_async_context(
                        streamable_http_client(self.server.url, http_client=client)
                    )
                session = await stack.enter_async_context(ClientSession(streams[0], streams[1]))
                await session.initialize()
                cursor = None
                for _ in range(20):
                    response = await session.list_tools(cursor=cursor)
                    self.tools.extend(response.tools)
                    cursor = response.nextCursor
                    if not cursor:
                        break
                else:
                    raise ValueError("插件工具列表分页超过限制")
                self._ready.set()
                while True:
                    item = await self._queue.get()
                    if item is None:
                        return
                    active, name, arguments = item
                    if active.cancelled():
                        active = None
                        continue
                    # On timeout retire the whole connection: MCP may deliver a late response.
                    result = await asyncio.wait_for(
                        session.call_tool(name, arguments),
                        settings.ARTEMIS_PLUGIN_TOOL_TIMEOUT_SECONDS,
                    )
                    if not active.done():
                        active.set_result(result)
                    active = None
        except BaseException as exc:
            self._error = exc
            if active is not None and self.server.transport == "streamable_http":
                self.stop_uncertain = True
        finally:
            self._ready.set()
            if active is not None and not active.done():
                active.set_exception(
                    RuntimeError(
                        "插件调用中断或超时，执行结果不确定；请先观察状态，不能盲目重放操作"
                    )
                )
            while not self._queue.empty():
                item = self._queue.get_nowait()
                if item is not None and not item[0].done():
                    item[0].set_exception(RuntimeError("插件连接已关闭，此调用尚未执行"))

    async def call(self, name: str, arguments: dict):
        if self._owner is None or self._owner.done() or self._error is not None:
            raise RuntimeError("插件连接已关闭，不能继续执行")
        future = asyncio.get_running_loop().create_future()
        self._queue.put_nowait((future, name, arguments))
        try:
            # Bound the caller as well as the protocol operation: transport context
            # cleanup may itself stall after the inner request times out.
            return await asyncio.wait_for(
                future,
                settings.ARTEMIS_PLUGIN_TOOL_TIMEOUT_SECONDS
                + settings.ARTEMIS_PLUGIN_SHUTDOWN_TIMEOUT_SECONDS,
            )
        except TimeoutError:
            await self.aclose()
            raise RuntimeError("插件响应或清理超时，执行结果不确定") from None
        except asyncio.CancelledError:
            # Guidance/pause must stop the transport before the task releases shared resources.
            await self.aclose()
            raise

    async def aclose(self):
        owner = self._owner
        if owner is None or owner.done():
            return
        owner.cancel()
        done, _ = await asyncio.wait(
            {owner}, timeout=settings.ARTEMIS_PLUGIN_SHUTDOWN_TIMEOUT_SECONDS
        )
        if not done:
            self.stop_uncertain = True
            raise RuntimeError("插件尚未确认停止，不能释放其占用资源")