90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
import asyncio
|
|
import logging
|
|
|
|
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])
|
|
return result.structured_content or result.content
|
|
|
|
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 |