agents/common/youtrack_mcp_client.py

import asyncio
import json
import logging
from typing import Optional, Any

import httpx
from contextlib import AsyncExitStack
from mcp import Client, MCPError
from mcp.client.streamable_http import streamable_http_client

from common.as_type import _extract_text, _looks_like_not_found

logger = logging.getLogger("youtrack_mcp")
logger.setLevel(logging.DEBUG)

class IssueNotFound(Exception):
    """Raised when a YouTrack issue doesn't exist or isn't accessible."""

class YouTrackMCPClient:
    """Wraps an MCP connection to YouTrack, owning the lifecycle in a dedicated task."""

    def __init__(self, endpoint_url: str, token: str, proxy: str | None = None):
        self.endpoint_url = endpoint_url
        self.token = token
        self.proxy = proxy
        self._client = None
        self._lifecycle_task = None
        self._ready = asyncio.Event()
        self._stop = asyncio.Event()

    async def _run(self):
        """Single task that owns the entire MCP client lifecycle."""
        async with AsyncExitStack() as stack:
            http = await stack.enter_async_context(
                httpx.AsyncClient(
                    headers={"Authorization": f"Bearer {self.token}"},
                    proxy=self.proxy,
                    timeout=httpx.Timeout(30.0, connect=10.0, read=300.0),
                    follow_redirects=True,
                )
            )
            transport = streamable_http_client(
                self.endpoint_url,
                http_client=http,
            )
            self._client = await stack.enter_async_context(Client(transport))
            self._ready.set()
            await self._stop.wait()  # keep the context alive

    async def connect(self):
        """Start the lifecycle task and wait until the client is ready."""
        self._lifecycle_task = asyncio.create_task(self._run())
        await self._ready.wait()
        return self

    async def close(self):
        """Signal the lifecycle task to exit and wait for it to finish."""
        if self._lifecycle_task is None:
            return
        self._stop.set()
        await self._lifecycle_task
        self._lifecycle_task = None

    async def call_tool(self, name: str, arguments: dict):
        logger.debug("MCP CALL: %s with %s", name, arguments)
        result = await self._call_tool_or_raise(name, arguments)
        logger.debug("MCP RAW RESULT: %s", repr(result)[:500])
        if result.structured_content is not None:
            return result.structured_content

        # Convert content blocks to JSON-serializable form
        if result.content:
            out = []
            for block in result.content:
                if getattr(block, "type", None) == "text":
                    out.append({"type": "text", "text": block.text})
                else:
                    # fallback for other block types
                    out.append({"type": getattr(block, "type", "unknown"),
                                "data": str(block)})
            # If it's a single text block, unwrap to plain string
            if len(out) == 1 and out[0]["type"] == "text":
                return out[0]["text"]
            return out

        return None

    async def _call_tool_or_raise(self, name: str, arguments: dict):
        try:
            result = await self._client.call_tool(name, arguments)
            logger.debug("call RESULT: %s", result)
        except MCPError as e:
            # Drill into cause chain for HTTP status
            cause = e.__cause__
            while cause:
                if isinstance(cause, httpx.HTTPStatusError):
                    if cause.response.status_code == 404:
                        raise IssueNotFound(f"{arguments} not found") from e
                    raise
                cause = cause.__cause__
            raise

        # MCP-level error flag
        if getattr(result, "is_error", False):
            text = _extract_text(result)
            if _looks_like_not_found(text):
                raise IssueNotFound(text)
            raise RuntimeError(f"MCP tool {name} failed: {text}")

        return result

    async def list_tools(self):
        """List all available tools from the MCP server."""
        return await self._client.list_tools()

    @classmethod
    def extract_comment_id(cls, result: Any) -> Optional[str]:
        """
        Best-effort extraction of the created comment's ID from a
        CallToolResult (or any MCP-ish response).

        Returns None if the ID can't be determined. Callers MUST tolerate
        None by falling back to marker-based lookup on the next Gather.
        """
        # 1. Unwrap the MCP envelope.
        content = getattr(result, "content", None)
        if content is None and isinstance(result, dict):
            content = result.get("content")

        if not content:
            return None

        # 2. Take the first text block.
        first = content[0]
        text = getattr(first, "text", None)
        if text is None and isinstance(first, dict):
            text = first.get("text")

        if not text:
            return None

        # 3. Try JSON first.
        try:
            data = json.loads(text or "{}")
        except json.JSONDecodeError:
            # 4. Not JSON — assume the server returned a bare ID.
            return text.strip() or None

        # 5. Common shapes: {"id": ...} or {"comment": {"id": ...}}.
        if isinstance(data, dict):
            if "id" in data:
                return str(data["id"])
            if "comment" in data and isinstance(data["comment"], dict):
                return str(data["comment"].get("id")) or None
            # YouTrack often uses "idReadable" for the human-facing ID.
            for key in ("idReadable", "commentId", "comment_id"):
                if key in data:
                    return str(data[key])

        return None