diff --git a/.changeset/per-request-timeouts.md b/.changeset/per-request-timeouts.md new file mode 100644 index 0000000..b54f0a9 --- /dev/null +++ b/.changeset/per-request-timeouts.md @@ -0,0 +1,15 @@ +--- +'@runpod/mcp-server': minor +--- + +Give every outbound request a client-side deadline — every tool call, the OAuth handshake on the hosted server, and the install wizard's key check. node-fetch applies no timeout of its own, so a Runpod host that accepted the connection and then went silent — a wedged worker, a load balancer holding the socket — left a tool call pending forever, and on the hosted server that ended as a bare 504 when Vercel reaped the function at its 60s limit. + +Requests now abort after 30 seconds with a named `RequestTimeoutError` naming the API that went quiet, the deadline it was given, and what to do next. `runsync-endpoint` is the one call that legitimately asks the server to hold a connection open, so it derives its deadline from the `wait` it requested (the server's own 90-second default when `wait` is omitted) rather than being truncated. Successful tool output is unchanged. + +On the hosted transport the cap is a single budget for the whole invocation rather than a fresh allowance per request, because several tools make more than one call — `get-job-status` adds a queued-job diagnosis, `deploy-hub-repo` and `set-endpoint-gpus` read before they write, `update-endpoint` reads the current scaler before patching it — and two full deadlines back to back outlived the platform even with each one bounded. Each request is clamped to what is left, so a stall anywhere in a handler surfaces as the named error instead of a 504. The queued-job diagnosis, which only decorates a status that is already in hand, is bounded at 5 seconds so it cannot spend a budget the reply itself needs. + +`stream-job` bounds each poll by the wait the server was asked to hold rather than by its whole budget. A deadline set to the budget meant one wedged socket consumed the entire run in a single attempt — 45 seconds on the hosted server, five minutes on stdio — so the loop's retry path never ran and a stall returned nothing at all. Each poll now gets the hold plus a round trip, which both clears a reply already in flight and leaves the budget room to reconnect; a run that ends on the budget reports `pollingTimedOut` with the last error rather than discarding it. + +The hosted OAuth routes are bounded too. `/token` polls the flash backend for an approval, and one silent socket there hung the whole handshake until the platform reaped it — the worst place for a blank error, since the user has no credential yet to retry with. Each backend call now has its own 10-second deadline (override with `MCP_FLASH_TIMEOUT_MS`), never exceeding what is left of a 45-second poll budget, and a stall is reported as a named error naming the operation, the host and the deadline. The `runpod-mcp install` wizard's key verification gets the same treatment, so it can no longer sit at "Verifying…" indefinitely. + +A timed-out GraphQL read is also no longer described as a possible write. The advice keys off the HTTP method and GraphQL is always POST on the wire, so `list-gpu-types` timing out used to tell the agent the call may have landed and to "check with the matching list-/get- tool first" — which is the tool that just failed. Only actual mutations carry that warning now. diff --git a/CLAUDE.md b/CLAUDE.md index 3e2a0a8..f66d41e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -43,6 +43,7 @@ The hosted HTTP path (`api/index.ts` + `src/http.ts`) reads these, all optional - `RUNPOD_AUTHED_GRAPHQL_URL`: override the GraphQL host for **authenticated** operations with no REST equivalent — `deploy-hub-repo` and `set-endpoint-gpus` (default `https://api.runpod.io/graphql`). These send the caller's API key as a Bearer token, so only point this at a host you trust with it; on the hosted server that key is a per-user OAuth-minted one. - `RUNPOD_API_KEY_NAME`: name for the minted key (default `runpod-mcp`; set to `""` to omit for a backend without the `apiKeyName` argument). - `MCP_VERBOSE_LOGS`: set to `true` to log OAuth request ids (live auth codes) for debugging. +- `MCP_FLASH_TIMEOUT_MS`: deadline in milliseconds for one call to the flash auth backend during the OAuth flow (default `10000`). Raise it for a slow backend; a value that is not a positive number is ignored, so the deadline cannot be disabled by a typo. The `/token` poll additionally caps itself at 45 seconds total, below the function's `maxDuration`. - `MCP_SKIP_CREDENTIAL_CHECK`: set to the exact string `true` to disable the hosted pre-flight credential verification (dead bearers then surface as tool-level 401 errors instead of an HTTP 401 re-auth signal). Use this if the pre-flight itself is ever causing outages. Note the pre-flight ALSO self-disables when a REST/Serverless host is overridden without a matching `RUNPOD_GRAPHQL_URL`, since it would otherwise validate the key against the wrong environment and reject every request. The build produces `dist/stdio.*`, `dist/http.*`, and `dist/tools.*`. Because `package.json` has `"type": "module"`, always use `dist/stdio.mjs` when running the built local server with `node`. diff --git a/api/index.ts b/api/index.ts index 60223f5..00cdcca 100644 --- a/api/index.ts +++ b/api/index.ts @@ -94,6 +94,61 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } +// node-fetch applies no timeout of its own, and this file runs under +// vercel.json's 60s maxDuration for api/index.ts. A flash backend that accepts +// the connection and goes quiet would otherwise hang the OAuth handshake until +// the platform reaps it — a blank 504 on the one flow the user cannot retry +// their way out of, since they have no credential yet. Tool requests get the +// same treatment through createHttpClient (src/_shared/http.ts); these calls +// build their own request, so they are bounded here. +// Exported so the parsing test asserts the real fallback, not a copy of it. +export const DEFAULT_FLASH_GRAPHQL_TIMEOUT_MS = 10_000; +// Whole-poll budget for /token, below maxDuration so the OAuth error response +// is serialized rather than reaped. 20 attempts spaced 2s apart can otherwise +// reach the platform limit on latency alone, with or without a wedged socket. +// Exported for the test that checks it against vercel.json — the same +// cross-file invariant HTTP_LONG_POLL_BUDGET_MS is pinned by, and the same +// silent failure if maxDuration moves and this does not. +export const TOKEN_POLL_BUDGET_MS = 45_000; +const TOKEN_POLL_INTERVAL_MS = 2_000; +// Budget below which /token stops polling instead of squeezing in one more +// read. Reading an APPROVED request consumes the code atomically upstream, so a +// read we abandon mid-flight can burn it — the user's retry then gets "already +// used" rather than a key. Stopping answers authorization_pending, which is +// retryable and consumes nothing, so it is the better trade for the last few +// seconds of a budget. +const MIN_TOKEN_POLL_REMAINDER_MS = 5_000; +// AbortSignal.timeout takes a uint32: a fractional or out-of-range delay throws +// ERR_OUT_OF_RANGE synchronously, which would 500 BOTH OAuth routes and make +// signing in impossible. Just above int32 it does not throw at all — it warns +// and fires immediately, turning a dial someone raised into a 1ms deadline. +const MAX_FLASH_TIMEOUT_MS = 2_147_483_647; + +/** + * Deadline for one flash-backend call. `MCP_FLASH_TIMEOUT_MS` overrides it — + * an ops dial for a slow backend, and how the timeout tests reach this path in + * milliseconds instead of ten real seconds. Anything that is not a positive + * integer inside the timer's range is ignored, so no value of this variable can + * disable the deadline or break the routes that depend on it. Exported so the + * test for that promise calls the real parser. + */ +export function getFlashTimeoutMs(): number { + const override = Number(process.env.MCP_FLASH_TIMEOUT_MS); + return Number.isSafeInteger(override) && + override > 0 && + override <= MAX_FLASH_TIMEOUT_MS + ? override + : DEFAULT_FLASH_GRAPHQL_TIMEOUT_MS; +} + +// Deadline for one /token poll: the per-call deadline, never past what is left +// of the whole poll. The sibling of streamPollTimeoutMs in src/tools/jobs.ts. +// No floor is needed — the loop below will not start a read at all once the +// budget is down to MIN_TOKEN_POLL_REMAINDER_MS. +function tokenPollDeadlineMs(remainingMs: number): number { + return Math.min(getFlashTimeoutMs(), remainingMs); +} + /** * Name for the minted Runpod API key, shown in the user's dashboard. Defaults * to "runpod-mcp" so keys minted through this server are identifiable and @@ -120,15 +175,34 @@ interface FlashAuthRequestStatus { * non-JSON responses (e.g. an SST live-debug notice when the dev session is * down) as a clear error instead of a cryptic JSON parse failure. */ -async function flashGraphql(query: string, field: string): Promise { +async function flashGraphql( + query: string, + field: string, + requestedTimeoutMs?: number +): Promise { const url = getRunpodGraphqlUrl(); - const response = await fetch(url, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ query }), - }); + const timeoutMs = requestedTimeoutMs ?? getFlashTimeoutMs(); + // Covers the body drain too: a backend can send headers and then stall. + const signal = AbortSignal.timeout(timeoutMs); + let response: Awaited>; + let text: string; + try { + response = await fetch(url, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ query }), + signal, + }); + text = await response.text(); + } catch (error) { + // `signal.aborted`, not the error name — node-fetch and undici label an + // abort differently, and a transport failure must keep its own message. + if (!signal.aborted) throw error; + throw new Error( + `${field} got no response from ${url} after ${timeoutMs}ms — the Runpod API may be unavailable. Start the sign-in again.` + ); + } - const text = await response.text(); let result: { data?: Record; errors?: Array<{ message: string }> }; try { result = JSON.parse(text); @@ -175,12 +249,16 @@ async function createFlashAuthRequest(codeChallenge: string): Promise { * Read the current status of a flash auth request (guest query). Once the user * approves it in the console, the backend mints and returns a Runpod API key. */ -async function getFlashAuthStatus(id: string): Promise { +async function getFlashAuthStatus( + id: string, + timeoutMs?: number +): Promise { return flashGraphql( `query { flashAuthRequestStatus(flashAuthRequestId: ${JSON.stringify( id )}) { id status apiKey codeChallenge codeChallengeMethod } }`, - 'flashAuthRequestStatus' + 'flashAuthRequestStatus', + timeoutMs ); } @@ -447,8 +525,27 @@ async function handleToken( // once (backend: model/src/flash/authRequests.ts), so this poll cannot hand // the same key out twice. const maxAttempts = 20; - for (let attempt = 0; attempt < maxAttempts; attempt++) { - const status = await getFlashAuthStatus(code); + const startedAt = Date.now(); + const remainingBudgetMs = () => + TOKEN_POLL_BUDGET_MS - (Date.now() - startedAt); + // Stops on whichever comes first: the attempt count, or a budget too thin + // to cover a read we would rather not start (MIN_TOKEN_POLL_REMAINDER_MS). + // Both fall through to the authorization_pending answer below. + for ( + let attempt = 0; + attempt < maxAttempts && + remainingBudgetMs() >= MIN_TOKEN_POLL_REMAINDER_MS; + attempt++ + ) { + // Deliberately NOT wrapped in try/continue: an APPROVED read consumes + // the code atomically upstream, so a read that failed on our side may + // already have minted and burned the key. Retrying would then read + // CONSUMED-with-no-key and report "already used" — a worse answer than + // the honest failure. One failed read ends the poll. + const status = await getFlashAuthStatus( + code, + tokenPollDeadlineMs(remainingBudgetMs()) + ); console.log('oauth_token_poll', { attempt, status: status.status, @@ -503,8 +600,10 @@ async function handleToken( return; } - // PENDING — wait and retry. - if (attempt < maxAttempts - 1) await sleep(2000); + // PENDING — wait and retry. Always the full interval: skipping the sleep + // near the end of the budget bought the user no extra approval time and + // just re-read the backend back to back. + if (attempt < maxAttempts - 1) await sleep(TOKEN_POLL_INTERVAL_MS); } tokenError( diff --git a/src/_shared/http.ts b/src/_shared/http.ts index 3d3f932..42b207f 100644 --- a/src/_shared/http.ts +++ b/src/_shared/http.ts @@ -19,6 +19,11 @@ interface RequestInitLike { method: string; headers: Record; body?: string; + // Required so an unbounded request through THIS client is a compile error + // rather than a test. It says nothing about requests built elsewhere — the + // GraphQL helper (tools/runtime.ts), the flash auth calls (api/index.ts) and + // the install wizard each bound their own fetch. + signal: AbortSignal; } type FetchLike = ( @@ -74,6 +79,68 @@ export class HttpError extends Error { } } +// ============== REQUEST DEADLINE ============== +// node-fetch has no timeout of its own, so a server that accepts the connection +// then goes silent leaves the request pending until the platform reaps the +// function — a bare 504, with whatever the tool had collected thrown away. +// Same mechanism as credential-check.ts and backend.ts. +// +// 30s is generous for a control-plane call that answers sub-second, and leaves +// the hosted function room to serialize a real error. runsync-endpoint asks the +// server to hold the connection open, so it raises its own (tools/jobs.ts). +export const DEFAULT_REQUEST_TIMEOUT_MS = 30_000; + +// Floor for a deadline shrunk by a nearly-spent invocation budget (see +// `maxTimeoutMs`). A 0ms deadline aborts before the socket opens and reports +// "no response after 0ms", which reads as a server fault rather than a budget +// we had already used up; a second is enough for a warm host to answer, and +// overshooting an exhausted budget by 1s is strictly better than the bare 504 +// the whole deadline exists to replace. +export const MIN_REQUEST_TIMEOUT_MS = 1_000; + +// A bare AbortError names neither the API nor the deadline, and the MCP SDK +// hands that string straight to the agent. +export class RequestTimeoutError extends Error { + readonly timeoutMs: number; + readonly method: string; + constructor(prefix: string, timeoutMs: number, method: string) { + // We abandoned the request, we did not undo it, and nothing sends an + // idempotency key — so a timed-out write may well have landed. Retry advice + // has to key off the method or an agent double-creates a billed resource. + 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 — 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; + } +} + +// Covers the whole exchange, not just the connect: a server can send headers +// then stall mid-body. Shared with the GraphQL helper in tools/runtime.ts, +// which builds its own request but has the same failure mode. +export async function withRequestTimeout( + errorPrefix: string, + timeoutMs: number, + method: string, + run: (signal: AbortSignal) => Promise +): Promise { + const signal = AbortSignal.timeout(timeoutMs); + try { + return await run(signal); + } catch (error) { + // `signal.aborted`, not an error name: node-fetch, undici and the test + // fakes each label an abort differently. A real status beats a timeout. + if (error instanceof HttpError || !signal.aborted) throw error; + throw new RequestTimeoutError(errorPrefix, timeoutMs, method); + } +} + // A response whose content-type marks it as JSON — including v2's RFC-9457 // `application/problem+json` error bodies. We match the `+json`/`/json` shape, // NOT the literal substring `application/json` (which `problem+json` does not @@ -104,11 +171,50 @@ function buildRequestHeaders( }; } +// Trailing and optional so existing call sites are untouched and only the +// calls that need a different deadline mention one. +export interface RequestOptions { + // Clamped by the client's `maxTimeoutMs`. Only runsync-endpoint sets it. + timeoutMs?: number; +} + +// A per-request deadline bounds one stalled socket; it does not bound a handler +// that makes several calls in a row. get-job-status (status, then the queued-job +// diagnosis), deploy-hub-repo (catalog, then the saveEndpoint mutation) and +// update-endpoint (read the scaler, then PATCH) each issue two, so two default +// deadlines back to back outlast the platform and the 504 comes back. +// +// So the hosted ceiling is a function, not a constant: it reports what is left +// of the whole invocation, and every request is clamped to that. The first call +// may take the lot; the second only gets the remainder. +export type TimeoutCeiling = number | (() => number); + +function resolveCeiling(ceiling: TimeoutCeiling | undefined): number { + if (ceiling === undefined) return Infinity; + // Only a thunk gets the floor. It is the one that decays as the invocation is + // spent and so can reach zero on its own; a static ceiling is a number + // someone chose, and silently raising it would be the surprise. + if (typeof ceiling !== 'function') return ceiling; + return Math.max(ceiling(), MIN_REQUEST_TIMEOUT_MS); +} + +// Honoring an override past the platform limit just trades the named timeout +// back for the 504 it exists to replace. Shared with the GraphQL helper in +// tools/runtime.ts, which builds its own request but is under the same budget. +export function clampTimeout( + requestedMs: number | undefined, + ceiling: TimeoutCeiling | undefined, + defaultMs: number = DEFAULT_REQUEST_TIMEOUT_MS +): number { + return Math.min(requestedMs ?? defaultMs, resolveCeiling(ceiling)); +} + export interface HttpClient { ( url: string, method?: string, - body?: Record + body?: Record, + options?: RequestOptions ): Promise; } @@ -119,42 +225,70 @@ export function createHttpClient(deps: { // Distinct per backend so error messages stay attributable // ("Runpod API Error" / "Runpod Serverless API Error"). errorPrefix: string; + // Deadline for calls that pass no `timeoutMs` of their own. + defaultTimeoutMs?: number; + // Ceiling on the default and on any override, for a caller under a platform + // deadline it cannot outlive (HTTP_TRANSPORT_BUDGET_MS in tools/runtime.ts). + // A thunk re-reads the remaining budget per call. Unset = no ceiling. + maxTimeoutMs?: TimeoutCeiling; }): HttpClient { + const defaultTimeoutMs = deps.defaultTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS; + return async function request( url: string, method: string = 'GET', - body?: Record + body?: Record, + options?: RequestOptions ): Promise { - const init: RequestInitLike = { + 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 - : undefined - ); - } - - // 204 / empty / non-JSON → a uniform success marker (matches today's helpers - // and covers pod-action responses that return no JSON body). - return isJsonContentType(response.headers.get('content-type')) - ? response.json() - : { success: true, status: response.status }; + async (signal) => { + const init: RequestInitLike = { + 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) { + // Status and hint first: the body is awaited as a constructor + // argument, so a deadline firing mid-drain would reject before the + // HttpError exists and lose the 429 retry hint / 401 re-auth signal. + const hint = + response.status === 429 + ? rateLimitHint( + response.headers.get('ratelimit'), + response.headers.get('retry-after') + ) + : response.status === 401 + ? EXPIRED_CREDENTIAL_HINT + : undefined; + const text = await response + .text() + .catch( + () => '' + ); + throw new HttpError(deps.errorPrefix, response.status, text, hint); + } + + // 204 / empty / non-JSON → a uniform success marker (matches today's + // helpers and covers pod-action responses that return no JSON body). + return isJsonContentType(response.headers.get('content-type')) + ? response.json() + : { success: true, status: response.status }; + } + ); }; } diff --git a/src/install/verify-key.ts b/src/install/verify-key.ts new file mode 100644 index 0000000..b0d1e1e --- /dev/null +++ b/src/install/verify-key.ts @@ -0,0 +1,36 @@ +// Split out of wizard.ts, which imports @clack/prompts — and @clack/core reaches +// for `styleText` from node:util, which does not exist on Node 18. Importing the +// wizard from a test therefore fails to load the whole file on our oldest +// supported runtime. This has no interactive dependencies, so the deadline below +// can be tested directly. + +// Interactive: the user is watching a spinner, so fail fast enough to retype a +// key rather than wait out a network stall. +export const VERIFY_API_KEY_TIMEOUT_MS = 10_000; + +// Verify the API key works by calling a read-only REST endpoint. Returns true +// on success, false on auth failure, and null when the check itself failed +// (offline, timed out, etc.) so we can warn without blocking. The deadline is a +// parameter so a test can drive the stall in milliseconds; the wizard always +// takes the default. +export async function verifyApiKey( + apiKey: string, + timeoutMs: number = VERIFY_API_KEY_TIMEOUT_MS +): Promise { + try { + const base = process.env.RUNPOD_REST_API_URL ?? 'https://rest.runpod.io/v1'; + const response = await fetch(`${base}/pods`, { + headers: { Authorization: `Bearer ${apiKey}` }, + // Without this the wizard sits at "Verifying…" forever against a host + // that accepts the connection and goes quiet, with no way out but ^C. + // An abort lands in the catch below and is reported as "check failed", + // which is exactly what it is. + signal: AbortSignal.timeout(timeoutMs), + }); + if (response.ok) return true; + if (response.status === 401 || response.status === 403) return false; + return null; + } catch { + return null; + } +} diff --git a/src/install/wizard.ts b/src/install/wizard.ts index 068d1fb..9c0e697 100644 --- a/src/install/wizard.ts +++ b/src/install/wizard.ts @@ -6,6 +6,7 @@ import { type AddMode, type McpClient, } from './clients.js'; +import { verifyApiKey } from './verify-key.js'; const API_KEYS_URL = 'https://www.runpod.io/console/user/settings'; @@ -29,23 +30,6 @@ function openBrowser(url: string): void { } } -// Verify the API key works by calling a read-only REST endpoint. Returns true -// on success, false on auth failure, and null when the check itself failed -// (offline, etc.) so we can warn without blocking. -async function verifyApiKey(apiKey: string): Promise { - try { - const base = process.env.RUNPOD_REST_API_URL ?? 'https://rest.runpod.io/v1'; - const response = await fetch(`${base}/pods`, { - headers: { Authorization: `Bearer ${apiKey}` }, - }); - if (response.ok) return true; - if (response.status === 401 || response.status === 403) return false; - return null; - } catch { - return null; - } -} - function bail(message = 'Cancelled.'): never { p.cancel(message); process.exit(0); diff --git a/src/tools/jobs.ts b/src/tools/jobs.ts index f41a370..1d18a52 100644 --- a/src/tools/jobs.ts +++ b/src/tools/jobs.ts @@ -27,21 +27,33 @@ export function clearQueuedJobDiagnosisCache(): void { // An http deployment is assumed to sit behind a gateway that reaps the request // mid-flight; Runpod's hosted one does, at 60s (vercel.json maxDuration). Wait -// longer than the budget and the gateway kills the call before the tool's own -// timeout path runs: bare 504, collected output discarded. 45s leaves room for -// the credential pre-flight and v2 probe (4s each). Exported for the test that -// checks it against vercel.json; stdio has no deadline. +// longer and the gateway kills the call before the tool's own timeout path +// runs: bare 504, collected output discarded. 45s leaves room for the 4s +// credential pre-flight that precedes dispatch (the v2 probe does NOT apply — +// stdio-only, see backend.ts). Exported for the test that checks it against +// vercel.json; stdio has no deadline. export const HTTP_LONG_POLL_BUDGET_MS = 45_000; -const STDIO_STREAM_BUDGET_MS = 5 * 60 * 1000; +// Exported so the poll-deadline tests measure attempts against the real stdio +// budget rather than a copy of it. +export const STDIO_STREAM_BUDGET_MS = 5 * 60 * 1000; // Upstream defaults, mirrored not derived — re-check against ai-api // (pkg/api/runsync.go, pkg/api/stream.go) if the service changes. const RUNSYNC_UPSTREAM_DEFAULT_WAIT_MS = 90_000; +// What an empty GET /stream holds for when no ?wait= is sent. stdio omits the +// query, so this — not the http value below — is the hold its polls bracket. +// Exported so those tests read the hold instead of restating it. +export const STREAM_UPSTREAM_DEFAULT_WAIT_MS = 10_000; // Caps how long an empty /stream may hold the poll. Left at the server's 10s // default, a chunk-sparse job overshoots the 45s budget to ~54s, since the -// budget is only checked between polls. Accepted range 1000–300000. -const HTTP_STREAM_POLL_WAIT_MS = 1000; +// budget is only checked between polls. Accepted range 1000–300000. Exported +// for the same reason as its stdio sibling above. +export const HTTP_STREAM_POLL_WAIT_MS = 1000; // Exported so the budget tests tick the real interval, not a copy of it. export const STREAM_JOB_POLL_INTERVAL_MS = 1000; +// Consecutive failures that end the poll. Only reachable while each attempt is +// bounded well inside the budget — see streamPollTimeoutMs. Exported so the +// guard asserting the budget fits that many attempts counts the real cap. +export const MAX_CONSECUTIVE_STREAM_ERRORS = 5; // Shared by stream-job (stop polling) and runsync (nothing was lost to the // clamp), so the two can't drift on what "finished" means. const TERMINAL_STATUSES = new Set([ @@ -70,6 +82,161 @@ function formatBudget(ms: number): string { : `${Math.round(ms / 1000)} seconds`; } +// A client deadline must outlast the wait the server was asked to hold, or we +// abort the reply we are waiting for; this is the round-trip room on top. Used +// by both long-poll callers (runsync's ?wait=, stream-job's per-poll hold). +// Must also fit under the http ceiling, or the clamp lands the deadline back on +// the hold (asserted in tests/http.test.ts). +export const UPSTREAM_HOLD_SLACK_MS = 5_000; +// So the last poll of a nearly-spent budget can still answer. Exported so the +// floor test asserts the real value. +export const MIN_STREAM_POLL_TIMEOUT_MS = 2_000; +// The queued-job diagnosis is enrichment on an answer we already have, and it is +// discarded on any failure. The shared invocation budget stops it from causing a +// 504, but spending 30 of the caller's seconds to decorate a reply that was +// ready is still the wrong trade — the status belongs to the caller, the hint is +// a bonus. Short enough that losing it costs little. +export const QUEUED_DIAGNOSIS_TIMEOUT_MS = 5_000; + +// Deadline for ONE /stream poll: the hold it brackets, bounded by what is left +// of the budget, and never below MIN_STREAM_POLL_TIMEOUT_MS. The budget is a +// ceiling here, never the target — +// +// above the hold — a poll that aborts before the server's own wait elapses +// kills a reply that was on its way, every time. +// below the budget — a deadline set TO the budget lets one wedged socket +// spend the entire run in a single attempt, so +// MAX_CONSECUTIVE_STREAM_ERRORS never engages and the +// caller gets nothing back from a stall the retry loop was +// built to survive. +// +// The floor only bites once the budget is nearly spent, where a 0ms deadline +// would report a server fault for time we had already used. On the hosted +// transport it is not the last word either: clampTimeout then measures the +// result against the remaining invocation budget, whose own floor is +// MIN_REQUEST_TIMEOUT_MS (src/_shared/http.ts). +export function streamPollTimeoutMs( + remainingMs: number, + holdMs: number +): number { + return Math.max( + Math.min(remainingMs, holdMs + UPSTREAM_HOLD_SLACK_MS), + MIN_STREAM_POLL_TIMEOUT_MS + ); +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +// Everything about a stream-job run that depends on the transport, chosen in +// one place. The query and the hold MUST agree — a poll that asks the server +// for one wait while its deadline brackets another either aborts replies in +// flight or waits far past what it should — and bundling the budget with them +// means a wrong transport here cannot slip past the budget tests. +// stdio sends no ?wait= (fewer requests, no platform deadline to race), so its +// hold is the server's own default rather than one we picked. +export function streamPollPlan(hosted: boolean): { + budgetMs: number; + holdMs: number; + query: string; +} { + 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: '', + }; +} + +// Both unfinished exits (budget spent, error cap) say the same two things: what +// stopped the run, and that /stream drains what it hands out — so calling again +// resumes where this run stopped rather than replaying from the start. Shared so +// an agent that hits the error cap is not left with chunks and no way forward. +const 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: number, + lastError: string | undefined +): Record { + 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: string): Record { + 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}`, + }; +} + +// Poll until the job reaches a terminal status, the budget runs out, or the API +// fails MAX_CONSECUTIVE_STREAM_ERRORS times in a row. Lifted out of the handler +// so those three exits read in one screen, and so a test can drive the loop +// through `poll` without an MCP server around it. +export async function collectJobStream(deps: { + poll: (timeoutMs: number) => Promise>; + budgetMs: number; + holdMs: number; +}): Promise<{ result: Record; chunks: unknown[] }> { + const { poll, budgetMs, holdMs } = deps; + const startedAt = Date.now(); + const elapsed = () => Date.now() - startedAt; + const chunks: unknown[] = []; + // Never mutated in place: `poll` is caller-supplied now that this is + // exported, and the last reply is not ours to annotate. + let result: Record = {}; + let consecutiveErrors = 0; + let lastError: string | undefined; + + while (true) { + try { + const reply = await poll( + streamPollTimeoutMs(budgetMs - elapsed(), holdMs) + ); + // A success clears both: only CONSECUTIVE failures end the run, so a + // flaky endpoint that answers in between keeps streaming — and the error + // reported at the end is one that was still happening at the end. + consecutiveErrors = 0; + lastError = undefined; + if (Array.isArray(reply.stream)) chunks.push(...reply.stream); + result = reply; + if (TERMINAL_STATUSES.has(reply.status as string)) 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 }; +} + export function registerJobTools(server: McpServer, rt: ToolRuntime): void { const { jsonReply, serverlessRequest, backendFor, callRestUrl } = rt; // Built per instance so each caller is told only the budget that applies to @@ -109,7 +276,10 @@ export function registerJobTools(server: McpServer, rt: ToolRuntime): void { 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', + undefined, + { timeoutMs: QUEUED_DIAGNOSIS_TIMEOUT_MS } )) as | { summary?: WorkerSummary; @@ -297,7 +467,14 @@ export function registerJobTools(server: McpServer, rt: ToolRuntime): void { endpointId, path, 'POST', - body as Record + body as Record, + { + // 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, + } ); // Unmarked, a reply from a job 45s in is identical to one from a job that @@ -377,65 +554,19 @@ export function registerJobTools(server: McpServer, rt: ToolRuntime): void { }, { title: 'Stream job', ...READ_ONLY }, async (params) => { - const MAX_POLL_TIME_MS = hosted - ? HTTP_LONG_POLL_BUDGET_MS - : STDIO_STREAM_BUDGET_MS; - const MAX_CONSECUTIVE_ERRORS = 5; - const allChunks: unknown[] = []; - let finalResult: Record = {}; - let consecutiveErrors = 0; - let lastError: string | undefined; - const startTime = Date.now(); - // stdio keeps the server's default hold: fewer requests, no deadline to race. - const streamQuery = - rt.transport === 'http' ? `?wait=${HTTP_STREAM_POLL_WAIT_MS}` : ''; - const streamPath = `/stream/${params.jobId}${streamQuery}`; - - while (true) { - try { - const result = (await serverlessRequest( - params.endpointId, - streamPath - )) as Record; - - consecutiveErrors = 0; - - if (Array.isArray(result.stream)) { - allChunks.push(...result.stream); - } - - finalResult = result; - - if (TERMINAL_STATUSES.has(result.status as string)) { - 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; - // /stream drains what it hands out, so calling again resumes where - // this run stopped rather than replaying from the start. - finalResult.note = `Polling stopped after ${formatBudget(MAX_POLL_TIME_MS)} with the job possibly still running. 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).`; - // Surface the most recent error (if any) instead of discarding it — - // the last poll may have been failing (e.g. job expired) even though - // earlier polls succeeded. - if (lastError) finalResult.lastError = lastError; - break; - } - - await new Promise((resolve) => - setTimeout(resolve, STREAM_JOB_POLL_INTERVAL_MS) - ); - } - - return jsonReply({ ...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', undefined, { + timeoutMs, + }) as Promise>, + }); + + return jsonReply({ ...result, stream: chunks }); } ); diff --git a/src/tools/runtime.ts b/src/tools/runtime.ts index de457ee..5965728 100644 --- a/src/tools/runtime.ts +++ b/src/tools/runtime.ts @@ -2,9 +2,13 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import fetch from 'node-fetch'; import { randomUUID } from 'node:crypto'; import { + clampTimeout, createHttpClient, EXPIRED_CREDENTIAL_HINT, HttpError, + withRequestTimeout, + type RequestOptions, + type TimeoutCeiling, } from '../_shared/http.js'; import { rateLimitHint } from '../_shared/rate-limit.js'; import { buildTrackingHeaders } from '../_shared/tracking.js'; @@ -98,6 +102,32 @@ function trackingHeaders( }); } +// A BACKSTOP for a wedged socket, not a budget. It has to clear the wait +// runsync asks the server to hold (or we abort a reply in flight) and still +// fire before the platform — measured from the INVOCATION, which starts with +// the 4s credential pre-flight, not from our request. That is what makes 55s +// wrong. Both bounds asserted in tests/http.test.ts. +// +// 0s invocation starts +// 4s credential pre-flight, worst case, before any tool request +// 54s this backstop fires +// 60s vercel.json maxDuration +export const HTTP_TRANSPORT_BUDGET_MS = 50_000; + +// Spent across the whole invocation, not per request: a handler that issues two +// calls (get-job-status, deploy-hub-repo, update-endpoint) would otherwise get +// two full deadlines and outlive the platform anyway. Safe to anchor at runtime +// construction because on http the server, and so this runtime, is built per +// request and disposed with the response (src/http.ts). +// Exported for the test that pins the reset: if registerTools were ever hoisted +// out of the http request handler to save cold-start time, this window would +// keep decaying across requests and every request after the first would run on +// the floor — a silent, load-shaped failure rather than an obvious one. +export function invocationBudget(totalMs: number): () => number { + const startedAt = Date.now(); + return () => totalMs - (Date.now() - startedAt); +} + // The fetch implementation the unified client uses. Defaults to node-fetch; // tests inject a fake to capture outbound requests offline (the A4 seam). export type HttpFetch = Parameters[0]['fetch']; @@ -154,6 +184,13 @@ export interface ToolDeps { // whole reader (and so bypasses the 401 observer); this replaces only its // transport, which is what lets a test drive an SSE 401 through the observer. sseFetch?: SseFetch; + // Test seam: shrink the deadline so a suite need not wait out the real one. + defaultTimeoutMs?: number; + // Ceiling on every request's deadline, including one a tool asks to lengthen + // (stream-job's per-poll hold, runsync's ?wait=). The hosted transport + // installs the remaining invocation budget here; a test pins it low so a + // stalled socket aborts in milliseconds instead of seconds. + maxTimeoutMs?: TimeoutCeiling; } // A bounded Server-Sent-Events read. Returns the raw accumulated stream text @@ -187,23 +224,27 @@ export interface ToolRuntime { variables?: Record ) => Promise; // Authenticated v1 REST call, path-relative to the v1 base (e.g. `/endpoints`). + runpodRequest: ( endpoint: string, method?: string, - body?: Record + body?: Record, + options?: RequestOptions ) => Promise; // Authenticated Serverless runtime call (api.runpod.ai/v2/{endpointId}{path}). serverlessRequest: ( endpointId: string, path: string, method?: string, - body?: Record + body?: Record, + options?: RequestOptions ) => Promise; // Authenticated REST call to a fully-resolved URL (the adapter builds the URL). callRestUrl: ( url: string, method?: string, - body?: Record + body?: Record, + options?: RequestOptions ) => Promise; // Resolve a resource's v1/v2 backend descriptor for the current env/transport. backendFor: (resource: Resource) => Backend; @@ -236,6 +277,49 @@ async function graphqlRequest( tracking: () => Record, fetchImpl: HttpFetch, url: string, + options?: { + variables?: Record; + apiKey?: string; + // See http.ts. + timeoutMs?: number; + // Remaining invocation budget, same thunk the REST clients are capped by. + // deploy-hub-repo spends two GraphQL calls in one invocation. + maxTimeoutMs?: TimeoutCeiling; + } +): Promise { + // Builds its own request rather than going through createHttpClient, so the + // deadline is wired separately — otherwise list-gpu-types / get-capacity + // against a wedged host hang the way the REST calls used to. + 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) + ); +} + +// Anonymous operations (`{ gpuTypes { ... } }`) and the explicit `query` keyword +// are both reads; only `mutation` writes. Leading whitespace and comments are +// the shapes that actually occur in this file's template literals. +export function isGraphqlMutation(query: string): boolean { + 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: string, + tracking: () => Record, + fetchImpl: HttpFetch, + url: string, + signal: AbortSignal, options?: { variables?: Record; apiKey?: string; @@ -252,6 +336,7 @@ async function graphqlRequest( query, ...(options?.variables ? { variables: options.variables } : {}), }), + signal, }); // HTTP-level failures used to fall through to response.json() and surface @@ -364,20 +449,38 @@ export function createToolRuntime( const rawFetch = deps.fetch ?? defaultFetch; const httpFetch = observeUnauthorized(rawFetch); + // Deadline policy for all three JSON clients. One budget object shared by all + // of them, so a handler that crosses clients (get-job-status: serverless, then + // REST for the diagnosis) draws down a single allowance. + const remainingHttpBudgetMs = invocationBudget(HTTP_TRANSPORT_BUDGET_MS); + const timeouts = { + defaultTimeoutMs: deps.defaultTimeoutMs, + maxTimeoutMs: + deps.maxTimeoutMs ?? + (ctx.transport === 'http' ? remainingHttpBudgetMs : undefined), + }; + // v1 REST client (path-relative to the v1 base). const v1Client = createHttpClient({ apiKey: ctx.apiKey, fetch: httpFetch, tracking, errorPrefix: 'Runpod API Error', + ...timeouts, }); const runpodRequest = ( endpoint: string, method: string = 'GET', - body?: Record + body?: Record, + options?: RequestOptions ) => withApiErrorLog('Error calling Runpod API:', () => - v1Client(`${restV1Base(process.env as Env)}${endpoint}`, method, body) + v1Client( + `${restV1Base(process.env as Env)}${endpoint}`, + method, + body, + options + ) ); // Serverless runtime client (endpointId + path against the serverless base). @@ -386,18 +489,21 @@ export function createToolRuntime( fetch: httpFetch, tracking, errorPrefix: 'Runpod Serverless API Error', + ...timeouts, }); const serverlessRequest = ( endpointId: string, path: string, method: string = 'GET', - body?: Record + body?: Record, + options?: RequestOptions ) => withApiErrorLog('Error calling Runpod Serverless API:', () => serverlessClient( `${serverlessBase(process.env as Env)}/${endpointId}${path}`, method, - body + body, + options ) ); @@ -408,14 +514,16 @@ export function createToolRuntime( fetch: httpFetch, tracking, errorPrefix: 'Runpod API Error', + ...timeouts, }); const callRestUrl = ( url: string, method: string = 'GET', - body?: Record + body?: Record, + options?: RequestOptions ): Promise => withApiErrorLog('Error calling Runpod API:', () => - restClient(url, method, body) + restClient(url, method, body, options) ); // Bounded SSE reader for stream-pod-logs / stream-worker-logs. Uses node-fetch @@ -467,13 +575,19 @@ export function createToolRuntime( 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: string) => graphqlRequest( query, tracking, // rawFetch, not httpFetch: this call carries no credential (see above). rawFetch, - publicGraphqlBase(process.env as Env) + publicGraphqlBase(process.env as Env), + { + timeoutMs: deps.defaultTimeoutMs, + maxTimeoutMs: timeouts.maxTimeoutMs, + } ), graphqlAuthed: (query: string, variables?: Record) => graphqlRequest( @@ -482,7 +596,12 @@ export function createToolRuntime( httpFetch, // NOT publicGraphqlBase: this call carries the caller's API key. authedGraphqlBase(process.env as Env), - { variables, apiKey: ctx.apiKey } + { + variables, + apiKey: ctx.apiKey, + timeoutMs: deps.defaultTimeoutMs, + maxTimeoutMs: timeouts.maxTimeoutMs, + } ), runpodRequest, serverlessRequest, diff --git a/tests/handlers.test.ts b/tests/handlers.test.ts index 024bc50..5c36c97 100644 --- a/tests/handlers.test.ts +++ b/tests/handlers.test.ts @@ -24,9 +24,21 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { registerTools } from '../src/tools.js'; import { clearQueuedJobDiagnosisCache, + collectJobStream, + streamPollPlan, + streamPollTimeoutMs, HTTP_LONG_POLL_BUDGET_MS, + HTTP_STREAM_POLL_WAIT_MS, + MAX_CONSECUTIVE_STREAM_ERRORS, + QUEUED_DIAGNOSIS_TIMEOUT_MS, + STDIO_STREAM_BUDGET_MS, STREAM_JOB_POLL_INTERVAL_MS, + STREAM_UPSTREAM_DEFAULT_WAIT_MS, + UPSTREAM_HOLD_SLACK_MS, } from '../src/tools/jobs.js'; +import { isGraphqlMutation } from '../src/tools/runtime.js'; +import { TOKEN_POLL_BUDGET_MS } from '../api/index.js'; +import { DEFAULT_REQUEST_TIMEOUT_MS } from '../src/_shared/http.js'; // ============== Handler integration / outbound-request golden ============== // Drives the REAL registerTools against a fake McpServer (captures handlers) and @@ -41,6 +53,111 @@ interface OutboundRecord { method: string; body?: string; headers?: Record; + signal?: AbortSignal; +} + +// Accepts the connection then stays quiet, rejecting on the caller's signal +// the way a real fetch does. +function stall(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(resolve, ms); + signal?.addEventListener('abort', () => { + clearTimeout(timer); + reject(signal.reason ?? new Error('aborted')); + }); + }); +} + +// A deterministic clock for the poll loops. Two clocks matter: setTimeout (the +// 1s sleep between polls) and Date.now (the budget check). setTimeout is mocked +// via node:test mock timers — Node >=20.11 takes `{apis}`, Node 18 takes a bare +// array, so try both. Date is NOT mockable on Node 18 at all, so it is stubbed +// by hand on every version and advanced in lockstep with the ticks. +// `t` is node:test's own TestContext, NOT a hand-rolled structural type: an +// `enable: (o: unknown) => void` field cannot accept the real signature under +// strictFunctionTypes (a function taking `unknown` is not assignable from one +// taking MockTimersOptions), so that annotation compiled only while tests were +// excluded from type-check. The Node 18 bare-array form is handled by the cast +// in the catch below. +function useFakeClock(t: TestContext) { + const enable = t.mock?.timers?.enable; + if (typeof enable !== 'function') { + throw new Error( + 'node:test mock timers unavailable (needs Node >=18.19) — this test cannot run on this runtime' + ); + } + // Captured before the mock replaces the global: AbortSignal.timeout runs on + // an internal timer the mock does not touch, so a test that waits out a REAL + // deadline needs a real delay to wait with. + const realSetTimeout = globalThis.setTimeout; + try { + t.mock.timers.enable({ apis: ['setTimeout'] }); + } catch (err) { + // Node 18 takes a bare array and rejects the object form with a TypeError + // naming the `timers` argument. Anything else is a real failure, not a + // signature mismatch, so don't swallow it behind a second attempt. + if (!(err instanceof TypeError)) throw err; + (t.mock.timers.enable as unknown as (apis: string[]) => void)([ + 'setTimeout', + ]); + } + let fakeNow = 0; + const realNow = Date.now; + Date.now = () => fakeNow; + // Registered here rather than left to each caller's try/finally: a test that + // throws mid-drive would otherwise leave every later test in this file on a + // frozen clock. + t.after(() => { + Date.now = realNow; + }); + return { + // Move both clocks together: a pending sleep fires and the budget ages by + // the same amount, which is what the real runtime does. + advance(ms: number) { + fakeNow += ms; + t.mock.timers.tick(ms); + }, + // Age the budget WITHOUT firing timers, for a test that wants the loop to + // notice an exhausted budget on its next check rather than poll again. + jump(ms: number) { + fakeNow += ms; + }, + // setImmediate is not mocked, so this lets in-flight microtasks (the fake + // fetch and its json()) settle before the clock moves again. + settle: () => new Promise((resolve) => setImmediate(resolve)), + // Wait on the condition, not on a guess about how long a real deadline + // takes: a fixed sleep sized to a 40ms abort is the kind of test that only + // fails on a loaded CI box. + async waitFor( + predicate: () => boolean, + description: string, + timeoutMs = 5_000 + ) { + const deadline = realNow() + timeoutMs; + while (!predicate()) { + if (realNow() > deadline) { + throw new Error( + `timed out after ${timeoutMs}ms waiting for ${description}` + ); + } + await new Promise((resolve) => realSetTimeout(resolve, 5)); + } + // Let whatever the predicate observed finish settling. + await new Promise((resolve) => setImmediate(resolve)); + }, + }; +} + +// node:test has no default timeout, so a loop that fails to terminate would +// stall every remaining test in this file until CI's job timeout instead of +// failing here. Every direct drive of a poll loop goes through this. +async function settled(pending: Promise, what: string): Promise { + const done = await Promise.race([ + pending.then(() => true), + new Promise((resolve) => setImmediate(() => resolve(false))), + ]); + if (!done) throw new Error(`${what} was still running when the drive ended`); + return pending; } function harness(opts?: { @@ -74,6 +191,14 @@ function harness(opts?: { ) => Promise<{ raw: string; truncated: boolean }>; // Observer for the ToolContext 401 hook (hosted credential invalidation). onUnauthorized?: () => void; + // Deadline seam + a server that stalls `delayMs`, so the timeout is + // exercised at millisecond scale instead of the real 30s. + defaultTimeoutMs?: number; + delayMs?: number; + // Ceiling seam. Unlike defaultTimeoutMs it also binds a tool that asks for a + // LONGER deadline (stream-job's per-poll hold, runsync's wait), which is the + // only way to exercise those at millisecond scale. + maxTimeoutMs?: number; // Transport under the SSE reader, so a test can drive an SSE 401 through the // observer (injecting `streamSse` replaces the reader and bypasses it). sseStatus?: number; @@ -98,14 +223,21 @@ function harness(opts?: { const steps = opts?.steps ? [...opts.steps] : null; const fakeFetch = async ( url: string, - init: { method: string; headers: Record; body?: string } + init: { + method: string; + headers: Record; + body?: string; + signal?: AbortSignal; + } ) => { outbound.push({ url, method: init.method, body: init.body, headers: init.headers, + signal: init.signal, }); + if (opts?.delayMs) await stall(opts.delayMs, init.signal); const step = steps?.shift(); const status = step?.status ?? opts?.status ?? 200; const jsonBody = step @@ -141,6 +273,10 @@ function harness(opts?: { fetch: fakeFetch as NonNullable< Parameters[2] >['fetch'], + ...(opts?.defaultTimeoutMs + ? { defaultTimeoutMs: opts.defaultTimeoutMs } + : {}), + ...(opts?.maxTimeoutMs ? { maxTimeoutMs: opts.maxTimeoutMs } : {}), ...(opts?.streamSse ? { streamSse: opts.streamSse } : {}), ...(opts?.sseStatus ? { @@ -720,10 +856,7 @@ describe('pod routing under RUNPOD_REST_VERSION=v2', () => { }); assert.equal(outbound.length, 2); // 1) template GET - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/templates/tpl_1' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/templates/tpl_1'); assert.equal(outbound[0].method, 'GET'); // 2) pod POST with the template's container config folded in assert.equal(outbound[1].url, 'https://api.runpod.io/v2/pods'); @@ -976,10 +1109,7 @@ describe('template / network-volume / registry routing under v2', () => { jsonBody: { networkVolumes: [] }, }); await handlers.get('list-network-volumes')!({}); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/network-volumes' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/network-volumes'); }); }); @@ -991,10 +1121,7 @@ describe('template / network-volume / registry routing under v2', () => { size: 50, dataCenterId: 'EU-RO-1', }); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/network-volumes' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/network-volumes'); const body = JSON.parse(outbound[0].body!); assert.equal(body.dataCenter, 'EU-RO-1'); assert.equal('dataCenterId' in body, false); @@ -1019,10 +1146,7 @@ describe('template / network-volume / registry routing under v2', () => { templateId: 't_1', imageName: 'img2', }); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/templates/t_1' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/templates/t_1'); assert.equal(outbound[0].method, 'PATCH'); const body = JSON.parse(outbound[0].body!); assert.equal(body.image, 'img2'); // v2 mapper maps it @@ -1239,10 +1363,7 @@ describe('catalog routing (B5)', () => { const out = (await handlers.get('list-gpu-types')!({ includeAvailability: false, })) as { content: Array<{ text: string }> }; - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/catalog/gpus' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/catalog/gpus'); // no availability data → nothing filtered out assert.equal(JSON.parse(out.content[0].text).items.length, 2); }); @@ -1346,10 +1467,7 @@ describe('catalog routing (B5)', () => { const out = (await handlers.get('list-cpu-types')!({})) as { content: Array<{ text: string }>; }; - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/catalog/cpus' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/catalog/cpus'); assert.deepEqual(JSON.parse(out.content[0].text).items, [ { id: 'cpu5c' }, ]); @@ -1606,89 +1724,45 @@ describe('hosted HTTP transport clamps long-poll budgets', () => { assert.equal(marker.effectiveMs, HTTP_LONG_POLL_BUDGET_MS); }); - // Drive stream-job's poll loop on a fully deterministic clock and return what - // it produced. Two clocks matter: setTimeout (the 1s poll sleep) and Date.now - // (the budget check). setTimeout is mocked via node:test mock timers — Node - // >=20.11 takes `{apis}`, Node 18 takes a bare array, so try both. Date is NOT - // mockable on Node 18 at all, so the budget clock is stubbed by hand on every - // version and advanced in lockstep with the ticks. - // `t` is node:test's own TestContext, NOT a hand-rolled structural type: an - // `enable: (o: unknown) => void` field cannot accept the real signature under - // strictFunctionTypes (a function taking `unknown` is not assignable from one - // taking MockTimersOptions), so that annotation compiled only while tests were - // excluded from type-check. The Node 18 bare-array form is handled by the cast - // in the catch below. async function driveStreamJobToBudget( t: TestContext, opts: { transport: 'stdio' | 'http'; endpointId: string; ticks: number } ) { - const enable = t.mock?.timers?.enable; - if (typeof enable !== 'function') { - throw new Error( - 'node:test mock timers unavailable (needs Node >=18.19) — this test cannot run on this runtime' - ); - } - try { - t.mock.timers.enable({ apis: ['setTimeout'] }); - } catch (err) { - // Node 18 takes a bare array and rejects the object form with a TypeError - // naming the `timers` argument. Anything else is a real failure, not a - // signature mismatch, so don't swallow it behind a second attempt. - if (!(err instanceof TypeError)) throw err; - (t.mock.timers.enable as unknown as (apis: string[]) => void)([ - 'setTimeout', - ]); - } - let fakeNow = 0; - const realNow = Date.now; - Date.now = () => fakeNow; - try { - const { handlers, outbound } = harness({ - transport: opts.transport, - // Never terminal — only the budget can end the loop. - jsonBody: { status: 'IN_PROGRESS', stream: [{ chunk: 'x' }] }, - }); - const pending = handlers.get('stream-job')!({ - endpointId: opts.endpointId, - jobId: 'jH', - }); - // setImmediate is NOT mocked, so each round lets the in-flight poll (fake - // fetch + json()) settle before advancing both clocks by one real poll - // interval — imported, not hardcoded, so halving the interval in the - // source can't double the live request rate with these tests still green. - for (let i = 0; i < opts.ticks; i++) { - await new Promise((r) => setImmediate(r)); - fakeNow += STREAM_JOB_POLL_INTERVAL_MS; - t.mock.timers.tick(STREAM_JOB_POLL_INTERVAL_MS); - } - await new Promise((r) => setImmediate(r)); - // The loop MUST have ended by now. If a regression widened the budget it - // would still be polling, and a bare `await pending` would hang forever — - // node:test has no default timeout, so that stalls this file's remaining - // ~90 tests until CI's job timeout instead of failing here. - const ended = await Promise.race([ - pending.then(() => true), - new Promise((r) => setImmediate(() => r(false))), - ]); - if (!ended) { - throw new Error( - `stream-job was still polling after ${opts.ticks} ticks (${ - (opts.ticks * STREAM_JOB_POLL_INTERVAL_MS) / 1000 - }s of fake time) — its budget is larger than this test expects` - ); - } - const out = (await pending) as { content: Array<{ text: string }> }; - return { - outbound, - result: JSON.parse(out.content[0].text) as { - pollingTimedOut?: boolean; - note?: string; - stream: unknown[]; - }, - }; - } finally { - Date.now = realNow; + const clock = useFakeClock(t); + const { handlers, outbound } = harness({ + transport: opts.transport, + // Never terminal — only the budget can end the loop. + jsonBody: { status: 'IN_PROGRESS', stream: [{ chunk: 'x' }] }, + }); + const pending = handlers.get('stream-job')!({ + endpointId: opts.endpointId, + jobId: 'jH', + }); + // setImmediate is NOT mocked, so each round lets the in-flight poll (fake + // fetch + json()) settle before advancing both clocks by one real poll + // interval — imported, not hardcoded, so halving the interval in the + // source can't double the live request rate with these tests still green. + for (let i = 0; i < opts.ticks; i++) { + await clock.settle(); + clock.advance(STREAM_JOB_POLL_INTERVAL_MS); } + await clock.settle(); + // The loop MUST have ended by now; if a regression widened the budget it + // would still be polling. Hence `settled` rather than a bare await. + const out = (await settled( + pending, + `stream-job after ${opts.ticks} ticks (${ + (opts.ticks * STREAM_JOB_POLL_INTERVAL_MS) / 1000 + }s of fake time; its budget is larger than this test expects)` + )) as { content: Array<{ text: string }> }; + return { + outbound, + result: JSON.parse(out.content[0].text) as { + pollingTimedOut?: boolean; + note?: string; + stream: unknown[]; + }, + }; } it('stream-job on http: stops polling at 45s and returns collected chunks with the resume note', async (t) => { @@ -1753,20 +1827,261 @@ describe('hosted HTTP transport clamps long-poll budgets', () => { const vercel = JSON.parse( readFileSync(new URL('../vercel.json', import.meta.url), 'utf8') ) as { functions?: Record }; - const maxDuration = Object.values(vercel.functions ?? {}).find( - (f) => typeof f.maxDuration === 'number' - )?.maxDuration; + // Keyed to the function both budgets actually run in, not to the first + // entry that happens to declare a limit: adding a second function ahead of + // it would otherwise assert against a number that governs neither. + const maxDuration = vercel.functions?.['api/index.ts']?.maxDuration; assert.ok( typeof maxDuration === 'number', - 'vercel.json no longer declares a maxDuration — the budget below is derived from it' + 'vercel.json no longer declares a maxDuration for api/index.ts — both budgets below are derived from it' ); - // 8s covers the credential pre-flight (4s) and the v2 probe (4s), both of - // which run before the tool does; the rest is serialization + cold start. + // The credential pre-flight (4s) runs before dispatch; the rest is + // serialization and cold start. The v2 probe is stdio-only startup wiring, + // so it is not in this path despite what an earlier version of this + // comment said. const PRE_FLIGHT_HEADROOM_MS = 8_000; assert.ok( HTTP_LONG_POLL_BUDGET_MS + PRE_FLIGHT_HEADROOM_MS <= maxDuration * 1000, `HTTP_LONG_POLL_BUDGET_MS (${HTTP_LONG_POLL_BUDGET_MS}) + pre-flight headroom (${PRE_FLIGHT_HEADROOM_MS}) exceeds maxDuration (${maxDuration}s). Adjust the budget in src/tools/jobs.ts to match.` ); + // The OAuth /token poll runs in the same function under the same limit and + // is the same 45s for the same reason, so it is pinned here rather than + // left as a comment claiming an invariant nothing checks. + assert.ok( + TOKEN_POLL_BUDGET_MS + PRE_FLIGHT_HEADROOM_MS <= maxDuration * 1000, + `TOKEN_POLL_BUDGET_MS (${TOKEN_POLL_BUDGET_MS}) leaves no room under maxDuration (${maxDuration}s) to serialize the OAuth error. Adjust it in api/index.ts.` + ); + }); +}); + +describe('stream-job poll loop', () => { + it('bounds each poll by the hold it brackets, not by the whole budget', async () => { + // The deadline the loop actually asks for, observed rather than inferred. + // The regression this pins let one wedged socket spend the entire budget in + // a single attempt: the retry loop below never ran, and a stall returned + // nothing at all. + const asked: number[] = []; + const { result, chunks } = await collectJobStream({ + budgetMs: HTTP_LONG_POLL_BUDGET_MS, + holdMs: HTTP_STREAM_POLL_WAIT_MS, + poll: async (timeoutMs) => { + asked.push(timeoutMs); + return { status: 'COMPLETED', stream: [{ chunk: 'a' }] }; + }, + }); + assert.deepEqual(asked, [ + HTTP_STREAM_POLL_WAIT_MS + UPSTREAM_HOLD_SLACK_MS, + ]); + assert.ok( + asked[0] < HTTP_LONG_POLL_BUDGET_MS, + 'a poll deadline equal to the budget leaves nothing for a second attempt' + ); + assert.equal(result.status, 'COMPLETED'); + assert.deepEqual(chunks, [{ chunk: 'a' }]); + }); + + it('retries a failing poll on a fresh deadline and gives up after the error cap', async (t) => { + // Each attempt is bounded well inside the budget, so the counter is + // reachable — that is the whole point of not spending the budget on one + // socket. Fails immediately (no stall) so the loop, not the transport, is + // what this measures. + const clock = useFakeClock(t); + const asked: number[] = []; + const pending = collectJobStream({ + budgetMs: STDIO_STREAM_BUDGET_MS, + holdMs: STREAM_UPSTREAM_DEFAULT_WAIT_MS, + poll: (timeoutMs) => { + asked.push(timeoutMs); + return Promise.reject( + new Error('Runpod Serverless API Error: no response after 15000ms') + ); + }, + }); + // One tick per sleep between attempts; the budget is nowhere near spent, + // so only the error cap can end this. + for (let i = 0; i < MAX_CONSECUTIVE_STREAM_ERRORS; i++) { + await clock.settle(); + clock.advance(STREAM_JOB_POLL_INTERVAL_MS); + } + await clock.settle(); + const { result, chunks } = await settled( + pending, + 'the error-cap poll loop' + ); + assert.equal(asked.length, MAX_CONSECUTIVE_STREAM_ERRORS); + for (const deadline of asked) { + assert.equal( + deadline, + STREAM_UPSTREAM_DEFAULT_WAIT_MS + UPSTREAM_HOLD_SLACK_MS + ); + } + assert.match( + String(result.error), + /Polling aborted after 5 consecutive errors/ + ); + // Same resume guidance the budget exit gives: chunks collected so far are + // useless to an agent that is not told it can call stream-job again. + assert.match(String(result.note), /Call stream-job again/); + assert.deepEqual(chunks, []); + }); + + it('picks budget, hold and ?wait= together per transport', () => { + // Bundled on purpose: the query and the hold have to agree, or a poll asks + // the server for one wait while its deadline brackets another. Because the + // budget travels with them, a wrong transport here also breaks the budget + // tests above — which is what makes those tests cover this choice. + const http = streamPollPlan(true); + assert.equal(http.budgetMs, HTTP_LONG_POLL_BUDGET_MS); + assert.equal(http.holdMs, HTTP_STREAM_POLL_WAIT_MS); + assert.equal( + http.query, + `?wait=${http.holdMs}`, + 'the hold the deadline brackets must be the wait actually sent' + ); + + const stdio = streamPollPlan(false); + assert.equal(stdio.budgetMs, STDIO_STREAM_BUDGET_MS); + assert.equal( + stdio.holdMs, + STREAM_UPSTREAM_DEFAULT_WAIT_MS, + 'stdio sends no wait, so it brackets the server default — not the http cap' + ); + assert.equal(stdio.query, ''); + // The regression this exists for: giving stdio the http hold shortens its + // deadline to 6s against a server holding 10s, so every poll aborts a + // reply in flight and the run dies at the error cap in ~30s. + assert.ok( + streamPollTimeoutMs(stdio.budgetMs, stdio.holdMs) > + streamPollTimeoutMs(http.budgetMs, http.holdMs), + 'stdio brackets a longer hold, so its poll deadline must be longer' + ); + }); + + it('reports only a trailing error, not one the job recovered from', async (t) => { + // lastError rides along to the budget exit. Left uncleared, a blip in the + // first second is still reported next to pollingTimedOut five minutes of + // healthy streaming later, and reads as the reason the run stopped. + const clock = useFakeClock(t); + let attempt = 0; + const pending = collectJobStream({ + budgetMs: 4 * STREAM_JOB_POLL_INTERVAL_MS, + holdMs: STREAM_UPSTREAM_DEFAULT_WAIT_MS, + poll: () => { + attempt++; + return attempt === 1 + ? Promise.reject(new Error('one early blip')) + : Promise.resolve({ status: 'IN_PROGRESS', stream: [] }); + }, + }); + for (let i = 0; i < 6; i++) { + await clock.settle(); + clock.advance(STREAM_JOB_POLL_INTERVAL_MS); + } + const { result } = await settled(pending, 'the stale-error poll loop'); + assert.equal(result.pollingTimedOut, true); + assert.equal( + result.lastError, + undefined, + 'an error the job recovered from must not be reported as the trailing one' + ); + }); + + it('counts only CONSECUTIVE failures, so a job that recovers keeps streaming', async (t) => { + // Without the reset, a long stream over a flaky endpoint dies at the fifth + // failure however many good polls sat between them — and every other test + // here still passes, because none of them ever succeeds after a failure. + const clock = useFakeClock(t); + const script: Array<'fail' | 'chunk' | 'done'> = [ + 'fail', + 'fail', + 'fail', + 'chunk', + 'fail', + 'fail', + 'fail', + 'done', + ]; + let attempt = 0; + const pending = collectJobStream({ + budgetMs: STDIO_STREAM_BUDGET_MS, + holdMs: STREAM_UPSTREAM_DEFAULT_WAIT_MS, + poll: () => { + const step = script[attempt++]; + if (step === 'fail') return Promise.reject(new Error('transient 502')); + return Promise.resolve( + step === 'done' + ? { status: 'COMPLETED', stream: [{ chunk: 'last' }] } + : { status: 'IN_PROGRESS', stream: [{ chunk: 'mid' }] } + ); + }, + }); + for (let i = 0; i < script.length; i++) { + await clock.settle(); + clock.advance(STREAM_JOB_POLL_INTERVAL_MS); + } + await clock.settle(); + const { result, chunks } = await settled(pending, 'the recovery poll loop'); + assert.equal( + attempt, + script.length, + `stopped after ${attempt} polls; six failures spread across a recovery must not trip the ${MAX_CONSECUTIVE_STREAM_ERRORS}-error cap` + ); + assert.equal(result.status, 'COMPLETED'); + assert.equal(result.error, undefined); + assert.deepEqual(chunks, [{ chunk: 'mid' }, { chunk: 'last' }]); + }); + + it('a wedged poll aborts on its own deadline, and the run ends with pollingTimedOut + lastError', async (t) => { + // End to end through the real handler and the real deadline: the server + // accepts the connection and never answers, so only the AbortSignal the + // client attaches can end the poll. The ceiling seam shrinks that deadline + // to milliseconds; everything else (retry, budget check, reply shape) is + // production code. + const POLL_DEADLINE_MS = 40; + const clock = useFakeClock(t); + const { handlers, outbound } = harness({ + transport: 'http', + maxTimeoutMs: POLL_DEADLINE_MS, + // Far longer than the deadline: the abort has to be what ends it. + delayMs: 60_000, + }); + const pending = handlers.get('stream-job')!({ + endpointId: 'ep_h', + jobId: 'jW', + }) as Promise<{ content: Array<{ text: string }> }>; + + // Waited for in real time, because AbortSignal.timeout does not run on the + // mocked clock — and waited on the recorded signal, not on `outbound`, + // which is appended when the request goes out rather than when it aborts. + const aborted = (n: number) => () => outbound[n]?.signal?.aborted === true; + await clock.waitFor(aborted(0), 'the first poll to abort'); + // The loop found budget left and parked in its (mocked) sleep. Spend the + // budget while it is parked, then release the sleep: the second poll still + // goes out — proving the loop reconnects rather than waiting a wedged + // socket out — and the check after it ends the run. + clock.jump(HTTP_LONG_POLL_BUDGET_MS); + clock.advance(STREAM_JOB_POLL_INTERVAL_MS); + await clock.waitFor(aborted(1), 'the retry to go out and abort'); + + const out = await pending; + const result = JSON.parse(out.content[0].text) as { + pollingTimedOut?: boolean; + lastError?: string; + stream: unknown[]; + }; + assert.equal( + outbound.length, + 2, + `polled ${outbound.length} times; a wedged socket must be abandoned and retried, not waited out` + ); + assert.equal(result.pollingTimedOut, true); + assert.match( + result.lastError ?? '', + new RegExp(`no response after ${POLL_DEADLINE_MS}ms`), + 'the stall must be reported as this poll’s deadline, not discarded' + ); + assert.match(result.lastError ?? '', /Runpod Serverless API Error/); + assert.deepEqual(result.stream, []); }); }); @@ -1982,10 +2297,7 @@ describe('endpoint routing under RUNPOD_REST_VERSION=v2', () => { endpointId: 'ep_1', includeTemplate: true, }); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/serverless/ep_1' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/serverless/ep_1'); }); }); @@ -2178,10 +2490,7 @@ describe('endpoint routing under RUNPOD_REST_VERSION=v2', () => { }); assert.equal(outbound.length, 2); assert.equal(outbound[0].method, 'GET'); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/serverless/ep_1' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/serverless/ep_1'); assert.equal(outbound[1].method, 'PATCH'); assert.deepEqual(JSON.parse(outbound[1].body!).scaling, { type: 'REQUEST_COUNT', @@ -2250,10 +2559,7 @@ describe('endpoint routing under RUNPOD_REST_VERSION=v2', () => { workersMax: 5, imageName: 'img:3', }); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/serverless/ep_1' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/serverless/ep_1'); assert.equal(outbound[0].method, 'PATCH'); const body = JSON.parse(outbound[0].body!); assert.deepEqual(body.workers, { max: 5 }); @@ -2266,10 +2572,7 @@ describe('endpoint routing under RUNPOD_REST_VERSION=v2', () => { await withV2(async () => { const { handlers, outbound } = harness({ jsonBody: {} }); await handlers.get('delete-endpoint')!({ endpointId: 'ep_1' }); - assert.equal( - outbound[0].url, - 'https://api.runpod.io/v2/serverless/ep_1' - ); + assert.equal(outbound[0].url, 'https://api.runpod.io/v2/serverless/ep_1'); assert.equal(outbound[0].method, 'DELETE'); }); }); @@ -2359,10 +2662,7 @@ describe('log streaming tools (v2-only)', () => { const out = await handlers.get('stream-pod-logs')!({ podId: 'pod_1' }); assert.equal(calls.length, 1); // `source: both` is the default → NO source query param. - assert.equal( - calls[0].url, - 'https://api.runpod.io/v2/pods/pod_1/logs' - ); + assert.equal(calls[0].url, 'https://api.runpod.io/v2/pods/pod_1/logs'); assert.equal(calls[0].maxWaitMs, 5000); assert.equal(calls[0].maxBytes, 256 * 1024); const body = parseText(out); @@ -3046,6 +3346,42 @@ describe('get-job-status — queued-job worker diagnosis', () => { }); }); + // The diagnosis decorates a status we already have, so it is given a short + // deadline of its own (QUEUED_DIAGNOSIS_TIMEOUT_MS) rather than the 30s + // default: spending the caller's remaining budget to enrich a reply that was + // already ready is the wrong trade, and on http it is budget the next call + // needs. Whatever it costs, failing must not cost the status. + it('a failing diagnosis is swallowed — the status the caller asked for still comes back', async () => { + await withV2(async () => { + const { handlers, outbound } = harness({ + steps: [ + { jsonBody: queued }, + { status: 500, text: 'workers listing unavailable' }, + ], + }); + const out = await handlers.get('get-job-status')!({ + endpointId: 'ep', + jobId: 'j1', + }); + assert.equal( + outbound.length, + 2, + 'expected the diagnosis to be attempted' + ); + const payload = parseText(out); + assert.equal(payload.status, 'IN_QUEUE'); + assert.equal(payload.workerHealth, undefined); + assert.equal(payload.hint, undefined); + }); + }); + + it('is bounded well under the ordinary request deadline', () => { + assert.ok( + QUEUED_DIAGNOSIS_TIMEOUT_MS < DEFAULT_REQUEST_TIMEOUT_MS, + `the diagnosis is discardable enrichment; at ${QUEUED_DIAGNOSIS_TIMEOUT_MS}ms it is no cheaper than the ${DEFAULT_REQUEST_TIMEOUT_MS}ms default it exists to undercut` + ); + }); + it('caches the diagnosis briefly: rapid polls reuse it instead of refetching workers', async () => { // Agents poll get-job-status in a loop while queued; without the cache every // poll fired a second workers call for an answer that changes on the order @@ -4082,3 +4418,107 @@ describe('onUnauthorized ignores the unauthenticated GraphQL path', () => { assert.equal(fired, 0, 'a no-credential 401 invalidated the verdict'); }); }); + +// The deadline through the real tool wiring — including the one tool that is +// SUPPOSED to wait a long time and must not be truncated by the default. +describe('per-request deadline', () => { + it('an ordinary tool gives up on a silent server with an actionable error', async () => { + const { handlers } = harness({ defaultTimeoutMs: 20, delayMs: 5_000 }); + await assert.rejects( + () => handlers.get('list-pods')!({}), + (err: unknown) => { + assert.equal((err as Error).name, 'RequestTimeoutError'); + assert.match( + (err as Error).message, + /^Runpod API Error: no response after 20ms / + ); + return true; + } + ); + }); + + it('runsync-endpoint outlasts the default, using the wait it asked the server for', async () => { + // Default squeezed to 20ms against a server that takes 120ms to answer. + // runsync asked for wait=5000, so its own deadline is 15s and the slow + // answer must still land — the default would have thrown the job away. + const { handlers } = harness({ + jsonBody: { id: 'job_1', status: 'COMPLETED' }, + defaultTimeoutMs: 20, + delayMs: 120, + }); + const out = (await handlers.get('runsync-endpoint')!({ + endpointId: 'ep_t', + input: { x: 1 }, + wait: 5000, + })) as { content: { text: string }[] }; + assert.deepEqual(JSON.parse(out.content[0].text), { + id: 'job_1', + status: 'COMPLETED', + }); + }); + + it('runsync-endpoint without an explicit wait gets the same long deadline (the server still waits its own 90s default)', async () => { + const { handlers } = harness({ + jsonBody: { id: 'job_2', status: 'COMPLETED' }, + defaultTimeoutMs: 20, + delayMs: 120, + }); + const out = (await handlers.get('runsync-endpoint')!({ + endpointId: 'ep_t', + input: { x: 1 }, + })) as { content: { text: string }[] }; + assert.deepEqual(JSON.parse(out.content[0].text), { + id: 'job_2', + status: 'COMPLETED', + }); + }); + + // GraphQL is POST on the wire whatever it carries, and the retry advice keys + // off the method. Reporting POST for a catalog read tells the agent to "check + // with the matching list-/get- tool first" — which is the tool that just + // failed. + it('a timed-out GraphQL read is described as safe to retry', async () => { + const { handlers } = harness({ defaultTimeoutMs: 20, delayMs: 5_000 }); + await assert.rejects( + () => handlers.get('list-gpu-types')!({}), + (err: unknown) => { + assert.equal((err as Error).name, 'RequestTimeoutError'); + assert.match((err as Error).message, /for GET /); + assert.match((err as Error).message, /retrying is safe/); + assert.doesNotMatch((err as Error).message, /may have SUCCEEDED/); + return true; + } + ); + }); + + // The write half is the saveEndpoint mutation inside set-endpoint-gpus and + // deploy-hub-repo, and both send a read first — so the classifier is pinned + // directly rather than by stalling a handler's second call. + it('classifies the operation shapes this file actually sends', () => { + assert.equal( + isGraphqlMutation('mutation saveEndpoint($input: EndpointInput!) { id }'), + true + ); + assert.equal( + isGraphqlMutation(` + mutation saveEndpoint($input: EndpointInput!) { + saveEndpoint(input: $input) { id } + } + `), + true, + 'the real call sites are indented template literals' + ); + // A read must not inherit the write warning: anonymous operations and the + // explicit `query` keyword are both reads. + assert.equal( + isGraphqlMutation('query { myself { endpoints { id } } }'), + false + ); + assert.equal(isGraphqlMutation('{ gpuTypes { id displayName } }'), false); + assert.equal( + isGraphqlMutation('# mutation, eventually\nquery { gpuTypes { id } }'), + false, + 'a comment mentioning mutation is not one' + ); + }); +}); diff --git a/tests/http.test.ts b/tests/http.test.ts index 1c2b0c1..b2e13b2 100644 --- a/tests/http.test.ts +++ b/tests/http.test.ts @@ -1,11 +1,35 @@ -import { describe, it } from 'node:test'; +import { describe, it, before, after } from 'node:test'; import assert from 'node:assert/strict'; +import { readFileSync } from 'node:fs'; +import http from 'node:http'; +import nodeFetch from 'node-fetch'; -import { createHttpClient, HttpError } from '../src/_shared/http.js'; +import { + clampTimeout, + createHttpClient, + DEFAULT_REQUEST_TIMEOUT_MS, + HttpError, + MIN_REQUEST_TIMEOUT_MS, + RequestTimeoutError, +} from '../src/_shared/http.js'; import { sanitizeUaToken, buildTrackingHeaders, } from '../src/_shared/tracking.js'; +import { + HTTP_TRANSPORT_BUDGET_MS, + invocationBudget, +} from '../src/tools/runtime.js'; +import { + HTTP_LONG_POLL_BUDGET_MS, + HTTP_STREAM_POLL_WAIT_MS, + MAX_CONSECUTIVE_STREAM_ERRORS, + MIN_STREAM_POLL_TIMEOUT_MS, + STDIO_STREAM_BUDGET_MS, + STREAM_UPSTREAM_DEFAULT_WAIT_MS, + streamPollTimeoutMs, + UPSTREAM_HOLD_SLACK_MS, +} from '../src/tools/jobs.js'; // ---- fake response/fetch builders (no network) ---- interface FakeResponseOpts { @@ -35,13 +59,23 @@ function fakeResponse(opts: FakeResponseOpts) { type Captured = { url?: string; - init?: { method: string; headers: Record; body?: string }; + init?: { + method: string; + headers: Record; + body?: string; + signal?: AbortSignal; + }; }; function fakeFetch(resp: ReturnType, captured?: Captured) { return async ( url: string, - init: { method: string; headers: Record; body?: string } + init: { + method: string; + headers: Record; + body?: string; + signal?: AbortSignal; + } ) => { if (captured) { captured.url = url; @@ -291,6 +325,255 @@ describe('createHttpClient — response handling', () => { }); }); +// Pins the deadline, the per-call override runsync needs, the platform +// ceiling, and the error text the agent actually reads. Why it exists: see the +// REQUEST DEADLINE block in _shared/http.ts. +describe('createHttpClient — request deadline', () => { + // Rejects on the caller's signal, as a real fetch does; a fake that ignored + // it would make the deadline unobservable. + // AbortSignal.timeout's timer is unref'd — fine in production, where a real + // socket holds the loop, but these fakes do no I/O. Without something ref'd + // pending, Node 18 drains the loop and CANCELS the test rather than letting + // the deadline fire. Released as soon as the abort lands. + const KEEP_ALIVE_MS = 500; + const onAbort = (signal?: AbortSignal): Promise => { + const keepAlive = setTimeout(() => {}, KEEP_ALIVE_MS); + const promise = new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => + reject(signal.reason ?? new Error('aborted')) + ); + }); + promise.catch(() => {}).finally(() => clearTimeout(keepAlive)); + return promise; + }; + + // Accepts the connection and then says nothing — the wedged-server case. + const silentFetch = (_url: string, init: { signal?: AbortSignal }) => + onAbort(init.signal); + + // Answers, but slowly — the shape of a legitimately long call (runsync). + const slowFetch = + (ms: number, body: unknown) => + (_url: string, init: { signal?: AbortSignal }) => + Promise.race([ + onAbort(init.signal), + new Promise>((resolve) => + setTimeout(() => resolve(fakeResponse({ jsonBody: body })), ms) + ), + ]); + + const mkClient = (extra: { + fetch: Parameters[0]['fetch']; + defaultTimeoutMs?: number; + maxTimeoutMs?: Parameters[0]['maxTimeoutMs']; + }) => + createHttpClient({ + apiKey: 'k', + tracking: noTracking, + errorPrefix: 'Runpod API Error', + ...extra, + }); + + it('attaches a deadline signal to every request', async () => { + const cap: Captured = {}; + const client = mkClient({ + fetch: fakeFetch(fakeResponse({ jsonBody: {} }), cap), + }); + await client('http://x'); + assert.ok( + cap.init?.signal instanceof AbortSignal, + 'no request may go out unbounded' + ); + assert.equal(cap.init?.signal.aborted, false); + }); + + it('the default deadline is well under the hosted 60s function limit', () => { + assert.ok( + DEFAULT_REQUEST_TIMEOUT_MS < 60_000, + 'a default at or past maxDuration would still surface as a bare 504' + ); + }); + + it('the default fires on a silent server and names itself, the API, and the method', async () => { + const client = mkClient({ fetch: silentFetch, defaultTimeoutMs: 5 }); + await assert.rejects( + () => client('http://x'), + (err: unknown) => { + assert.ok(err instanceof RequestTimeoutError); + assert.equal(err.name, 'RequestTimeoutError'); + assert.equal(err.timeoutMs, 5); + assert.equal(err.method, 'GET'); + assert.match(err.message, /^Runpod API Error: no response after 5ms /); + return true; + } + ); + }); + + // The message is the agent's entire input on a timeout, and we abandoned the + // request without undoing it. On a write it may well have landed, and there + // is no idempotency key upstream to deduplicate a retry — so "retry" is only + // ever safe advice for a read. + it('a timed-out write warns against blind retry; a read says retrying is safe', async () => { + const client = mkClient({ fetch: silentFetch, defaultTimeoutMs: 5 }); + + await assert.rejects( + () => client('http://x', 'POST', { name: 'p' }), + (err: unknown) => { + assert.ok(err instanceof RequestTimeoutError); + assert.equal(err.method, 'POST'); + assert.match(err.message, /may have SUCCEEDED upstream/); + assert.match(err.message, /do not retry blindly/); + return true; + } + ); + + await assert.rejects( + () => client('http://x', 'GET'), + (err: unknown) => { + assert.ok(err instanceof RequestTimeoutError); + assert.match(err.message, /retrying is safe/); + assert.equal( + /SUCCEEDED/.test(err.message), + false, + 'a read must not carry the write warning' + ); + return true; + } + ); + + for (const method of ['PATCH', 'PUT', 'DELETE']) { + await assert.rejects( + () => client('http://x', method, { a: 1 }), + (err: unknown) => + err instanceof RequestTimeoutError && + /do not retry blindly/.test(err.message), + `${method} must be treated as a write` + ); + } + }); + + it('the error prefix follows the client (serverless vs rest)', async () => { + const client = createHttpClient({ + apiKey: 'k', + fetch: silentFetch, + tracking: noTracking, + errorPrefix: 'Runpod Serverless API Error', + defaultTimeoutMs: 5, + }); + await assert.rejects( + () => client('http://x'), + /Runpod Serverless API Error: no response after 5ms / + ); + }); + + it('a per-call timeoutMs shortens the deadline', async () => { + const client = mkClient({ fetch: silentFetch, defaultTimeoutMs: 10_000 }); + await assert.rejects( + () => client('http://x', 'GET', undefined, { timeoutMs: 5 }), + (err: unknown) => + err instanceof RequestTimeoutError && err.timeoutMs === 5 + ); + }); + + it('a per-call timeoutMs lengthens it — the runsync case, where the default would truncate a wait the caller asked for', async () => { + const client = mkClient({ + fetch: slowFetch(60, { id: 'job_1' }), + defaultTimeoutMs: 5, + }); + const out = await client('http://x', 'POST', {}, { timeoutMs: 5_000 }); + assert.deepEqual(out, { id: 'job_1' }); + }); + + it('maxTimeoutMs caps an override (the hosted platform budget)', async () => { + const client = mkClient({ + fetch: silentFetch, + defaultTimeoutMs: 10_000, + maxTimeoutMs: 5, + }); + await assert.rejects( + () => client('http://x', 'GET', undefined, { timeoutMs: 300_000 }), + (err: unknown) => + err instanceof RequestTimeoutError && err.timeoutMs === 5 + ); + }); + + // A per-request deadline bounds one wedged socket. It does NOT bound a handler + // that makes two calls, and several do (get-job-status, deploy-hub-repo, + // update-endpoint) — so a thunk ceiling reports what is left of the whole + // invocation and each request is clamped to the remainder. + // Asserted on clampTimeout rather than through a client: the deadline a + // request was given is otherwise only observable by waiting for it to fire, + // and the floor is a full second by design. + it('a decaying ceiling hands each successive call only the remainder', () => { + let remaining = 50_000; + const ceiling = () => remaining; + assert.equal(clampTimeout(undefined, ceiling, 30_000), 30_000); + remaining = 20_000; // as if that first 30s call had been spent in full + assert.equal( + clampTimeout(undefined, ceiling, 30_000), + 20_000, + 'a second call handed a fresh 30s is exactly the two-call 504 this closes' + ); + }); + + it('an exhausted ceiling floors rather than aborting before it connects', () => { + // A 0ms or negative deadline reports "no response after 0ms", which reads + // as a server fault rather than a budget we had already spent. + assert.equal( + clampTimeout(undefined, () => 0), + MIN_REQUEST_TIMEOUT_MS + ); + assert.equal( + clampTimeout(undefined, () => -5_000), + MIN_REQUEST_TIMEOUT_MS + ); + }); + + it('a static ceiling is left exactly as configured — only the decaying thunk is floored', () => { + // The floor exists for a budget that ran out on its own. A number someone + // wrote down is not that, and silently raising it would be the surprise + // (this suite's own 5ms ceilings depend on it). + assert.equal(clampTimeout(300_000, 5), 5); + assert.equal( + clampTimeout(300_000, () => 5), + MIN_REQUEST_TIMEOUT_MS + ); + }); + + it('the client re-reads the ceiling on every request instead of capturing it once', async () => { + // The counterpart to the arithmetic above: a ceiling resolved once at + // construction would let every request in an invocation claim the full + // allowance no matter what earlier ones had already spent. + let reads = 0; + const client = mkClient({ + fetch: fakeFetch(fakeResponse({ jsonBody: {} })), + maxTimeoutMs: () => { + reads += 1; + return 30_000; + }, + }); + await client('http://x'); + await client('http://x'); + assert.equal(reads, 2); + }); + + it('a network failure is not relabelled as a timeout', async () => { + const client = mkClient({ + fetch: async () => { + throw new Error('connect ECONNREFUSED'); + }, + }); + await assert.rejects( + () => client('http://x'), + (err: unknown) => { + assert.ok(!(err instanceof RequestTimeoutError)); + assert.match((err as Error).message, /ECONNREFUSED/); + return true; + } + ); + }); +}); + describe('tracking headers', () => { it('sanitizeUaToken strips reserved chars and bounds length', () => { assert.equal(sanitizeUaToken('claude (code)'), 'claude__code_'); @@ -455,3 +738,232 @@ describe('createHttpClient — 429 hint wiring', () => { ); }); }); + +// The backstop is only correct relative to two numbers in two other files. It +// must stay under the platform limit (or it never fires) and above any +// server-side wait a tool asks for (or it aborts a reply that was in flight). +describe('the http request backstop sits between the waits it brackets', () => { + const maxDuration = ( + JSON.parse( + readFileSync(new URL('../vercel.json', import.meta.url), 'utf8') + ) as { functions?: Record } + ).functions?.['api/index.ts']?.maxDuration; + + it('leaves the platform room for the pre-flight that runs before any tool request', () => { + assert.ok( + typeof maxDuration === 'number', + 'vercel.json no longer declares a maxDuration for api/index.ts' + ); + // The platform clock starts before ours: src/http.ts awaits the credential + // pre-flight (4s, credential-check.ts) before dispatching the tool, and it + // burns that full budget in the same outage that stalls the tool call. The + // remainder covers cold start and serializing the error. Measuring from the + // request instead of the invocation is what put this at 55s originally. + const PRE_FLIGHT_MS = 4_000; + const SERIALIZE_AND_COLD_START_MS = 5_000; + assert.ok( + HTTP_TRANSPORT_BUDGET_MS + PRE_FLIGHT_MS + SERIALIZE_AND_COLD_START_MS <= + maxDuration * 1000, + `HTTP_TRANSPORT_BUDGET_MS (${HTTP_TRANSPORT_BUDGET_MS}) + pre-flight (${PRE_FLIGHT_MS}) + serialization (${SERIALIZE_AND_COLD_START_MS}) exceeds maxDuration (${maxDuration}s) — the backstop would fire after the platform already reaped the function` + ); + }); + + it('the budget window is per construction, so each request starts it over', async () => { + // The whole scheme rests on the runtime being rebuilt per request, which on + // http it is (src/http.ts builds a single-use server + registerTools inside + // the handler). Pinned here because hoisting that out — a tempting + // cold-start optimization — would leave every request after the first + // running on the floor. + const first = invocationBudget(1_000); + const before = first(); + await new Promise((resolve) => setTimeout(resolve, 25)); + const after = first(); + assert.ok( + after < before, + 'the window must decay as the invocation is spent' + ); + + const second = invocationBudget(1_000); + assert.ok( + second() > after, + 'a freshly constructed runtime must get a fresh allowance, not the previous request’s remainder' + ); + }); + + it('leaves a second call something to spend, since several handlers make two', () => { + // get-job-status, deploy-hub-repo, update-endpoint and set-endpoint-gpus all + // issue two requests. The shared budget is what keeps their worst case under + // the platform limit, but that only helps if the ceiling is bigger than one + // full default — otherwise the first call takes everything and the second is + // left on the floor, which is a timeout dressed up as a diagnosis. + assert.ok( + HTTP_TRANSPORT_BUDGET_MS > DEFAULT_REQUEST_TIMEOUT_MS, + `HTTP_TRANSPORT_BUDGET_MS (${HTTP_TRANSPORT_BUDGET_MS}) is not above one ${DEFAULT_REQUEST_TIMEOUT_MS}ms default, so a two-call handler's second request starts already exhausted` + ); + }); + + it('clears the longest server-side wait by the slack runsync actually adds', () => { + // runsync asks the Serverless API to hold the connection open for + // HTTP_LONG_POLL_BUDGET_MS and sets its deadline that much plus + // UPSTREAM_HOLD_SLACK_MS. Assert against both real constants: restating + // a smaller literal here would let the ceiling silently clamp the slack + // away — the deadline would land on the hold and abort a reply in flight. + assert.ok( + HTTP_TRANSPORT_BUDGET_MS >= + HTTP_LONG_POLL_BUDGET_MS + UPSTREAM_HOLD_SLACK_MS, + `HTTP_TRANSPORT_BUDGET_MS (${HTTP_TRANSPORT_BUDGET_MS}) is below the ${HTTP_LONG_POLL_BUDGET_MS}ms hold plus its ${UPSTREAM_HOLD_SLACK_MS}ms slack, so the clamp eats the slack` + ); + }); +}); + +describe('stream-job poll deadline', () => { + // The transports differ only in the hold each poll brackets: http caps it with + // ?wait=, stdio takes the server's default. Both are read from the source, so + // moving either constant moves these assertions with it. + const CASES = [ + { + transport: 'http', + budgetMs: HTTP_LONG_POLL_BUDGET_MS, + holdMs: HTTP_STREAM_POLL_WAIT_MS, + }, + { + transport: 'stdio', + budgetMs: STDIO_STREAM_BUDGET_MS, + holdMs: STREAM_UPSTREAM_DEFAULT_WAIT_MS, + }, + ] as const; + + for (const { transport, budgetMs, holdMs } of CASES) { + it(`on ${transport}: clears the hold it brackets, so a reply in flight is never aborted`, () => { + // A deadline at or below the hold cannot be met: the server does not + // answer until its wait elapses, and the reply still needs a round trip. + const deadline = streamPollTimeoutMs(budgetMs, holdMs); + assert.ok( + deadline > holdMs, + `poll deadline ${deadline}ms does not clear the ${holdMs}ms hold, so every slow poll aborts a response already on its way` + ); + }); + + it(`on ${transport}: leaves room for the retry loop instead of spending the budget in one attempt`, () => { + // The regression this pins: a deadline set TO the remaining budget means + // one wedged socket consumes the whole run, MAX_CONSECUTIVE_STREAM_ERRORS + // never engages, and a stall that the loop was built to survive returns + // nothing. Reconnecting is the only recovery, so the budget has to fit + // several attempts. + const deadline = streamPollTimeoutMs(budgetMs, holdMs); + assert.ok( + deadline < budgetMs, + `poll deadline ${deadline}ms is the entire ${budgetMs}ms budget — a single wedged poll ends the run` + ); + const attempts = Math.floor(budgetMs / deadline); + assert.ok( + attempts >= MAX_CONSECUTIVE_STREAM_ERRORS, + `a wedged socket allows only ${attempts} attempts inside the ${budgetMs}ms budget, below the ${MAX_CONSECUTIVE_STREAM_ERRORS} the error counter needs to ever fire` + ); + }); + } + + it('never outlives what is left of the budget', () => { + // Between the hold and the budget the budget wins: overshooting it just + // hands the platform reaper the timeout we were trying to report. + // Above the floor (which is what a nearly-spent budget hits) and below the + // hold-derived deadline, so the budget is the only thing that can be + // binding here. + const remaining = MIN_STREAM_POLL_TIMEOUT_MS + 1; + assert.ok(remaining < HTTP_STREAM_POLL_WAIT_MS + UPSTREAM_HOLD_SLACK_MS); + assert.equal( + streamPollTimeoutMs(remaining, HTTP_STREAM_POLL_WAIT_MS), + remaining + ); + }); + + it('floors a spent budget rather than asking for 0ms', () => { + // A 0ms (or negative) deadline aborts before the socket opens and reports + // "no response after 0ms", which reads as a server fault rather than time + // we had already spent. The budget check after the poll ends the loop. + for (const remaining of [0, -5_000]) { + assert.equal( + streamPollTimeoutMs(remaining, HTTP_STREAM_POLL_WAIT_MS), + MIN_STREAM_POLL_TIMEOUT_MS + ); + } + }); +}); + +// ============== real node-fetch, real socket ============== +// Everything above injects a fake fetch, and `defaultFetch = fetch as HttpFetch` +// (src/tools/runtime.ts) is a CAST — nothing type-checks the init object against +// node-fetch's own RequestInit. If node-fetch ignored the `signal` field, every +// deadline in this release would be inert and every test above would still pass. +// So this one drives the real client, over a real socket, against a server that +// answers and one that never does. +describe('createHttpClient against real node-fetch', () => { + let server: http.Server; + let baseUrl: string; + let stall = false; + const stalled: http.ServerResponse[] = []; + + before(async () => { + server = http.createServer((_req, res) => { + if (stall) { + // Accept, send nothing. The shape node-fetch has no answer for. + stalled.push(res); + return; + } + res.setHeader('content-type', 'application/json'); + res.end(JSON.stringify({ ok: true })); + }); + await new Promise((resolve) => + server.listen(0, '127.0.0.1', () => resolve()) + ); + const address = server.address(); + assert.ok(address && typeof address === 'object'); + baseUrl = `http://127.0.0.1:${address.port}`; + }); + + after(async () => { + for (const res of stalled) res.destroy(); + await new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())) + ); + }); + + const client = () => + createHttpClient({ + apiKey: 'rpa_test', + // The production default, not a fake: this is the whole point. + fetch: nodeFetch as unknown as Parameters< + typeof createHttpClient + >[0]['fetch'], + tracking: () => ({}), + errorPrefix: 'Runpod API Error', + }); + + it('passes a normal response through untouched', async () => { + stall = false; + assert.deepEqual(await client()(`${baseUrl}/ok`), { ok: true }); + }); + + // Bounded: if node-fetch ever stops honoring the signal this hangs rather + // than fails, and node:test has no default timeout. + it( + 'aborts a wedged socket and names it, rather than hanging', + { timeout: 5_000 }, + async () => { + stall = true; + const startedAt = Date.now(); + await assert.rejects( + client()(`${baseUrl}/wedged`, 'GET', undefined, { timeoutMs: 150 }), + (error: unknown) => { + assert.ok(error instanceof RequestTimeoutError); + assert.match(error.message, /no response after 150ms for GET/); + return true; + } + ); + assert.ok( + Date.now() - startedAt < 5_000, + 'node-fetch did not honor the signal — the deadline is inert in production' + ); + } + ); +}); diff --git a/tests/install-clients.test.ts b/tests/install-clients.test.ts index 059e47d..2d792bf 100644 --- a/tests/install-clients.test.ts +++ b/tests/install-clients.test.ts @@ -3,7 +3,9 @@ import assert from 'node:assert/strict'; import * as fs from 'fs'; import * as os from 'os'; import * as path from 'path'; +import http from 'node:http'; +import { verifyApiKey } from '../src/install/verify-key.js'; import { claudeCandidatePaths, pickClaudeBinary, @@ -1761,3 +1763,47 @@ describe('config edits land where the client actually reads', () => { assert.match(result.message ?? '', /not an absolute path/); }); }); + +// ============== install wizard: key verification deadline ============== +// The wizard's own fetch, which does not go through createHttpClient and so +// needed its own deadline. Interactive code with no other test surface: the +// failure mode is a spinner that never resolves and no way out but ^C. +describe('verifyApiKey deadline', () => { + // Bounded: without the deadline this hangs rather than fails, and node:test + // has no default timeout. + it( + 'reports a stalled host as "check failed" instead of hanging', + { timeout: 5_000 }, + async (t) => { + const stalled: http.ServerResponse[] = []; + const server = http.createServer((_req, res) => { + // Accept and go quiet — the shape node-fetch has no answer for. + stalled.push(res); + }); + await new Promise((resolve) => + server.listen(0, '127.0.0.1', () => resolve()) + ); + const address = server.address(); + assert.ok(address && typeof address === 'object'); + const previous = process.env.RUNPOD_REST_API_URL; + process.env.RUNPOD_REST_API_URL = `http://127.0.0.1:${address.port}/v1`; + t.after(async () => { + if (previous === undefined) delete process.env.RUNPOD_REST_API_URL; + else process.env.RUNPOD_REST_API_URL = previous; + for (const res of stalled) res.destroy(); + await new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())) + ); + }); + + const startedAt = Date.now(); + // null is the wizard's "could not check" verdict — it warns and lets the + // user continue, which is right for a host that never answered. + assert.equal(await verifyApiKey('rpa_test', 100), null); + assert.ok( + Date.now() - startedAt < 5_000, + 'the deadline did not end the stalled verification' + ); + } + ); +}); diff --git a/tests/oauth.test.ts b/tests/oauth.test.ts index a5f37d8..4f1723a 100644 --- a/tests/oauth.test.ts +++ b/tests/oauth.test.ts @@ -2,7 +2,10 @@ import assert from 'node:assert/strict'; import http from 'node:http'; import { after, before, beforeEach, describe, it } from 'node:test'; -import handler from '../api/index.js'; +import handler, { + DEFAULT_FLASH_GRAPHQL_TIMEOUT_MS, + getFlashTimeoutMs, +} from '../api/index.js'; type JsonObject = Record; @@ -45,10 +48,16 @@ const LOOPBACK_REDIRECT = 'http://127.0.0.1:8765/callback'; const originalGraphqlUrl = process.env.RUNPOD_GRAPHQL_URL; const originalConsoleBaseUrl = process.env.CONSOLE_BASE_URL; const originalApiKeyName = process.env.RUNPOD_API_KEY_NAME; +const originalFlashTimeout = process.env.MCP_FLASH_TIMEOUT_MS; let backend: http.Server; let backendRequests: string[] = []; let flashStatus: JsonObject; +// When set, the backend accepts the request and never answers — the wedged +// socket this file's timeout test exists for. Held so they can be destroyed in +// teardown; an open response would otherwise keep server.close() waiting. +let backendStalls = false; +const stalledResponses: http.ServerResponse[] = []; before(async () => { backend = http.createServer((req, res) => { @@ -59,6 +68,10 @@ before(async () => { query: string; }; backendRequests.push(payload.query); + if (backendStalls) { + stalledResponses.push(res); + return; + } res.setHeader('content-type', 'application/json'); if (payload.query.includes('createFlashAuthRequest')) { res.end( @@ -82,6 +95,7 @@ before(async () => { beforeEach(() => { backendRequests = []; + backendStalls = false; flashStatus = { id: 'code-1', status: 'APPROVED', @@ -92,12 +106,18 @@ beforeEach(() => { }); after(async () => { + for (const res of stalledResponses) res.destroy(); + stalledResponses.length = 0; + if (originalGraphqlUrl === undefined) delete process.env.RUNPOD_GRAPHQL_URL; else process.env.RUNPOD_GRAPHQL_URL = originalGraphqlUrl; if (originalConsoleBaseUrl === undefined) delete process.env.CONSOLE_BASE_URL; else process.env.CONSOLE_BASE_URL = originalConsoleBaseUrl; if (originalApiKeyName === undefined) delete process.env.RUNPOD_API_KEY_NAME; else process.env.RUNPOD_API_KEY_NAME = originalApiKeyName; + if (originalFlashTimeout === undefined) + delete process.env.MCP_FLASH_TIMEOUT_MS; + else process.env.MCP_FLASH_TIMEOUT_MS = originalFlashTimeout; await new Promise((resolve, reject) => backend.close((error) => (error ? reject(error) : resolve())) @@ -225,4 +245,81 @@ describe('OAuth PKCE endpoints', () => { token_type: 'Bearer', }); }); + + // Bounded: without the deadline this hangs rather than fails, and node:test + // has no default timeout — one regression would stall the whole file. + it( + 'answers with a named error when the flash backend accepts and goes silent', + { timeout: 10_000 }, + async () => { + // node-fetch applies no timeout of its own, so before this deadline the + // poll below sat on a wedged socket until Vercel reaped the function at + // maxDuration: a blank 504 on the one flow with no credential yet, so the + // user cannot even retry into a working state. + backendStalls = true; + process.env.MCP_FLASH_TIMEOUT_MS = '150'; + try { + const startedAt = Date.now(); + const res = await token(VERIFIER); + const elapsed = Date.now() - startedAt; + + assert.equal(res.statusCode, 500); + assert.equal(res.body?.error, 'server_error'); + // Names the operation, the host and the deadline — a bare AbortError + // names none of them, and this string is what the client shows. + assert.match( + String(res.body?.error_description), + /flashAuthRequestStatus got no response from http:\/\/127\.0\.0\.1:\d+\/graphql after 150ms/ + ); + assert.ok( + elapsed < 5_000, + `took ${elapsed}ms — the deadline did not end the wedged poll` + ); + // One attempt: a read that may have consumed the code upstream is not + // silently retried. + assert.equal(backendRequests.length, 1); + } finally { + delete process.env.MCP_FLASH_TIMEOUT_MS; + } + } + ); + + it('ignores a MCP_FLASH_TIMEOUT_MS that is not a positive number', () => { + // CLAUDE.md promises a typo cannot disable the deadline, which is only + // true while every non-positive parse falls back to the default. + try { + // Fractional and out-of-range values are the dangerous ones: + // AbortSignal.timeout takes a uint32, throws ERR_OUT_OF_RANGE on a + // fraction (500ing BOTH OAuth routes), and above int32 fires + // immediately instead — turning the dial up would turn it off. + const bads = [ + 'abc', + '0', + '-1', + '', + ' ', + 'NaN', + 'Infinity', + '10000.5', + '5000000000', + '3000000000', + '9007199254740993', + ]; + for (const bad of bads) { + process.env.MCP_FLASH_TIMEOUT_MS = bad; + assert.equal( + getFlashTimeoutMs(), + DEFAULT_FLASH_GRAPHQL_TIMEOUT_MS, + `MCP_FLASH_TIMEOUT_MS=${JSON.stringify(bad)} must fall back to the default` + ); + } + process.env.MCP_FLASH_TIMEOUT_MS = '250'; + assert.equal(getFlashTimeoutMs(), 250, 'a positive integer is honored'); + // What the accepted values are actually handed to. A value this rejects + // would throw here instead, synchronously, on every OAuth request. + assert.doesNotThrow(() => AbortSignal.timeout(getFlashTimeoutMs())); + } finally { + delete process.env.MCP_FLASH_TIMEOUT_MS; + } + }); });