Prefect MCP Server

from typing import Any, Dict, List, Optional, Union
from uuid import UUID

import mcp.types as types
from prefect import get_client
from prefect.states import Cancelled, Completed, Failed, Pending, Running, Scheduled

from .envs import PREFECT_API_URL
from .server import mcp


def get_flow_run_url(flow_run_id: str) -> str:
    base_url = PREFECT_API_URL.replace("/api", "")
    return f"{base_url}/flow-runs/{flow_run_id}"


@mcp.tool
async def get_flow_runs(
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    flow_name: Optional[str] = None,
    state_type: Optional[str] = None,
    state_name: Optional[str] = None,
    deployment_id: Optional[str] = None,
    tags: Optional[List[str]] = None,
    start_time_before: Optional[str] = None,
    start_time_after: Optional[str] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Get a list of flow runs with optional filtering.
    
    Args:
        limit: Maximum number of flow runs to return
        offset: Number of flow runs to skip
        flow_name: Filter by flow name
        state_type: Filter by state type (e.g., "RUNNING", "COMPLETED", "FAILED")
        state_name: Filter by state name
        deployment_id: Filter by deployment ID
        tags: Filter by tags
        start_time_before: ISO formatted datetime string
        start_time_after: ISO formatted datetime string
        
    Returns:
        A list of flow runs with their details
    """
    async with get_client() as client:
        # Build filter parameters
        filters = {}
        if flow_name:
            filters["flow_name"] = {"like_": f"%{flow_name}%"}
        if state_type:
            filters["state"] = {"type": {"any_": [state_type.upper()]}}
        if state_name:
            filters["state"] = {"name": {"any_": [state_name]}}
        if deployment_id:
            filters["deployment_id"] = {"eq_": UUID(deployment_id)}
        if tags:
            filters["tags"] = {"all_": tags}
        if start_time_after:
            filters["start_time"] = {"ge_": start_time_after}
        if start_time_before:
            if "start_time" in filters:
                filters["start_time"]["le_"] = start_time_before
            else:
                filters["start_time"] = {"le_": start_time_before}
        
        flow_runs = await client.read_flow_runs(
            limit=limit,
            offset=offset,
            **filters
        )
        
        # Add UI links to each flow run
        flow_runs_result = {
            "flow_runs": [
                {
                    **flow_run.dict(),
                    "ui_url": get_flow_run_url(str(flow_run.id))
                }
                for flow_run in flow_runs
            ]
        }
        
        return [types.TextContent(type="text", text=str(flow_runs_result))]


@mcp.tool
async def get_flow_run(
    flow_run_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Get details of a specific flow run by ID.
    
    Args:
        flow_run_id: The flow run UUID
        
    Returns:
        Flow run details
    """
    async with get_client() as client:
        flow_run = await client.read_flow_run(UUID(flow_run_id))
        
        # Add UI link
        flow_run_dict = flow_run.dict()
        flow_run_dict["ui_url"] = get_flow_run_url(flow_run_id)
        
        return [types.TextContent(type="text", text=str(flow_run_dict))]


@mcp.tool
async def get_flow_runs_by_flow(
    flow_id: str,
    limit: Optional[int] = None,
    offset: Optional[int] = None,
    state_type: Optional[str] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Get flow runs for a specific flow.
    
    Args:
        flow_id: The flow UUID
        limit: Maximum number of flow runs to return
        offset: Number of flow runs to skip
        state_type: Filter by state type (e.g., "RUNNING", "COMPLETED", "FAILED")
        
    Returns:
        A list of flow runs for the specified flow
    """
    async with get_client() as client:
        # Build filter parameters
        filters = {"flow_id": {"eq_": UUID(flow_id)}}
        if state_type:
            filters["state"] = {"type": {"any_": [state_type.upper()]}}
        
        flow_runs = await client.read_flow_runs(
            limit=limit,
            offset=offset,
            **filters
        )
        
        # Add UI links to each flow run
        flow_runs_result = {
            "flow_runs": [
                {
                    **flow_run.dict(),
                    "ui_url": get_flow_run_url(str(flow_run.id))
                }
                for flow_run in flow_runs
            ]
        }
        
        return [types.TextContent(type="text", text=str(flow_runs_result))]


@mcp.tool
async def restart_flow_run(
    flow_run_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Restart a flow run.
    
    Args:
        flow_run_id: The flow run UUID
        
    Returns:
        Details of the new flow run
    """
    async with get_client() as client:
        flow_run_id_uuid = UUID(flow_run_id)
        new_flow_run = await client.create_flow_run_from_flow_run(flow_run_id_uuid)
        
        new_flow_run_dict = new_flow_run.dict()
        new_flow_run_dict["ui_url"] = get_flow_run_url(str(new_flow_run.id))
        
        return [types.TextContent(type="text", text=str(new_flow_run_dict))]


@mcp.tool
async def cancel_flow_run(
    flow_run_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Cancel a flow run.
    
    Args:
        flow_run_id: The flow run UUID
        
    Returns:
        Confirmation message
    """
    async with get_client() as client:
        await client.set_flow_run_state(
            flow_run_id=UUID(flow_run_id),
            state=Cancelled(message="Cancelled via MCP")
        )
        
        return [types.TextContent(type="text", text=f"Flow run '{flow_run_id}' cancelled successfully.")]


@mcp.tool
async def delete_flow_run(
    flow_run_id: str,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Delete a flow run.
    
    Args:
        flow_run_id: The flow run UUID
        
    Returns:
        Confirmation message
    """
    async with get_client() as client:
        await client.delete_flow_run(UUID(flow_run_id))
        
        return [types.TextContent(type="text", text=f"Flow run '{flow_run_id}' deleted successfully.")]


@mcp.tool
async def set_flow_run_state(
    flow_run_id: str,
    state: str,
    message: Optional[str] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Set a flow run's state.
    
    Args:
        flow_run_id: The flow run UUID
        state: The new state to set (e.g., "SCHEDULED", "RUNNING", "COMPLETED", "FAILED")
        message: An optional message explaining the state change
        
    Returns:
        Result of the state change operation
    """
    async with get_client() as client:
        state_obj = None
        if state.upper() == "SCHEDULED":
            state_obj = Scheduled(message=message)
        elif state.upper() == "RUNNING":
            state_obj = Running(message=message)
        elif state.upper() == "COMPLETED":
            state_obj = Completed(message=message)
        elif state.upper() == "FAILED":
            state_obj = Failed(message=message)
        elif state.upper() == "PENDING":
            state_obj = Pending(message=message)
        elif state.upper() == "CANCELLED":
            state_obj = Cancelled(message=message)
        else:
            return [types.TextContent(
                type="text", 
                text=f"Invalid state '{state}'. Must be one of: SCHEDULED, RUNNING, COMPLETED, FAILED, PENDING, CANCELLED"
            )]
        
        result = await client.set_flow_run_state(
            flow_run_id=UUID(flow_run_id),
            state=state_obj
        )
        
        return [types.TextContent(type="text", text=str(result.dict()))]