@runpod/mcp-server 3.2.0 → 3.4.0

package/dist/{chunk-WGK7D2YR.mjs → chunk-G4URDODF.mjs}

@@ -151,6 +151,28 @@ var HttpError = class extends Error {
     this.body = body;
   }
 };
+var DEFAULT_REQUEST_TIMEOUT_MS = 3e4;
+var MIN_REQUEST_TIMEOUT_MS = 1e3;
+var RequestTimeoutError = class extends Error {
+  constructor(prefix, timeoutMs, method) {
+    const write = method !== "GET" && method !== "HEAD";
+    super(
+      `${prefix}: no response after ${timeoutMs}ms for ${method} (the API may be overloaded, or still be working on it). ${write ? "This request may have SUCCEEDED upstream \u2014 do not retry blindly. Check with the matching list-/get- tool first and retry only if it did not take effect." : "This was a read and changed nothing, so retrying is safe."}`
+    );
+    this.name = "RequestTimeoutError";
+    this.timeoutMs = timeoutMs;
+    this.method = method;
+  }
+};
+async function withRequestTimeout(errorPrefix, timeoutMs, method, run) {
+  const signal = AbortSignal.timeout(timeoutMs);
+  try {
+    return await run(signal);
+  } catch (error) {
+    if (error instanceof HttpError || !signal.aborted) throw error;
+    throw new RequestTimeoutError(errorPrefix, timeoutMs, method);
+  }
+}
 function isJsonContentType(contentType) {
   if (!contentType) return false;
   const ct = contentType.toLowerCase();
@@ -166,28 +188,49 @@ function buildRequestHeaders(apiKey, tracking) {
     ...tracking()
   };
 }
+function resolveCeiling(ceiling) {
+  if (ceiling === void 0) return Infinity;
+  if (typeof ceiling !== "function") return ceiling;
+  return Math.max(ceiling(), MIN_REQUEST_TIMEOUT_MS);
+}
+function clampTimeout(requestedMs, ceiling, defaultMs = DEFAULT_REQUEST_TIMEOUT_MS) {
+  return Math.min(requestedMs ?? defaultMs, resolveCeiling(ceiling));
+}
 function createHttpClient(deps) {
-  return async function request(url, method = "GET", body) {
-    const init = {
+  const defaultTimeoutMs = deps.defaultTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS;
+  return async function request(url, method = "GET", body, options) {
+    const timeoutMs = clampTimeout(
+      options?.timeoutMs,
+      deps.maxTimeoutMs,
+      defaultTimeoutMs
+    );
+    return withRequestTimeout(
+      deps.errorPrefix,
+      timeoutMs,
       method,
-      headers: buildRequestHeaders(deps.apiKey, deps.tracking)
-    };
-    if (body && methodSendsBody(method)) {
-      init.body = JSON.stringify(body);
-    }
-    const response = await deps.fetch(url, init);
-    if (!response.ok) {
-      throw new HttpError(
-        deps.errorPrefix,
-        response.status,
-        await response.text(),
-        response.status === 429 ? rateLimitHint(
-          response.headers.get("ratelimit"),
-          response.headers.get("retry-after")
-        ) : response.status === 401 ? EXPIRED_CREDENTIAL_HINT : void 0
-      );
-    }
-    return isJsonContentType(response.headers.get("content-type")) ? response.json() : { success: true, status: response.status };
+      async (signal) => {
+        const init = {
+          method,
+          headers: buildRequestHeaders(deps.apiKey, deps.tracking),
+          signal
+        };
+        if (body && methodSendsBody(method)) {
+          init.body = JSON.stringify(body);
+        }
+        const response = await deps.fetch(url, init);
+        if (!response.ok) {
+          const hint = response.status === 429 ? rateLimitHint(
+            response.headers.get("ratelimit"),
+            response.headers.get("retry-after")
+          ) : response.status === 401 ? EXPIRED_CREDENTIAL_HINT : void 0;
+          const text = await response.text().catch(
+            () => "<body unavailable: the connection closed while reading it>"
+          );
+          throw new HttpError(deps.errorPrefix, response.status, text, hint);
+        }
+        return isJsonContentType(response.headers.get("content-type")) ? response.json() : { success: true, status: response.status };
+      }
+    );
   };
 }
 
@@ -337,10 +380,15 @@ function containerConfigToV2(p) {
     mounts: persistentMount(p.volumeInGb, p.volumeMountPath)
   });
 }
