@runpod/mcp-server 1.1.0 → 1.2.0

package/dist/index.mjs

@@ -5,25 +5,77 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { z } from "zod";
 import fetch from "node-fetch";
+import { randomUUID } from "crypto";
 var API_BASE_URL = "https://rest.runpod.io/v1";
 var API_KEY = process.env.RUNPOD_API_KEY;
 if (!API_KEY) {
   console.error("RUNPOD_API_KEY environment variable is required");
   process.exit(1);
 }
+var MCP_SERVER_NAME = "RunPod API Server";
+var MCP_SERVER_VERSION = true ? "1.2.0" : "dev";
+var MCP_TRANSPORT = "stdio";
+function setTransport(t) {
+  MCP_TRANSPORT = t;
+}
 var server = new McpServer({
-  name: "RunPod API Server",
-  version: "1.0.0",
+  name: MCP_SERVER_NAME,
+  version: MCP_SERVER_VERSION,
   capabilities: {
     resources: {},
     tools: {}
   }
 });
+var SESSION_ID = randomUUID();
+function resolveClientInfo() {
+  const info = server.server.getClientVersion();
+  if (info?.name) {
+    return { name: info.name, version: info.version || "unknown" };
+  }
+  return { name: "unknown", version: "unknown" };
+}
+function sanitizeUaToken(value) {
+  return value.replace(/[()<>@,;:\\"/\[\]?={}\s]/g, "_").slice(0, 64);
+}
+function buildUserAgent() {
+  const { name, version } = resolveClientInfo();
+  const safeName = sanitizeUaToken(name);
+  const safeVersion = sanitizeUaToken(version);
+  return `runpod-mcp-server/${MCP_SERVER_VERSION} (caller=mcp; client=${safeName}; client_version=${safeVersion}; transport=${MCP_TRANSPORT})`;
+}
+function trackingHeaders() {
+  return {
+    "User-Agent": buildUserAgent(),
+    "X-Runpod-Session-Id": SESSION_ID
+  };
+}
+var GRAPHQL_URL = "https://api.runpod.io/graphql";
+async function graphqlRequest(query) {
+  const response = await fetch(GRAPHQL_URL, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      ...trackingHeaders()
+    },
+    body: JSON.stringify({ query })
+  });
+  const result = await response.json();
+  if (result.errors && result.errors.length > 0) {
+    throw new Error(
+      `GraphQL Error: ${result.errors.map((e) => e.message).join(", ")}`
+    );
+  }
+  if (!result.data) {
+    throw new Error("No data returned from GraphQL query");
+  }
+  return result.data;
+}
 async function runpodRequest(endpoint, method = "GET", body) {
   const url = `${API_BASE_URL}${endpoint}`;
   const headers = {
     Authorization: `Bearer ${API_KEY}`,
-    "Content-Type": "application/json"
+    "Content-Type": "application/json",
+    ...trackingHeaders()
   };
   const options = {
     method,
@@ -48,6 +100,156 @@ async function runpodRequest(endpoint, method = "GET", body) {
     throw error;
   }
 }
+server.tool(
+  "list-gpu-types",
+  {
+    minMemoryGb: z.number().optional().describe("Filter to GPUs with at least this much VRAM in GB"),
+    secureCloudOnly: z.boolean().optional().describe("Filter to only GPUs available in secure cloud"),
+    communityCloudOnly: z.boolean().optional().describe("Filter to only GPUs available in community cloud"),
+    searchTerm: z.string().optional().describe(
+      "Search term to filter GPUs by name (e.g., 'A100', 'RTX 4090')"
+    ),
+    includeUnavailable: z.boolean().optional().describe(
+      "Include GPUs that are currently out of stock. Default is false"
+    )
+  },
+  async (params) => {
+    const data = await graphqlRequest(`
+      query {
+        gpuTypes {
+          id
+          displayName
+          memoryInGb
+          secureCloud
+          communityCloud
+          lowestPrice(input: { gpuCount: 1 }) {
+            stockStatus
+          }
+        }
+      }
+    `);
+    const stockPriority = {
+      High: 3,
+      Medium: 2,
+      Low: 1
+    };
+    const isAvailable = (gpu) => {
+      const status = gpu.lowestPrice?.stockStatus;
+      return !!status && status !== "Out";
+    };
+    let gpuTypes = data.gpuTypes.filter((gpu) => gpu.id !== "unknown");
+    if (!params.includeUnavailable) {
+      gpuTypes = gpuTypes.filter(isAvailable);
+    }
+    if (params.minMemoryGb) {
+      gpuTypes = gpuTypes.filter(
+        (gpu) => gpu.memoryInGb >= params.minMemoryGb
+      );
+    }
+    if (params.secureCloudOnly) {
+      gpuTypes = gpuTypes.filter((gpu) => gpu.secureCloud);
+    }
+    if (params.communityCloudOnly) {
+      gpuTypes = gpuTypes.filter((gpu) => gpu.communityCloud);
+    }
+    if (params.searchTerm) {
+      const term = params.searchTerm.toLowerCase();
+      gpuTypes = gpuTypes.filter(
+        (gpu) => gpu.id.toLowerCase().includes(term) || gpu.displayName.toLowerCase().includes(term)
+      );
+    }
+    gpuTypes.sort((a, b) => {
+      const aP = stockPriority[a.lowestPrice?.stockStatus || ""] || 0;
+      const bP = stockPriority[b.lowestPrice?.stockStatus || ""] || 0;
+      if (bP !== aP) return bP - aP;
+      return b.memoryInGb - a.memoryInGb;
+    });
+    const result = gpuTypes.map((gpu) => ({
+      id: gpu.id,
+      displayName: gpu.displayName,
+      memoryGb: gpu.memoryInGb,
+      secureCloud: gpu.secureCloud,
+      communityCloud: gpu.communityCloud,
+      stockStatus: gpu.lowestPrice?.stockStatus || "unavailable",
+      available: isAvailable(gpu)
+    }));
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "list-data-centers",
+  {
+    region: z.string().optional().describe(
+      "Filter by region/location (e.g., 'United States', 'Europe', 'Canada')"
+    )
+  },
+  async (params) => {
+    const data = await graphqlRequest(`
+      query {
+        dataCenters {
+          id
+          name
+          location
+        }
+      }
+    `);
+    let dataCenters = data.dataCenters;
+    if (params.region) {
+      const term = params.region.toLowerCase();
+      dataCenters = dataCenters.filter(
+        (dc) => dc.location.toLowerCase().includes(term)
+      );
+    }
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(dataCenters, null, 2)
+        }
+      ]
+    };
+  }
+);
+var SERVERLESS_API_BASE_URL = "https://api.runpod.ai/v2";
+async function serverlessRequest(endpointId, path, method = "GET", body) {
+  const url = `${SERVERLESS_API_BASE_URL}/${endpointId}${path}`;
+  const headers = {
+    Authorization: `Bearer ${API_KEY}`,
+    "Content-Type": "application/json",
+    ...trackingHeaders()
+  };
+  const options = {
+    method,
+    headers
+  };
+  if (body && (method === "POST" || method === "PATCH")) {
+    options.body = JSON.stringify(body);
+  }
+  try {
+    const response = await fetch(url, options);
+    if (!response.ok) {
+      const errorText = await response.text();
+      throw new Error(
+        `RunPod Serverless API Error: ${response.status} - ${errorText}`
+      );
+    }
+    const contentType = response.headers.get("content-type");
+    if (contentType && contentType.includes("application/json")) {
+      return await response.json();
+    }
+    return { success: true, status: response.status };
+  } catch (error) {
+    console.error("Error calling RunPod Serverless API:", error);
+    throw error;
+  }
+}
 server.tool(
   "list-pods",
   {
@@ -118,6 +320,7 @@ server.tool(
 );
 server.tool(
   "create-pod",
+  'Create a new GPU/CPU pod on RunPod. If the user does not specify an image, recommend the "Runpod Pytorch 2.8.0" image (runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404) as the default \u2014 it has the most up-to-date CUDA and PyTorch versions.',
   {
     name: z.string().optional().describe("Name for the pod"),
     imageName: z.string().describe("Docker image to use"),
@@ -342,17 +545,298 @@ server.tool(
     };
   }
 );
-server.tool("list-templates", {}, async () => {
-  const result = await runpodRequest("/templates");
-  return {
-    content: [
-      {
-        type: "text",
-        text: JSON.stringify(result, null, 2)
+var endpointIdSchema = z.string().regex(/^[a-zA-Z0-9_-]+$/, "Invalid endpoint ID format");
+var jobIdSchema = z.string().regex(/^[a-zA-Z0-9_-]+$/, "Invalid job ID format");
+var inputSchema = z.record(z.unknown()).describe(
+  "Input payload for the worker handler. The expected fields depend on the deployed model or worker."
+);
+var webhookSchema = z.string().url().optional().describe(
+  "Webhook URL to receive job completion notifications instead of polling"
+);
+var policySchema = z.object({
+  executionTimeout: z.number().optional().describe("Maximum execution time in milliseconds"),
+  lowPriority: z.boolean().optional().describe("Submit as a low-priority job"),
+  ttl: z.number().optional().describe("Time-to-live for the job result in milliseconds")
+}).optional().describe("Execution policy options");
+var s3ConfigSchema = z.object({
+  accessId: z.string().describe("S3 access key ID"),
+  accessSecret: z.string().describe("S3 secret access key"),
+  bucketName: z.string().describe("S3 bucket name"),
+  endpointUrl: z.string().describe("S3 endpoint URL")
+}).optional().describe("S3-compatible storage config for large outputs");
+server.tool(
+  "run-endpoint",
+  "Submit an asynchronous job to a Serverless endpoint. Returns a job ID immediately \u2014 use get-job-status to poll for results. Async results are available for 30 minutes after completion.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint to run"
+    ),
+    input: inputSchema,
+    webhook: webhookSchema,
+    policy: policySchema,
+    s3Config: s3ConfigSchema
+  },
+  async (params) => {
+    const { endpointId, ...body } = params;
+    const result = await serverlessRequest(
+      endpointId,
+      "/run",
+      "POST",
+      body
+    );
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "runsync-endpoint",
+  "Submit a synchronous job to a Serverless endpoint and wait for the result. Best for tasks completing within 90 seconds. If processing exceeds 90 seconds, the response returns a job ID to poll with get-job-status. Max payload: 20 MB. Results expire after 1 minute. Use the wait parameter to extend the server-side wait up to 5 minutes (300000 ms).",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint to run synchronously"
+    ),
+    input: inputSchema,
+    wait: z.number().min(1e3).max(3e5).optional().describe(
+      "How long in milliseconds the server should wait for a result before returning a job ID to poll (1000\u2013300000). Defaults to 90000 (90 seconds)."
+    ),
+    webhook: webhookSchema,
+    policy: policySchema,
+    s3Config: s3ConfigSchema
+  },
+  async (params) => {
+    const { endpointId, wait, ...body } = params;
+    const path = wait ? `/runsync?wait=${wait}` : "/runsync";
+    const result = await serverlessRequest(
+      endpointId,
+      path,
+      "POST",
+      body
+    );
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "get-job-status",
+  "Check the status of an asynchronous Serverless job. Returns the current status and output when complete. Job statuses: IN_QUEUE, IN_PROGRESS, COMPLETED, FAILED, CANCELLED, TIMED_OUT.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint the job belongs to"
+    ),
+    jobId: jobIdSchema.describe("ID of the job to check")
+  },
+  async (params) => {
+    const result = await serverlessRequest(
+      params.endpointId,
+      `/status/${params.jobId}`
+    );
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "stream-job",
+  "Retrieve all streaming output from a Serverless job by polling until the job reaches a terminal state. The worker must support streaming output. Polls /stream/{jobId} repeatedly and collects every chunk until status is COMPLETED, FAILED, CANCELLED, or TIMED_OUT.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint the job belongs to"
+    ),
+    jobId: jobIdSchema.describe("ID of the job to stream results from")
+  },
+  async (params) => {
+    const TERMINAL_STATUSES = /* @__PURE__ */ new Set([
+      "COMPLETED",
+      "FAILED",
+      "CANCELLED",
+      "TIMED_OUT"
+    ]);
+    const MAX_POLL_TIME_MS = 5 * 60 * 1e3;
+    const POLL_INTERVAL_MS = 1e3;
+    const MAX_CONSECUTIVE_ERRORS = 5;
+    const allChunks = [];
+    let finalResult = {};
+    let consecutiveErrors = 0;
+    const startTime = Date.now();
+    while (true) {
+      try {
+        const result = await serverlessRequest(
+          params.endpointId,
+          `/stream/${params.jobId}`
+        );
+        consecutiveErrors = 0;
+        if (Array.isArray(result.stream)) {
+          allChunks.push(...result.stream);
+        }
+        finalResult = result;
+        if (TERMINAL_STATUSES.has(result.status)) {
+          break;
+        }
+      } catch (error) {
+        consecutiveErrors++;
+        if (consecutiveErrors >= MAX_CONSECUTIVE_ERRORS) {
+          finalResult.error = `Polling aborted after ${MAX_CONSECUTIVE_ERRORS} consecutive errors: ${error instanceof Error ? error.message : String(error)}`;
+          break;
+        }
       }
-    ]
-  };
-});
+      if (Date.now() - startTime > MAX_POLL_TIME_MS) {
+        finalResult.pollingTimedOut = true;
+        finalResult.note = "Polling timed out after 5 minutes. Use get-job-status to check the job later.";
+        break;
+      }
+      await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
+    }
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify({ ...finalResult, stream: allChunks }, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "cancel-job",
+  "Cancel a Serverless job that is queued or in progress.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint the job belongs to"
+    ),
+    jobId: jobIdSchema.describe("ID of the job to cancel")
+  },
+  async (params) => {
+    const result = await serverlessRequest(
+      params.endpointId,
+      `/cancel/${params.jobId}`,
+      "POST"
+    );
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "retry-job",
+  "Retry a failed or timed-out Serverless job. Only works for jobs with FAILED or TIMED_OUT status. The previous output is removed and the job is requeued.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint the job belongs to"
+    ),
+    jobId: jobIdSchema.describe("ID of the job to retry")
+  },
+  async (params) => {
+    const result = await serverlessRequest(
+      params.endpointId,
+      `/retry/${params.jobId}`,
+      "POST"
+    );
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "endpoint-health",
+  "Get the health and operational status of a Serverless endpoint, including worker counts and job statistics.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint to check health for"
+    )
+  },
+  async (params) => {
+    const result = await serverlessRequest(params.endpointId, "/health");
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "purge-endpoint-queue",
+  "Remove all pending jobs from a Serverless endpoint queue. Only affects queued jobs \u2014 in-progress jobs continue running. Use this for error recovery or clearing outdated requests.",
+  {
+    endpointId: endpointIdSchema.describe(
+      "ID of the Serverless endpoint to purge the queue for"
+    )
+  },
+  async (params) => {
+    const result = await serverlessRequest(
+      params.endpointId,
+      "/purge-queue",
+      "POST"
+    );
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
+server.tool(
+  "list-templates",
+  `List available templates. By default returns only the user's own templates. Use includeRunpodTemplates to also include official RunPod templates. The recommended default template for new pods is "Runpod Pytorch 2.8.0" (ID: runpod-torch-v280) \u2014 it has the latest CUDA and PyTorch versions.`,
+  {
+    includeRunpodTemplates: z.boolean().optional().describe("Include official RunPod templates in the response"),
+    includePublicTemplates: z.boolean().optional().describe("Include community-made public templates in the response"),
+    includeEndpointBoundTemplates: z.boolean().optional().describe(
+      "Include templates bound to Serverless endpoints in the response"
+    )
+  },
+  async (params) => {
+    const queryParams = new URLSearchParams();
+    if (params.includeRunpodTemplates)
+      queryParams.set("includeRunpodTemplates", "true");
+    if (params.includePublicTemplates)
+      queryParams.set("includePublicTemplates", "true");
+    if (params.includeEndpointBoundTemplates)
+      queryParams.set("includeEndpointBoundTemplates", "true");
+    const query = queryParams.toString();
+    const result = await runpodRequest(`/templates${query ? `?${query}` : ""}`);
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2)
+        }
+      ]
+    };
+  }
+);
 server.tool(
   "get-template",
   {
@@ -612,4 +1096,7 @@ server.tool(
 );
 var transport = new StdioServerTransport();
 server.connect(transport);
+export {
+  setTransport
+};
 //# sourceMappingURL=index.mjs.map