codeapi-client 0.4.1__py3-none-any.whl

codeapi/client/_async/_mcp.py

@@ -0,0 +1,445 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, List, Literal
+
+import httpx
+
+from codeapi.types import (
+    CodeInfo,
+    CodeType,
+    CodeZip,
+    DataZip,
+    EnvVars,
+    Job,
+    JobStage,
+    JobStatus,
+    JobType,
+    JsonData,
+)
+
+if TYPE_CHECKING:
+    from fastmcp.mcp_config import MCPConfig
+
+    from . import AsyncClient
+
+
+class AsyncStoredMcpClient:
+    def __init__(self, client: AsyncClient):
+        self._client = client
+
+    async def run(
+        self,
+        code_id: str,
+        mcp_module: str,
+        env_vars: EnvVars | dict | None = None,
+        transport: Literal["streamable-http", "http", "sse"] = "streamable-http",
+        data_zip: DataZip | None = None,
+    ) -> Job:
+        """Runs a stored MCP server.
+
+        Args:
+            code_id (str): The code ID.
+            mcp_module (str): The name of the MCP server module.
+            env_vars (EnvVars | dict | None): Optional environment variables.
+            transport (Literal["streamable-http", "http", "sse"]): The transport protocol.
+            data_zip (DataZip | None): Optional data zip file.
+
+        Returns:
+            Job: The created job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/code/{code_id}/run/mcp"
+
+        files = self._client._prepare_files(data_zip=data_zip)
+        data: dict[str, Any] = {
+            "mcp_module": mcp_module,
+            "transport": transport,
+        }
+        if env_vars:
+            data["env_vars"] = EnvVars(env_vars).json_str
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    url,
+                    headers=self._client.api_key_header,
+                    files=files if files else None,
+                    data=data,
+                )
+                response.raise_for_status()
+                return Job(**response.json())
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def list_info(self) -> list[CodeInfo]:
+        """List all stored MCP code.
+
+        Returns:
+            list[CodeInfo]: List of MCP code info.
+        """
+        return await self._client.code.list_info(code_type=CodeType.MCP)
+
+    async def delete(self, code_id: str) -> str:
+        """Delete stored MCP code.
+
+        Args:
+            code_id (str): The code ID to delete.
+
+        Returns:
+            str: Deletion confirmation message.
+
+        Raises:
+            ValueError: If the code_id is not MCP code.
+        """
+        # Verify this is actually MCP code
+        code_info = await self._client.code.get_info(code_id)
+        if code_info.code_type != CodeType.MCP:
+            raise ValueError(
+                f"Code '{code_id}' is {code_info.code_type}, not MCP code. "
+                "Cannot delete non-MCP code from MCP client."
+            )
+
+        return await self._client.code.delete(code_id)
+
+
+class AsyncMcpConfigClient:
+    def __init__(self, client: AsyncClient):
+        self._client = client
+
+    async def run(
+        self,
+        mcp_config: MCPConfig | Path | str | dict,
+        transport: Literal["streamable-http", "http", "sse"] = "streamable-http",
+    ) -> Job:
+        """Runs a MCP server from MCP configuration.
+
+        Args:
+            mcp_config (MCPConfig | Path | str | dict): The MCP configuration. Can be:
+                - FastMCP MCPConfig object
+                - Path to a .json file containing MCP configuration
+                - JSON string containing MCP configuration
+                - Dictionary containing MCP configuration
+            transport (Literal["streamable-http", "http", "sse"]): The transport protocol.
+
+        Returns:
+            Job: The created job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/config/run/mcp"
+
+        if isinstance(mcp_config, (Path, str)):
+            if str(mcp_config).endswith(".json"):
+                config_path = Path(mcp_config)
+                with open(config_path, "r") as f:
+                    config_data = json.load(f)
+            else:
+                config_data = json.loads(str(mcp_config))
+        elif isinstance(mcp_config, dict):
+            config_data = mcp_config
+        else:
+            config_data = mcp_config.model_dump()
+
+        data = {
+            "transport": transport,
+        }
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    url,
+                    headers=self._client.api_key_header,
+                    json=config_data,
+                    params=data,
+                )
+                response.raise_for_status()
+                return Job(**response.json())
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+
+class AsyncMcpJobsClient:
+    def __init__(self, client):
+        self._client = client
+
+    async def list(
+        self,
+        job_status: JobStatus | None = None,
+        job_stage: JobStage | None = None,
+    ) -> List[Job]:
+        """List MCP jobs.
+
+        Args:
+            job_status (JobStatus | None): Filter by job status.
+            job_stage (JobStage | None): Filter by job stage.
+
+        Returns:
+            list[Job]: List of MCP jobs.
+        """
+        return await self._client.jobs.list(
+            job_type=JobType.RUN_MCP,
+            job_status=job_status,
+            job_stage=job_stage,
+        )
+
+    async def get_latest(self) -> Job | None:
+        """Get the most recent MCP job.
+
+        Returns:
+            Job | None: The most recent MCP job, or None if no jobs exist.
+        """
+        jobs = await self.list()
+        return jobs[0] if jobs else None
+
+    async def list_queued(self) -> List[Job]:
+        """Get all queued MCP jobs.
+
+        Returns:
+            List[Job]: List of queued MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.QUEUED)
+
+    async def list_scheduled(self) -> List[Job]:
+        """Get all scheduled MCP jobs.
+
+        Returns:
+            List[Job]: List of scheduled MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.SCHEDULED)
+
+    async def list_started(self) -> List[Job]:
+        """Get all started MCP jobs.
+
+        Returns:
+            list[Job]: List of started MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.STARTED)
+
+    async def list_deferred(self) -> List[Job]:
+        """Get all deferred MCP jobs.
+
+        Returns:
+            List[Job]: List of deferred MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.DEFERRED)
+
+    async def list_canceled(self) -> List[Job]:
+        """Get all canceled MCP jobs.
+
+        Returns:
+            List[Job]: List of canceled MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.CANCELED)
+
+    async def list_stopped(self) -> List[Job]:
+        """Get all stopped MCP jobs.
+
+        Returns:
+            List[Job]: List of stopped MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.STOPPED)
+
+    async def list_failed(self) -> List[Job]:
+        """Get all failed MCP jobs.
+
+        Returns:
+            list[Job]: List of failed MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.FAILED)
+
+    async def list_finished(self) -> List[Job]:
+        """Get all finished MCP jobs.
+
+        Returns:
+            list[Job]: List of finished MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.FINISHED)
+
+    async def list_timed_out(self) -> List[Job]:
+        """Get all timed out MCP jobs.
+
+        Returns:
+            List[Job]: List of timed out MCP jobs.
+        """
+        return await self.list(job_status=JobStatus.TIMEOUT)
+
+    async def list_pre_running(self) -> List[Job]:
+        """Get all pre-running MCP jobs.
+
+        Returns:
+            List[Job]: List of pre-running MCP jobs.
+        """
+        return await self.list(job_stage=JobStage.PRE_RUNNING)
+
+    async def list_running(self) -> List[Job]:
+        """Get all running MCP jobs.
+
+        Returns:
+            List[Job]: List of running MCP jobs.
+        """
+        return await self.list(job_stage=JobStage.RUNNING)
+
+    async def list_post_running(self) -> List[Job]:
+        """Get all post-running MCP jobs.
+
+        Returns:
+            List[Job]: List of post-running MCP jobs.
+        """
+        return await self.list(job_stage=JobStage.POST_RUNNING)
+
+
+class AsyncMcpClient:
+    def __init__(self, client: AsyncClient):
+        self._client = client
+        self.stored = AsyncStoredMcpClient(client)
+        self.config = AsyncMcpConfigClient(client)
+        self.jobs = AsyncMcpJobsClient(client)
+
+    async def run(
+        self,
+        code_zip: CodeZip,
+        mcp_module: str,
+        env_vars: EnvVars | dict | None = None,
+        transport: Literal["streamable-http", "http", "sse"] = "streamable-http",
+        data_zip: DataZip | None = None,
+    ) -> Job:
+        """Runs a MCP server from code zip.
+
+        Args:
+            code_zip (CodeZip): The code zip.
+            mcp_module (str): The name of the MCP server module.
+            env_vars (EnvVars | dict | None): Optional environment variables.
+            transport (Literal["streamable-http", "http", "sse"]): The transport protocol.
+            data_zip (DataZip | None): Optional data zip file.
+
+        Returns:
+            Job: The created job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/code/run/mcp"
+
+        files = self._client._prepare_files(code_zip=code_zip, data_zip=data_zip)
+        data: dict[str, Any] = {
+            "mcp_module": mcp_module,
+            "transport": transport,
+        }
+        if env_vars:
+            data["env_vars"] = EnvVars(env_vars).json_str
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    url,
+                    headers=self._client.api_key_header,
+                    files=files,
+                    data=data,
+                )
+                response.raise_for_status()
+                return Job(**response.json())
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def is_healthy(self, job_id: str) -> bool:
+        """Checks whether launched MCP server is healthy.
+
+        Args:
+            job_id (str): The ID of the MCP launch job.
+
+        Returns:
+            bool: True if MCP is healthy else False.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        return await self._client.jobs.is_healthy(job_id=job_id)
+
+    async def await_healthy(self, job_id: str, timeout: float | None = None) -> Job:
+        """Waits for a MCP server to become healthy.
+
+        Args:
+            job_id (str): The ID of the MCP run job.
+            timeout (float | None): Maximum time to wait in seconds. If None, waits indefinitely.
+
+        Returns:
+            Job: The job object of the MCP run job.
+
+        Raises:
+            HTTPException: If the request fails.
+            APIException: If the job enters stage POST_RUNNING unexpectedly.
+            TimeoutError: If the timeout is exceeded.
+        """
+        return await self._client.jobs.await_healthy(job_id=job_id, timeout=timeout)
+
+    def get_url(
+        self, job_id: str, transport: Literal["streamable-http", "http", "sse"] = "http"
+    ) -> str:
+        """Gets the URL for a MCP server.
+
+        Args:
+            job_id (str): The job ID.
+            transport (Literal["streamable-http", "http", "sse"]): The transport protocol.
+
+        Returns:
+            str: The MCP URL.
+        """
+        mcp_path = "sse" if transport == "sse" else "mcp"
+        return f"{self._client.get_proxy_url(job_id)}/{mcp_path}"
+
+    async def upload(
+        self,
+        code_zip: CodeZip,
+        code_name: str,
+        metadata: JsonData | dict | None = None,
+    ) -> str:
+        """Upload MCP code.
+
+        Args:
+            code_zip (CodeZip): The code zip.
+            code_name (str): The name of the code.
+            metadata (JsonData | dict | None): The JSON metadata of the code.
+
+        Returns:
+            str: The code ID.
+        """
+        return await self._client.code.upload(
+            code_zip=code_zip,
+            code_name=code_name,
+            code_type=CodeType.MCP,
+            metadata=metadata,
+        )
+
+    async def ingest(
+        self,
+        code_zip: CodeZip,
+        code_name: str,
+        metadata: JsonData | dict | None = None,
+        build_pexenv: bool = False,
+        pexenv_python: str | None = None,
+    ) -> Job:
+        """Ingest MCP code.
+
+        Args:
+            code_zip (CodeZip): The code zip.
+            code_name (str): The name of the code.
+            metadata (JsonData | dict | None): The JSON metadata of the code.
+            build_pexenv (bool): Whether to build the pex venv.
+            pexenv_python: (str | None): Python interpreter for the pex venv.
+
+        Returns:
+            Job: The code ingestion job.
+        """
+        return await self._client.code.ingest(
+            code_zip=code_zip,
+            code_name=code_name,
+            code_type=CodeType.MCP,
+            metadata=metadata,
+            build_pexenv=build_pexenv,
+            pexenv_python=pexenv_python,
+        )