-function gpuConfigToV2(gpuTypeIds, gpuCount) {
+function gpuConfigToV2(gpuTypeIds, gpuCount, allowedCudaVersions, minCudaVersion2) {
   const id = gpuTypeIds?.[0];
   if (id === void 0) return void 0;
-  return compact({ id, count: gpuCount });
+  return compact({
+    id,
+    count: gpuCount ?? 1,
+    allowedCudaVersions,
+    minCudaVersion: minCudaVersion2
+  });
 }
 function mapPodCreateToV2(params) {
   return compact({
@@ -361,7 +409,12 @@ function mapPodCreateToV2(params) {
     // singular on v2 (`gpu.id`): gpuConfigToV2 takes the first id; the
     // create-pod handler warns when more than one gpuTypeId was supplied.
     dataCenterIds: params.dataCenterIds?.length ? params.dataCenterIds : void 0,
-    gpu: gpuConfigToV2(params.gpuTypeIds, params.gpuCount)
+    gpu: gpuConfigToV2(
+      params.gpuTypeIds,
+      params.gpuCount,
+      params.allowedCudaVersions,
+      params.minCudaVersion
+    )
   });
 }
 function mapPodUpdateToV2(params) {
@@ -378,9 +431,15 @@ function defaultScalerType(endpointType) {
 function endpointScaling(scalerType, value) {
   return scalerType === "REQUEST_COUNT" ? { type: "REQUEST_COUNT", requestCount: value } : { type: "QUEUE_DELAY", queueDelay: value };
 }
-function endpointGpuConfig(pools, count) {
-  if (!pools?.length) return void 0;
-  return compact({ pools, count });
+function endpointGpuConfig(params, mode) {
+  const gpu = compact({
+    pools: params.gpuPoolIds?.length ? params.gpuPoolIds : void 0,
+    count: params.gpuCount,
+    allowedCudaVersions: params.allowedCudaVersions,
+    minCudaVersion: params.minCudaVersion
+  });
+  if (mode === "create" && !("pools" in gpu)) return void 0;
+  return Object.keys(gpu).length ? gpu : void 0;
 }
 function endpointWorkers(params) {
   const workers = compact({
@@ -390,7 +449,7 @@ function endpointWorkers(params) {
   });
   return Object.keys(workers).length ? workers : void 0;
 }
-function endpointCommonToV2(params) {
+function endpointCommonToV2(params, mode) {
   return compact({
     name: params.name,
     image: params.imageName,
@@ -399,7 +458,7 @@ function endpointCommonToV2(params) {
     ports: params.ports,
     env: params.env,
     registry: params.containerRegistryAuthId,
-    gpu: endpointGpuConfig(params.gpuPoolIds, params.gpuCount),
+    gpu: endpointGpuConfig(params, mode),
     workers: endpointWorkers(params),
     dataCenterIds: params.dataCenterIds?.length ? params.dataCenterIds : void 0,
     networkVolumes: params.networkVolumeIds?.length ? params.networkVolumeIds : void 0,
@@ -410,7 +469,7 @@ function endpointCommonToV2(params) {
 function mapEndpointCreateToV2(params) {
   const endpointType = params.endpointType ?? DEFAULT_ENDPOINT_TYPE;
   return compact({
-    ...endpointCommonToV2(params),
+    ...endpointCommonToV2(params, "create"),
     type: endpointType,
     scaling: endpointScaling(
       params.scalerType ?? defaultScalerType(endpointType),
@@ -420,13 +479,28 @@ function mapEndpointCreateToV2(params) {
 }
 function mapEndpointUpdateToV2(params) {
   return compact({
-    ...endpointCommonToV2(params),
+    ...endpointCommonToV2(params, "update"),
     scaling: params.scalerType ? endpointScaling(
       params.scalerType,
       params.scalerValue ?? DEFAULT_SCALER_VALUE
     ) : void 0
   });
 }
+var CUDA_MAJOR_MINOR = /^\d+\.\d+$/;
+function cudaConstraintError(params, opts = {}) {
+  const { allowedCudaVersions, minCudaVersion: minCudaVersion2 } = params;
+  if (allowedCudaVersions?.length && minCudaVersion2) {
+    return "allowedCudaVersions and minCudaVersion are mutually exclusive \u2014 pass either an exact set or an open-ended floor, not both.";
+  }
+  if (minCudaVersion2 !== void 0 && !CUDA_MAJOR_MINOR.test(minCudaVersion2) && !(opts.allowClear && minCudaVersion2 === "")) {
+    return `minCudaVersion must be major.minor (e.g. "12.0", not "12")${opts.allowClear ? "; an empty string clears the floor" : ""}.`;
+  }
+  const bad = allowedCudaVersions?.find((v) => !CUDA_MAJOR_MINOR.test(v));
+  if (bad !== void 0) {
+    return `allowedCudaVersions entries must be major.minor (e.g. "12.8"); got "${bad}". Discover valid values via get-capacity or list-gpu-types.`;
+  }
+  return void 0;
+}
 function mapNetworkVolumeCreateToV2(params) {
   return compact({
     name: params.name,
@@ -467,6 +541,42 @@ function podBodyFromTemplate(template) {
   }
   return out;
 }
+var SSH_PUBLIC_KEY_LINE = /^(ssh-[a-z0-9-]+|ecdsa-[a-z0-9-]+|sk-[a-z0-9@.-]+)\s+\S+/i;
+var PRIVATE_KEY_MARKER = /private[-\s]?key|PuTTY-User-Key-File/i;
+function normalizeSshPublicKey(input) {
+  if (PRIVATE_KEY_MARKER.test(input))
+    return {
+      ok: false,
+      error: 'sshPublicKey looks like a PRIVATE key. Pass the PUBLIC key instead \u2014 the contents of your ~/.ssh/<key>.pub file (e.g. "ssh-ed25519 AAAA\u2026").'
+    };
+  const lines = input.split(/\r?\n/).map((line) => line.trim()).filter((line) => line.length > 0);
+  if (lines.length === 0)
+    return {
+      ok: false,
+      error: "sshPublicKey is empty. Pass the contents of your ~/.ssh/<key>.pub file, or omit the parameter entirely to create the pod without SSH."
+    };
+  if (lines.some((line) => !SSH_PUBLIC_KEY_LINE.test(line)))
+    return {
+      ok: false,
+      error: 'sshPublicKey is not an SSH public key. Expected the CONTENTS of your ~/.ssh/<key>.pub file \u2014 a line like "ssh-ed25519 AAAAC3Nza\u2026 you@host" (ssh-*, ecdsa-*, or sk-* followed by the base64 key) \u2014 not a file path, a SHA256 fingerprint, or a bare base64 blob. Newline-separate multiple keys.'
+    };
+  return { ok: true, key: lines.join("\n") };
+}
+function applySshPublicKey(body, sshPublicKey) {
+  const ports = Array.isArray(body.ports) ? body.ports : [];
+  const hasSshPort = ports.some((p) => /^\s*22\/tcp\s*$/i.test(String(p)));
+  const env = { ...body.env ?? {} };
+  const key = sshPublicKey.trim();
+  if (!env.PUBLIC_KEY) env.PUBLIC_KEY = key;
+  else if (!env.PUBLIC_KEY.includes(key))
+    env.PUBLIC_KEY = `${env.PUBLIC_KEY}
+${key}`;
+  return {
+    ...body,
+    ports: hasSshPort ? ports : [...ports, "22/tcp"],
+    env
+  };
+}
 function mapTemplateUpdateToV2(params) {
   return {
     ...containerConfigToV2(params),
@@ -681,7 +791,7 @@ function resolveBackend(opts) {
 }
 
 // src/tools/runtime.ts
-var MCP_SERVER_VERSION = true ? "3.2.0" : "dev";
+var MCP_SERVER_VERSION = true ? "3.4.0" : "dev";
 var SESSION_ID = randomUUID();
 function trackingHeaders(server, ctx) {
   const info = server.server.getClientVersion();
@@ -696,6 +806,11 @@ function trackingHeaders(server, ctx) {
     sessionId: SESSION_ID
   });
 }
+var HTTP_TRANSPORT_BUDGET_MS = 5e4;
+function invocationBudget(totalMs) {
+  const startedAt = Date.now();
+  return () => totalMs - (Date.now() - startedAt);
+}
 var defaultFetch = fetch;
 async function withApiErrorLog(label, call) {
   try {
@@ -719,6 +834,23 @@ var DESTRUCTIVE = {
   openWorldHint: true
 };
 async function graphqlRequest(query, tracking, fetchImpl, url, options) {
+  return withRequestTimeout(
+    "Runpod GraphQL Error",
+    clampTimeout(options?.timeoutMs, options?.maxTimeoutMs),
+    // GraphQL is POST on the wire whatever it carries, but the retry advice keys
+    // off the method, and telling a caller its list-gpu-types "may have
+    // SUCCEEDED upstream — check with the matching list-/get- tool first" points
+    // it back at the tool that just failed. The operation keyword is what
+    // actually says whether anything could have been written.
+    isGraphqlMutation(query) ? "POST" : "GET",
+    (signal) => runGraphql(query, tracking, fetchImpl, url, signal, options)
+  );
+}
+function isGraphqlMutation(query) {
+  const firstMeaningful = query.split("\n").map((line) => line.trim()).find((line) => line.length > 0 && !line.startsWith("#"));
+  return /^mutation\b/.test(firstMeaningful ?? "");
+}
+async function runGraphql(query, tracking, fetchImpl, url, signal, options) {
   const response = await fetchImpl(url, {
     method: "POST",
     headers: {
@@ -729,7 +861,8 @@ async function graphqlRequest(query, tracking, fetchImpl, url, options) {
     body: JSON.stringify({
       query,
       ...options?.variables ? { variables: options.variables } : {}
-    })
+    }),
+    signal
   });
   if (!response.ok) {
     const bodyText = await response.text();
@@ -785,39 +918,53 @@ function createToolRuntime(server, ctx, deps = {}) {
   }
   const rawFetch = deps.fetch ?? defaultFetch;
   const httpFetch = observeUnauthorized(rawFetch);
+  const remainingHttpBudgetMs = invocationBudget(HTTP_TRANSPORT_BUDGET_MS);
+  const timeouts = {
+    defaultTimeoutMs: deps.defaultTimeoutMs,
+    maxTimeoutMs: deps.maxTimeoutMs ?? (ctx.transport === "http" ? remainingHttpBudgetMs : void 0)
+  };
   const v1Client = createHttpClient({
     apiKey: ctx.apiKey,
     fetch: httpFetch,
     tracking,
-    errorPrefix: "Runpod API Error"
+    errorPrefix: "Runpod API Error",
+    ...timeouts
   });
-  const runpodRequest = (endpoint, method = "GET", body) => withApiErrorLog(
+  const runpodRequest = (endpoint, method = "GET", body, options) => withApiErrorLog(
     "Error calling Runpod API:",
-    () => v1Client(`${restV1Base(process.env)}${endpoint}`, method, body)
+    () => v1Client(
+      `${restV1Base(process.env)}${endpoint}`,
+      method,
+      body,
+      options
+    )
   );
   const serverlessClient = createHttpClient({
     apiKey: ctx.apiKey,
     fetch: httpFetch,
     tracking,
-    errorPrefix: "Runpod Serverless API Error"
+    errorPrefix: "Runpod Serverless API Error",
+    ...timeouts
   });
-  const serverlessRequest = (endpointId, path, method = "GET", body) => withApiErrorLog(
+  const serverlessRequest = (endpointId, path, method = "GET", body, options) => withApiErrorLog(
     "Error calling Runpod Serverless API:",
     () => serverlessClient(
       `${serverlessBase(process.env)}/${endpointId}${path}`,
       method,
-      body
+      body,
+      options
     )
   );
   const restClient = createHttpClient({
     apiKey: ctx.apiKey,
     fetch: httpFetch,
     tracking,
-    errorPrefix: "Runpod API Error"
+    errorPrefix: "Runpod API Error",
+    ...timeouts
   });
-  const callRestUrl = (url, method = "GET", body) => withApiErrorLog(
+  const callRestUrl = (url, method = "GET", body, options) => withApiErrorLog(
     "Error calling Runpod API:",
-    () => restClient(url, method, body)
+    () => restClient(url, method, body, options)
   );
   const sseFetch = observeUnauthorized(
     deps.sseFetch ?? fetch
@@ -850,12 +997,18 @@ function createToolRuntime(server, ctx, deps = {}) {
   };
   return {
     jsonReply,
+    // No GraphQL caller overrides a deadline, but both are still capped by the
+    // shared budget: deploy-hub-repo spends one of each back to back.
     graphql: (query) => graphqlRequest(
       query,
       tracking,
       // rawFetch, not httpFetch: this call carries no credential (see above).
       rawFetch,
-      publicGraphqlBase(process.env)
+      publicGraphqlBase(process.env),
+      {
+        timeoutMs: deps.defaultTimeoutMs,
+        maxTimeoutMs: timeouts.maxTimeoutMs
+      }
     ),
     graphqlAuthed: (query, variables) => graphqlRequest(
       query,
@@ -863,7 +1016,12 @@ function createToolRuntime(server, ctx, deps = {}) {
       httpFetch,
       // NOT publicGraphqlBase: this call carries the caller's API key.
       authedGraphqlBase(process.env),
-      { variables, apiKey: ctx.apiKey }
+      {
+        variables,
+        apiKey: ctx.apiKey,
+        timeoutMs: deps.defaultTimeoutMs,
+        maxTimeoutMs: timeouts.maxTimeoutMs
+      }
     ),
     runpodRequest,
     serverlessRequest,
@@ -871,7 +1029,8 @@ function createToolRuntime(server, ctx, deps = {}) {
     backendFor,
     streamSse,
     podAction,
-    env: process.env
+    env: process.env,
+    transport: ctx.transport
   };
 }
 
@@ -940,6 +1099,7 @@ function capListResult(result, options, extra) {
 }
 
 // src/tools/catalog.ts
+var GPU_PRODUCTS = /* @__PURE__ */ new Set(["POD", "CLUSTER", "SERVERLESS"]);
 function registerCatalogTools(server, rt) {
   const { jsonReply: jsonReply2, graphql, callRestUrl, backendFor } = rt;
   server.tool(
@@ -958,15 +1118,50 @@ function registerCatalogTools(server, rt) {
       ),
       includeAvailability: z3.boolean().optional().describe(
         "Request realtime stock and annotate each GPU with an availability summary (HIGH/MEDIUM/LOW/NONE). Default true. Set false to skip the availability lookup \u2014 then out-of-stock GPUs cannot be filtered."
+      ),
+      product: z3.enum(["POD", "CLUSTER", "SERVERLESS"]).optional().describe(
+        "Product context for the availability lookup \u2014 the same GPU can be scarce for Pods and plentiful for Serverless. Default POD. Use SERVERLESS when picking a GPU for an endpoint, CLUSTER for Instant Clusters. Ignored when includeAvailability is false."
+      ),
+      minCudaVersion: z3.string().optional().describe(
+        'Scope the availability lookup to hosts at or above this CUDA version (v2 only). Integer major or major.minor \u2014 "12" means any 12.x release; unlike the create-pod/create-endpoint body fields, a bare major is accepted here because the filter only widens a read. Requires includeAvailability (the default).'
       )
     },
     { title: "List GPU types", ...READ_ONLY },
     async (params) => {
       const backend = backendFor("gpus");
+      if (params.minCudaVersion !== void 0 && backend.version !== "v2") {
+        return jsonReply2({
+          error: "The minCudaVersion filter is only supported on the v2 REST API. Set RUNPOD_REST_VERSION=v2, or use get-capacity for per-CUDA-version availability.",
+          status: 501
+        });
+      }
       if (backend.version === "v2") {
         const wantAvailability = params.includeAvailability !== false;
+        const product = GPU_PRODUCTS.has(String(params.product)) ? String(params.product) : "POD";
+        if (params.minCudaVersion !== void 0) {
+          if (!wantAvailability) {
+            return jsonReply2({
+              error: "minCudaVersion filters the availability lookup, so it requires includeAvailability (the default). Drop includeAvailability:false or the filter.",
+              status: 400
+            });
+          }
+          if (!/^\d+(\.\d+)?$/.test(params.minCudaVersion)) {
+            return jsonReply2({
+              error: 'minCudaVersion must be an integer major ("12") or major.minor ("12.1").',
+              status: 400
+            });
+          }
+        }
+        const query = new URLSearchParams();
+        if (wantAvailability) {
+          query.set("include", "AVAILABILITY");
+          query.set("product", product);
+          if (params.minCudaVersion !== void 0) {
+            query.set("minCudaVersion", params.minCudaVersion);
+          }
+        }
         const raw = await callRestUrl(
-          `${backend.base}${backend.list}${wantAvailability ? "?include=AVAILABILITY" : ""}`
+          `${backend.base}${backend.list}${query.size ? `?${query}` : ""}`
         );
         let gpus = backend.unwrap(raw);
         gpus = gpus.filter((g) => g.id !== "unknown");
@@ -1139,6 +1334,9 @@ function registerCatalogTools(server, rt) {
       gpuTypeId: z3.string().describe("ID of the GPU type to retrieve"),
       includeAvailability: z3.boolean().optional().describe(
         "Include realtime per-datacenter availability (HIGH/MEDIUM/LOW/NONE). Default true."
+      ),
+      product: z3.enum(["POD", "CLUSTER", "SERVERLESS"]).optional().describe(
+        "Product context for the availability lookup \u2014 stock differs by product. Default POD. Use SERVERLESS when picking a GPU for an endpoint, CLUSTER for Instant Clusters. Ignored when includeAvailability is false."
       )
     },
     { title: "Get GPU type", ...READ_ONLY },
@@ -1151,7 +1349,8 @@ function registerCatalogTools(server, rt) {
         });
       }
       const path = backend.get(encodeURIComponent(params.gpuTypeId));
-      const query = params.includeAvailability === false ? "" : "?include=AVAILABILITY";
+      const product = GPU_PRODUCTS.has(String(params.product)) ? String(params.product) : "POD";
+      const query = params.includeAvailability === false ? "" : `?include=AVAILABILITY&product=${product}`;
       const result = await callRestUrl(`${backend.base}${path}${query}`);
       return jsonReply2(result);
     }
@@ -2031,7 +2230,7 @@ function registerPodTools(server, rt) {
   );
   server.tool(
     "create-pod",
-    'Create a new GPU/CPU pod on Runpod. Pass gpuTypeIds for a GPU pod, or computeType:"CPU" for a CPU pod (CPU pods are served by the v1 API for now). Pass either imageName or templateId (templateId deploys from an existing template \u2014 its name, image, start command, ports, env, disk, volume, and registry credential are used as defaults, and any field you also pass explicitly replaces the template value for that field, e.g. passing env replaces the whole env set rather than merging; pass containerRegistryAuthId to override or supply a private-image credential, or an empty string to deploy with none). If the user specifies neither an image nor a template, 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.',
+    'Create a new GPU/CPU pod on Runpod. Pass gpuTypeIds for a GPU pod, or computeType:"CPU" for a CPU pod (CPU pods are served by the v1 API for now). Pass either imageName or templateId (templateId deploys from an existing template \u2014 its name, image, start command, ports, env, disk, volume, and registry credential are used as defaults, and any field you also pass explicitly replaces the template value for that field, e.g. passing env replaces the whole env set rather than merging; pass containerRegistryAuthId to override or supply a private-image credential, or an empty string to deploy with none). For full SSH access (including SCP/SFTP/rsync), pass sshPublicKey \u2014 it sets the PUBLIC_KEY env var and exposes 22/tcp, which official Runpod images turn into a running sshd with the key authorized. If the user specifies neither an image nor a template, 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: z7.string().optional().describe(
         "Name for the pod. Required on v2 unless deploying from a templateId (the template supplies its name as the default)."
@@ -2045,20 +2244,48 @@ function registerPodTools(server, rt) {
       ),
       cloudType: z7.enum(["SECURE", "COMMUNITY"]).optional().describe("SECURE or COMMUNITY cloud"),
       gpuTypeIds: z7.array(z7.string()).optional().describe("List of acceptable GPU types"),
-      gpuCount: z7.number().optional().describe("Number of GPUs"),
+      // v2 GpuConfig.count is `integer, minimum: 1`. Without these bounds a 0,
+      // a negative, or a fractional count reached the wire as-is and came back
+      // as a raw 422; reject it here with a message that names the field.
+      gpuCount: z7.number().int().min(1).optional().describe("Number of GPUs (minimum 1)"),
       containerDiskInGb: z7.number().optional().describe("Container disk size in GB"),
       volumeInGb: z7.number().optional().describe("Volume size in GB"),
       volumeMountPath: z7.string().optional().describe("Path to mount the volume"),
       ports: z7.array(z7.string()).optional().describe("Ports to expose (e.g., '8888/http', '22/tcp')"),
       env: z7.record(z7.string()).optional().describe("Environment variables"),
+      sshPublicKey: z7.string().optional().describe(
+        'SSH PUBLIC key(s) to authorize on the pod \u2014 the CONTENTS of your ~/.ssh/<key>.pub file, i.e. a literal "ssh-ed25519 AAAAC3Nza\u2026 you@host" line, NOT a path to the file, a fingerprint, or the private key; newline-separate multiple keys. Merges the key into the PUBLIC_KEY env var and ensures 22/tcp is exposed, enabling full SSH (SCP/SFTP/rsync) on any image that honors the PUBLIC_KEY convention \u2014 all runpod/* official images do. Extends template ports/env rather than replacing them.'
+      ),
       containerRegistryAuthId: z7.string().optional().describe(
         "Container registry credential ID for pulling a private image (see list-container-registry-auths). When deploying from a templateId, this overrides the credential the template carries; pass an empty string to deploy with NO credential (clears a template\u2019s inherited registry)."
       ),
-      dataCenterIds: z7.array(z7.string()).optional().describe("List of data centers")
+      dataCenterIds: z7.array(z7.string()).optional().describe("List of data centers"),
+      allowedCudaVersions: z7.array(z7.string()).optional().describe(
+        'Acceptable host CUDA versions as major.minor, e.g. ["12.8","12.6"] (v2, GPU pods only). Exact match \u2014 discover valid values per GPU type with get-capacity. Mutually exclusive with minCudaVersion.'
+      ),
+      minCudaVersion: z7.string().optional().describe(
+        'Lowest acceptable host CUDA version as major.minor, e.g. "12.4" (v2, GPU pods only). A bare major like "12" is rejected. Open-ended floor; mutually exclusive with a non-empty allowedCudaVersions.'
+      )
     },
     { title: "Create pod", ...WRITE },
     async (params) => {
       const backend = backendFor("pods");
+      const { sshPublicKey, ...podParams } = params;
+      let sshKey;
+      if (sshPublicKey !== void 0) {
+        const validated = normalizeSshPublicKey(sshPublicKey);
+        if (!validated.ok)
+          return jsonReply2({ error: validated.error, status: 400 });
+        sshKey = validated.key;
+      }
+      const sshNote = "sshPublicKey applied: PUBLIC_KEY env set and 22/tcp exposed. Full SSH (incl. SCP/SFTP/rsync) also requires an image that installs PUBLIC_KEY into authorized_keys and starts sshd (all runpod/* official images do) and a machine with a public IP \u2014 read the pod\u2019s portMappings for port 22 once it is running.";
+      const hasCudaParams = params.allowedCudaVersions !== void 0 || params.minCudaVersion !== void 0;
+      if (hasCudaParams && backend.version !== "v2") {
+        return jsonReply2({
+          error: "allowedCudaVersions/minCudaVersion are only supported on the v2 REST API. Set RUNPOD_REST_VERSION=v2.",
+          status: 501
+        });
+      }
       if (params.templateId && backend.version !== "v2") {
         return jsonReply2({
           error: "Deploying a pod from a templateId is only supported on the v2 REST API. Set RUNPOD_REST_VERSION=v2, or pass imageName instead.",
@@ -2094,6 +2321,16 @@ function registerPodTools(server, rt) {
             status: 400
           });
         }
+        if (hasCudaParams && !wantsGpu) {
+          return jsonReply2({
+            error: "allowedCudaVersions/minCudaVersion apply to GPU pods only \u2014 pass gpuTypeIds (see list-gpu-types).",
+            status: 400
+          });
+        }
+        if (hasCudaParams) {
+          const cudaError = cudaConstraintError(params);
+          if (cudaError) return jsonReply2({ error: cudaError, status: 400 });
+        }
         if (!wantsCpu && !params.name && !params.templateId) {
           return jsonReply2({
             error: "Provide a name for the pod (deploying from a templateId inherits the template name).",
@@ -2108,13 +2345,15 @@ function registerPodTools(server, rt) {
             });
           }
           try {
+            const cpuBody = sshKey ? applySshPublicKey(podParams, sshKey) : podParams;
             const result = await callRestUrl(
               `${restV1Base(env)}/pods`,
               "POST",
-              params
+              cpuBody
             );
             return jsonReply2({
               ...result,
+              ...sshKey ? { _ssh: sshNote } : {},
               _servedBy: "v1",
               _note: "CPU pod: created on the v1 API, because the v2 REST API does not support CPU pods yet."
             });
@@ -2131,7 +2370,7 @@ function registerPodTools(server, rt) {
           }
         }
       }
-      let body = backend.mapCreate(params);
+      let body = backend.mapCreate(podParams);
       if (params.templateId) {
         let template;
         try {
@@ -2149,19 +2388,21 @@ function registerPodTools(server, rt) {
         }
         body = { ...podBodyFromTemplate(template), ...body };
       }
+      if (sshKey) body = applySshPublicKey(body, sshKey);
       try {
         const result = await callRestUrl(
           `${backend.base}${backend.list}`,
           "POST",
           body
         );
+        const extras = {};
+        if (sshKey) extras._ssh = sshNote;
         if (backend.version === "v2" && (params.gpuTypeIds?.length ?? 0) > 1) {
-          return jsonReply2({
-            ...result,
-            _warning: `v2 accepts one GPU type per pod; used "${params.gpuTypeIds[0]}" and ignored ${params.gpuTypeIds.length - 1} other(s): ${params.gpuTypeIds.slice(1).join(", ")}.`
-          });
+          extras._warning = `v2 accepts one GPU type per pod; used "${params.gpuTypeIds[0]}" and ignored ${params.gpuTypeIds.length - 1} other(s): ${params.gpuTypeIds.slice(1).join(", ")}.`;
         }
-        return jsonReply2(result);
+        return jsonReply2(
+          Object.keys(extras).length ? { ...result, ...extras } : result
+        );
       } catch (error) {
         if (error instanceof HttpError) {
           return jsonReply2({ error: error.message, status: error.status });
@@ -2279,6 +2520,9 @@ function registerPodTools(server, rt) {
 import { z as z8 } from "zod";
 var MIN_QUEUE_DELAY = 0.5;
 var MIN_REQUEST_COUNT = 1;
+function isRequestCountScaler(type) {
+  return type === "REQUEST_COUNT";
+}
 var MIN_IDLE_TIMEOUT = 1;
 var MAX_IDLE_TIMEOUT = 3600;
 function scalerValueError(scalerType, scalerValue) {
@@ -2371,6 +2615,12 @@ function registerEndpointTools(server, rt) {
         'GPU pool names (v2, required). The `pool` field from list-gpu-types, e.g. ["AMPERE_80"]. NOT GPU type ids.'
       ),
       gpuCount: z8.number().optional().describe("GPUs per worker (v2)"),
+      allowedCudaVersions: z8.array(z8.string()).optional().describe(
+        'Acceptable host CUDA versions for worker placement as major.minor, e.g. ["12.8"] (v2). Exact match \u2014 see get-capacity for which versions have stock. Mutually exclusive with minCudaVersion.'
+      ),
+      minCudaVersion: z8.string().optional().describe(
+        'Lowest acceptable host CUDA version as major.minor, e.g. "12.4" (v2). A bare major like "12" is rejected. Mutually exclusive with a non-empty allowedCudaVersions.'
+      ),
       args: z8.string().optional().describe("Container start command/args (v2)"),
       containerDiskInGb: z8.number().optional().describe("Container disk size in GB (v2)"),
       ports: z8.array(z8.string()).optional().describe("Ports to expose (v2), e.g. ['8000/http']"),
@@ -2400,6 +2650,13 @@ function registerEndpointTools(server, rt) {
     { title: "Create endpoint", ...WRITE },
     async (params) => {
       const backend = backendFor("endpoints");
+      const hasCudaParams = params.allowedCudaVersions !== void 0 || params.minCudaVersion !== void 0;
+      if (hasCudaParams && backend.version !== "v2") {
+        return jsonReply2({
+          error: "allowedCudaVersions/minCudaVersion are only supported on the v2 REST API. Set RUNPOD_REST_VERSION=v2.",
+          status: 501
+        });
+      }
       if (backend.version === "v2") {
         if (!params.name) {
           return jsonReply2({
@@ -2432,6 +2689,10 @@ function registerEndpointTools(server, rt) {
         if (scalerError) {
           return jsonReply2({ error: scalerError, status: 400 });
         }
+        if (hasCudaParams) {
+          const cudaError = cudaConstraintError(params);
+          if (cudaError) return jsonReply2({ error: cudaError, status: 400 });
+        }
         const body = backend.mapCreate(params);
         const result2 = await callRestUrl(
           `${backend.base}${backend.list}`,
@@ -2469,7 +2730,7 @@ function registerEndpointTools(server, rt) {
   );
   server.tool(
     "update-endpoint",
-    "Update a Serverless endpoint's config. On v2 you can change image/disk/env/ports/registry/workers/scaling/networkVolumes/timeout/flashboot; on v1, scaling fields (worker min/max, idle timeout, scaler type/value, name). Only provided fields change. An endpoint's request routing (queue vs load balancer) is fixed at creation and cannot be changed here \u2014 recreate the endpoint instead.",
+    "Update a Serverless endpoint's config. On v2 you can change image/disk/env/ports/registry/workers/scaling/networkVolumes/timeout/flashboot, plus the GPU selection (gpuPoolIds/gpuCount) and its CUDA constraints (allowedCudaVersions/minCudaVersion) \u2014 CUDA-only or count-only changes do not require resending gpuPoolIds; on v1, scaling fields (worker min/max, idle timeout, scaler type/value, name). Only provided fields change. An endpoint's request routing (queue vs load balancer) is fixed at creation and cannot be changed here \u2014 recreate the endpoint instead. Note: passing gpuPoolIds replaces the GPU selection wholesale, which clears any GPU-type exclusions set elsewhere (console or set-endpoint-gpus) \u2014 the reply carries a _warning when that happens.",
     {
       endpointId: z8.string().describe("ID of the endpoint to update"),
       name: z8.string().optional().describe("New name for the endpoint"),
@@ -2488,6 +2749,12 @@ function registerEndpointTools(server, rt) {
       imageName: z8.string().optional().describe("New Docker image (v2)"),
       gpuPoolIds: z8.array(z8.string()).optional().describe('New GPU pool names (v2), e.g. ["AMPERE_80"]'),
       gpuCount: z8.number().optional().describe("New GPUs per worker (v2)"),
+      allowedCudaVersions: z8.array(z8.string()).optional().describe(
+        "New acceptable host CUDA versions as major.minor (v2). Changing only this does not require resending gpuPoolIds. An explicit [] clears the constraint. Mutually exclusive with minCudaVersion."
+      ),
+      minCudaVersion: z8.string().optional().describe(
+        'New lowest acceptable host CUDA version as major.minor (v2). An empty string clears the floor; a bare major like "12" is rejected. Mutually exclusive with a non-empty allowedCudaVersions.'
+      ),
       args: z8.string().optional().describe("New container args (v2)"),
       containerDiskInGb: z8.number().optional().describe("New container disk size in GB (v2)"),
       ports: z8.array(z8.string()).optional().describe("New ports (v2)"),
@@ -2502,14 +2769,41 @@ function registerEndpointTools(server, rt) {
       const { endpointId, ...updateParams } = params;
       const backend = backendFor("endpoints");
       const url = `${backend.base}${backend.get(endpointId)}`;
+      const hasCudaParams = updateParams.allowedCudaVersions !== void 0 || updateParams.minCudaVersion !== void 0;
       if (backend.version !== "v2") {
+        if (hasCudaParams) {
+          return jsonReply2({
+            error: "allowedCudaVersions/minCudaVersion are only supported on the v2 REST API. Set RUNPOD_REST_VERSION=v2.",
+            status: 501
+          });
+        }
         const body2 = backend.mapUpdate(updateParams);
         return jsonReply2(await callRestUrl(url, "PATCH", body2));
       }
+      if (hasCudaParams) {
+        const cudaError = cudaConstraintError(updateParams, {
+          allowClear: true
+        });
+        if (cudaError) return jsonReply2({ error: cudaError, status: 400 });
+      }
+      if (updateParams.gpuPoolIds !== void 0 && !updateParams.gpuPoolIds.length) {
+        return jsonReply2({
+          error: "gpuPoolIds cannot be empty \u2014 omit it to keep the current GPU selection.",
+          status: 400
+        });
+      }
       let scalerType = updateParams.scalerType;
-      if (scalerType === void 0 && updateParams.scalerValue !== void 0) {
+      const needScalerRead = scalerType === void 0 && updateParams.scalerValue !== void 0;
+      const needExclusionRead = (updateParams.gpuPoolIds?.length ?? 0) > 0;
+      let clearedExclusions = [];
+      if (needScalerRead || needExclusionRead) {
         const current = await callRestUrl(url);
-        scalerType = current?.scaling?.type === "REQUEST_COUNT" ? "REQUEST_COUNT" : "QUEUE_DELAY";
+        if (needScalerRead) {
+          scalerType = isRequestCountScaler(current?.scaling?.type) ? "REQUEST_COUNT" : "QUEUE_DELAY";
+        }
+        if (needExclusionRead && current?.gpu?.excludedTypes?.length) {
+          clearedExclusions = current.gpu.excludedTypes;
+        }
       }
       const scalerError = scalerValueError(
         scalerType,
@@ -2522,7 +2816,14 @@ function registerEndpointTools(server, rt) {
         ...updateParams,
         scalerType
       });
-      return jsonReply2(await callRestUrl(url, "PATCH", body));
+      const result = await callRestUrl(url, "PATCH", body);
+      if (clearedExclusions.length) {
+        return jsonReply2({
+          ...result,
+          _warning: `Supplying gpuPoolIds replaced the GPU selection wholesale, clearing ${clearedExclusions.length} GPU-type exclusion(s) previously set on this endpoint: ${clearedExclusions.join(", ")}. Re-apply them with set-endpoint-gpus if still wanted.`
+        });
+      }
+      return jsonReply2(result);
     }
   );
   server.tool(
@@ -2622,8 +2923,108 @@ import { z as z9 } from "zod";
 var DIAGNOSIS_TTL_MS = 15e3;
 var MAX_DIAGNOSIS_ENTRIES = 500;
 var diagnosisCache = /* @__PURE__ */ new Map();
+var HTTP_LONG_POLL_BUDGET_MS = 45e3;
+var STDIO_STREAM_BUDGET_MS = 5 * 60 * 1e3;
+var RUNSYNC_UPSTREAM_DEFAULT_WAIT_MS = 9e4;
+var STREAM_UPSTREAM_DEFAULT_WAIT_MS = 1e4;
+var HTTP_STREAM_POLL_WAIT_MS = 1e3;
+var STREAM_JOB_POLL_INTERVAL_MS = 1e3;
+var MAX_CONSECUTIVE_STREAM_ERRORS = 5;
+var TERMINAL_STATUSES = /* @__PURE__ */ new Set([
+  "COMPLETED",
+  "FAILED",
+  "CANCELLED",
+  "TIMED_OUT"
+]);
+function isUnfinishedJobReply(result) {
+  return result !== null && typeof result === "object" && !Array.isArray(result) && !("waitClamped" in result) && !TERMINAL_STATUSES.has(result.status ?? "");
+}
+function formatBudget(ms) {
+  return ms >= 12e4 ? `${Math.round(ms / 6e4)} minutes` : `${Math.round(ms / 1e3)} seconds`;
+}
+var UPSTREAM_HOLD_SLACK_MS = 5e3;
+var MIN_STREAM_POLL_TIMEOUT_MS = 2e3;
+var QUEUED_DIAGNOSIS_TIMEOUT_MS = 5e3;
+function streamPollTimeoutMs(remainingMs, holdMs) {
+  return Math.max(
+    Math.min(remainingMs, holdMs + UPSTREAM_HOLD_SLACK_MS),
+    MIN_STREAM_POLL_TIMEOUT_MS
+  );
+}
+function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+function streamPollPlan(hosted) {
+  return hosted ? {
+    budgetMs: HTTP_LONG_POLL_BUDGET_MS,
+    holdMs: HTTP_STREAM_POLL_WAIT_MS,
+    query: `?wait=${HTTP_STREAM_POLL_WAIT_MS}`
+  } : {
+    budgetMs: STDIO_STREAM_BUDGET_MS,
+    holdMs: STREAM_UPSTREAM_DEFAULT_WAIT_MS,
+    query: ""
+  };
+}
+var RESUME_ADVICE = "Call stream-job again to continue collecting output, get-job-status to check the job without streaming, or stream without a budget by calling the runtime API directly (GET https://api.runpod.ai/v2/{endpointId}/stream/{jobId} with a Bearer API key).";
+function budgetExhaustedNote(budgetMs, lastError) {
+  return {
+    pollingTimedOut: true,
+    note: `Polling stopped after ${formatBudget(
+      budgetMs
+    )} with the job possibly still running. ${RESUME_ADVICE}`,
+    // Surface the trailing error (if any) instead of discarding it — the last
+    // polls may have been failing (e.g. job expired) even though earlier ones
+    // succeeded. Cleared on success, so this is never a stale error from
+    // minutes of healthy streaming ago.
+    ...lastError ? { lastError } : {}
+  };
+}
+function errorCapNote(lastError) {
+  return {
+    error: `Polling aborted after ${MAX_CONSECUTIVE_STREAM_ERRORS} consecutive errors: ${lastError}`,
+    note: `Polling stopped after ${MAX_CONSECUTIVE_STREAM_ERRORS} consecutive errors with the job possibly still running. ${RESUME_ADVICE}`
+  };
+}
+async function collectJobStream(deps) {
+  const { poll, budgetMs, holdMs } = deps;
+  const startedAt = Date.now();
+  const elapsed = () => Date.now() - startedAt;
+  const chunks = [];
+  let result = {};
+  let consecutiveErrors = 0;
+  let lastError;
+  while (true) {
+    try {
+      const reply = await poll(
+        streamPollTimeoutMs(budgetMs - elapsed(), holdMs)
+      );
+      consecutiveErrors = 0;
+      lastError = void 0;
+      if (Array.isArray(reply.stream)) chunks.push(...reply.stream);
+      result = reply;
+      if (TERMINAL_STATUSES.has(reply.status)) break;
+    } catch (error) {
+      consecutiveErrors++;
+      lastError = error instanceof Error ? error.message : String(error);
+      if (consecutiveErrors >= MAX_CONSECUTIVE_STREAM_ERRORS) {
+        result = { ...result, ...errorCapNote(lastError) };
+        break;
+      }
+    }
+    if (elapsed() > budgetMs) {
+      result = { ...result, ...budgetExhaustedNote(budgetMs, lastError) };
+      break;
+    }
+    await sleep(STREAM_JOB_POLL_INTERVAL_MS);
+  }
+  return { result, chunks };
+}
 function registerJobTools(server, rt) {
   const { jsonReply: jsonReply2, serverlessRequest, backendFor, callRestUrl } = rt;
+  const hosted = rt.transport === "http";
+  const streamBudget = formatBudget(
+    hosted ? HTTP_LONG_POLL_BUDGET_MS : STDIO_STREAM_BUDGET_MS
+  );
   async function diagnoseQueuedJob(endpointId) {
     const cached = diagnosisCache.get(endpointId);
     if (cached && cached.expiresAt > Date.now()) {
@@ -2634,7 +3035,10 @@ function registerJobTools(server, rt) {
       const backend = backendFor("workers");
       if (backend.version !== "v2") return null;
       const raw = await callRestUrl(
-        `${backend.base}/serverless/${endpointId}/workers`
+        `${backend.base}/serverless/${endpointId}/workers`,
+        "GET",
+        void 0,
+        { timeoutMs: QUEUED_DIAGNOSIS_TIMEOUT_MS }
       );
       const summary = raw?.summary;
       const workers = raw?.workers ?? [];
@@ -2725,14 +3129,14 @@ function registerJobTools(server, rt) {
   );
   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).",
+    `Submit a synchronous job to a Serverless endpoint and wait for the result. Best for fast tasks: if the job outlives the wait, the response returns a job ID and a non-terminal status (IN_QUEUE or IN_PROGRESS) to poll with get-job-status. Max payload 20 MB; results expire after 1 minute. ${hosted ? `This server waits up to ${HTTP_LONG_POLL_BUDGET_MS} ms \u2014 it runs behind a 60-second gateway deadline, so that is both the default and the ceiling for the wait parameter. For longer jobs use run-endpoint + get-job-status, or POST https://api.runpod.ai/v2/{endpointId}/runsync?wait=300000 directly with a Bearer API key.` : "The wait parameter extends the server-side wait to up to 5 minutes (300000 ms); it defaults to 90000 (90 seconds)."}`,
     {
       endpointId: endpointIdSchema.describe(
         "ID of the Serverless endpoint to run synchronously"
       ),
       input: inputSchema,
       wait: z9.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)."
+        hosted ? `How long in milliseconds the server should wait for a result before returning a job ID to poll. Accepted range 1000\u2013300000, but this server clamps anything above ${HTTP_LONG_POLL_BUDGET_MS} down to ${HTTP_LONG_POLL_BUDGET_MS}, which is also the effective default.` : "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,
@@ -2741,13 +3145,36 @@ function registerJobTools(server, rt) {
     { title: "Run endpoint (sync)", ...WRITE },
     async (params) => {
       const { endpointId, wait, ...body } = params;
-      const path = wait ? `/runsync?wait=${wait}` : "/runsync";
+      const effectiveWait = hosted ? Math.min(
+        wait ?? RUNSYNC_UPSTREAM_DEFAULT_WAIT_MS,
+        HTTP_LONG_POLL_BUDGET_MS
+      ) : wait;
+      const waitQuery = effectiveWait === void 0 ? "" : `?wait=${effectiveWait}`;
+      const path = `/runsync${waitQuery}`;
       const result = await serverlessRequest(
         endpointId,
         path,
         "POST",
-        body
+        body,
+        {
+          // From the wait actually sent — on http the clamped 45s, not the
+          // 300s the caller may have asked for.
+          timeoutMs: (effectiveWait ?? RUNSYNC_UPSTREAM_DEFAULT_WAIT_MS) + UPSTREAM_HOLD_SLACK_MS
+        }
       );
+      const requestedWait = wait ?? RUNSYNC_UPSTREAM_DEFAULT_WAIT_MS;
+      if (hosted && requestedWait > HTTP_LONG_POLL_BUDGET_MS && isUnfinishedJobReply(result)) {
+        return jsonReply2({
+          ...result,
+          waitClamped: {
+            requestedMs: requestedWait,
+            effectiveMs: effectiveWait,
+            reason: `This server runs behind a 60-second gateway deadline, so a non-terminal status here means the job outlived ${formatBudget(
+              HTTP_LONG_POLL_BUDGET_MS
+            )}, not the requested wait \u2014 poll get-job-status rather than treating it as stuck.`
+          }
+        });
+      }
       return jsonReply2(result);
     }
   );
@@ -2781,7 +3208,7 @@ function registerJobTools(server, rt) {
   );
   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.",
+    `Retrieve streaming output from a Serverless job. The worker must support streaming output. Polls /stream/{jobId} and collects chunks until the status is COMPLETED, FAILED, CANCELLED, or TIMED_OUT, for up to ${streamBudget}${hosted ? " (this server runs behind a 60-second gateway deadline)" : ""}. If the budget expires first, returns the chunks collected so far with pollingTimedOut: true \u2014 call stream-job again to resume where it left off, or get-job-status to check without streaming.${hosted ? " To collect with no budget, poll GET https://api.runpod.ai/v2/{endpointId}/stream/{jobId} yourself with a Bearer API key until the status is terminal \u2014 each request drains only the chunks buffered so far, so a single request is not the whole output." : ""}`,
     {
       endpointId: endpointIdSchema.describe(
         "ID of the Serverless endpoint the job belongs to"
@@ -2790,51 +3217,16 @@ function registerJobTools(server, rt) {
     },
     { title: "Stream job", ...READ_ONLY },
     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;
-      let lastError;
-      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++;
-          lastError = error instanceof Error ? error.message : String(error);
-          if (consecutiveErrors >= MAX_CONSECUTIVE_ERRORS) {
-            finalResult.error = `Polling aborted after ${MAX_CONSECUTIVE_ERRORS} consecutive errors: ${lastError}`;
-            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.";
-          if (lastError) finalResult.lastError = lastError;
-          break;
-        }
-        await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
-      }
-      return jsonReply2({ ...finalResult, stream: allChunks });
+      const plan = streamPollPlan(hosted);
+      const streamPath = `/stream/${params.jobId}${plan.query}`;
+      const { result, chunks } = await collectJobStream({
+        budgetMs: plan.budgetMs,
+        holdMs: plan.holdMs,
+        poll: (timeoutMs) => serverlessRequest(params.endpointId, streamPath, "GET", void 0, {
+          timeoutMs
+        })
+      });
+      return jsonReply2({ ...result, stream: chunks });
     }
   );
   server.tool(
@@ -3381,4 +3773,4 @@ export {
   restVersionSummary,
   registerTools
 };
-//# sourceMappingURL=chunk-WGK7D2YR.mjs.map
+//# sourceMappingURL=chunk-G4URDODF.mjs.map