MCP Server Replicate

Instructions

Implementation Reference

• src/mcp_server_replicate/server.py:889-894 (handler)

  The MCP tool handler for 'cancel_prediction' that executes the cancellation logic using the ReplicateClient.

  async def cancel_prediction(prediction_id: str) -> dict[str, Any]:
      """Cancel a running prediction."""
      async with ReplicateClient(api_token=os.getenv("REPLICATE_API_TOKEN")) as client:
          response = await client.cancel_prediction(prediction_id)
          return await response.json()

• src/mcp_server_replicate/replicate_client.py:658-692 (helper)

  Supporting method in ReplicateClient that performs the HTTP POST request to the Replicate API endpoint for canceling a prediction.

  async def cancel_prediction(self, prediction_id: str) -> dict[str, Any]:
      """Cancel a running prediction.
      
      Args:
          prediction_id: The ID of the prediction to cancel
          
      Returns:
          Dict containing the updated prediction status
          
      Raises:
          ValueError: If the prediction is not found
          Exception: If the cancellation fails
      """
      if not self.client:
          raise RuntimeError("Client not initialized. Check error property for details.")
  
      try:
          response = await self.http_client.post(
              f"/predictions/{prediction_id}/cancel",
              headers={
                  "Authorization": f"Bearer {self.api_token}",
                  "Content-Type": "application/json",
              }
          )
          response.raise_for_status()
          return response.json()
          
      except httpx.HTTPStatusError as err:
          if err.response.status_code == 404:
              raise ValueError(f"Prediction not found: {prediction_id}")
          logger.error(f"Failed to cancel prediction: {str(err)}")
          raise Exception(f"Failed to cancel prediction: {str(err)}") from err
      except Exception as err:
          logger.error(f"Failed to cancel prediction: {str(err)}")
          raise Exception(f"Failed to cancel prediction: {str(err)}") from err

• src/mcp_server_replicate/models/prediction.py:28-42 (schema)

  Pydantic model defining the structure of prediction responses, used as output type/schema for the tool.

  class Prediction(BaseModel):
      """A prediction (model run) on Replicate."""
      id: str = Field(..., description="Unique identifier for this prediction")
      version: str = Field(..., description="Model version used for this prediction")
      status: PredictionStatus = Field(..., description="Current status of the prediction")
      input: Dict[str, Any] = Field(..., description="Input parameters used for the prediction")
      output: Optional[Any] = Field(None, description="Output from the prediction if completed")
      error: Optional[str] = Field(None, description="Error message if prediction failed")
      logs: Optional[str] = Field(None, description="Execution logs from the prediction")
      created_at: datetime
      started_at: Optional[datetime] = None
      completed_at: Optional[datetime] = None
      urls: Dict[str, str] = Field(..., description="Related API URLs for this prediction")
      metrics: Optional[Dict[str, float]] = Field(None, description="Performance metrics if available")
      stream_url: Optional[str] = Field(None, description="URL for streaming output if requested")