PMFlow/backend/app/services/mcp_config/runtime.py

from __future__ import annotations

import json
from typing import Any, Callable

import httpx

from ...utils.timezone import now_cn_naive
from ..external_query_service.schemas import ExternalEvidenceItem


class MicrosoftLearnMCPRuntime:
    def __init__(
        self,
        *,
        endpoint_url: str = "https://learn.microsoft.com/api/mcp",
        client_factory: Callable[[], httpx.AsyncClient] | None = None,
    ) -> None:
        self.endpoint_url = str(endpoint_url or "https://learn.microsoft.com/api/mcp").strip()
        self.client_factory = client_factory or (
            lambda: httpx.AsyncClient(timeout=20.0, trust_env=False, follow_redirects=True)
        )

    async def probe(self) -> dict[str, Any]:
        checked_at = now_cn_naive().isoformat()
        try:
            session_id, _ = await self._initialize()
            _, tools_payload = await self._post_jsonrpc(
                {
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list",
                    "params": {},
                },
                session_id=session_id,
            )
            tools = list((tools_payload.get("result") or {}).get("tools") or [])
            tool_names = [str(item.get("name") or "").strip() for item in tools if str(item.get("name") or "").strip()]
            return {
                "status": "ready",
                "checked_at": checked_at,
                "message": "PMFlow backend can directly reach the remote MCP server.",
                "tool_names": tool_names,
            }
        except Exception as exc:
            return {
                "status": "error",
                "checked_at": checked_at,
                "message": f"PMFlow backend probe failed: {exc}",
                "tool_names": [],
            }

    async def search(self, *, query: str, result_limit: int = 3) -> list[ExternalEvidenceItem]:
        session_id, _ = await self._initialize()
        _, payload = await self._post_jsonrpc(
            {
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "microsoft_docs_search",
                    "arguments": {
                        "query": str(query or "").strip(),
                    },
                },
            },
            session_id=session_id,
        )
        text_payload = ""
        for item in list((payload.get("result") or {}).get("content") or []):
            if str(item.get("type") or "").strip() == "text":
                text_payload = str(item.get("text") or "").strip()
                if text_payload:
                    break
        parsed = json.loads(text_payload) if text_payload else {}
        results = list(parsed.get("results") or [])
        retrieved_at = now_cn_naive().isoformat()
        evidence: list[ExternalEvidenceItem] = []
        for raw in results[: max(1, int(result_limit or 3))]:
            target_url = str(raw.get("contentUrl") or "").strip()
            if not target_url.startswith("http://") and not target_url.startswith("https://"):
                continue
            title = str(raw.get("title") or "").strip()
            snippet = str(raw.get("content") or "").strip()
            evidence.append(
                ExternalEvidenceItem(
                    source_type="mcp_search",
                    source_label="Microsoft Learn MCP",
                    title=title,
                    target_url=target_url,
                    snippet=snippet[:1800],
                    retrieved_at=retrieved_at,
                    freshness_score=0.86,
                    authority_score=0.98,
                    grounded_fields=["microsoft_docs_search"],
                    raw_metadata={
                        "provider": "microsoft_learn_mcp",
                        "result_id": str(raw.get("id") or "").strip() or None,
                    },
                )
            )
        return evidence

    async def _initialize(self) -> tuple[str | None, dict[str, Any]]:
        return await self._post_jsonrpc(
            {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-03-26",
                    "capabilities": {},
                    "clientInfo": {
                        "name": "pmflow",
                        "version": "0.1.0",
                    },
                },
            }
        )

    async def _post_jsonrpc(
        self,
        payload: dict[str, Any],
        *,
        session_id: str | None = None,
    ) -> tuple[str | None, dict[str, Any]]:
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if session_id:
            headers["mcp-session-id"] = session_id
        async with self.client_factory() as client:
            response = await client.post(self.endpoint_url, headers=headers, json=payload)
            response.raise_for_status()
        next_session = str(response.headers.get("mcp-session-id") or session_id or "").strip() or None
        body = response.text
        parsed = self._extract_payload(body)
        return next_session, parsed

    @staticmethod
    def _extract_payload(body: str) -> dict[str, Any]:
        data_lines: list[str] = []
        for line in str(body or "").splitlines():
            if line.startswith("data:"):
                data_lines.append(line[len("data:") :].lstrip())
        if not data_lines:
            raise ValueError("MCP streamable_http response missing data payload")
        return json.loads("\n".join(data_lines))