@wavespeed/mcp 1.0.7 → 1.0.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -33,13 +33,13 @@ Built on the same design as the open-source [`@wavespeed/cli`](https://github.co
33
33
  |---|---|
34
34
  | `list_models` | Search the live catalog by text or modality type |
35
35
  | `get_model_schema` | A model's real input schema (required, properties, defaults) |
36
- | `run_model` | Submit + wait; `@path` inputs upload automatically; returns output URLs |
36
+ | `run_model` | Submit + wait; rejects input keys the model doesn't accept; `@path` inputs upload automatically; returns output URLs |
37
37
  | `get_price` | Cost estimate with `unpriced_inputs` / `at_base_price` disclosure |
38
38
  | `get_balance` | Account credit balance |
39
39
  | `upload_file` | Local file → hosted URL (24h content-hash dedupe) |
40
- | `get_prediction` | Recover status/outputs of any run by id |
40
+ | `get_prediction` | Status/outputs of any run by id; optional `wait_seconds` to wait for it |
41
41
 
42
- If `run_model` hits its wait limit the task keeps running server-side — the error names the prediction id, and `get_prediction` picks it up.
42
+ Some generations take minutes to hours. `run_model` waits up to `wait_seconds` (default 50, max 1800) and then returns `done: false` with the prediction id — the task keeps running server-side, and `get_prediction` picks it up. While waiting, both tools send MCP progress notifications that carry the prediction id.
43
43
 
44
44
  ## Setup
45
45
 
package/dist/lib/api.js CHANGED
@@ -14,10 +14,11 @@ const { version: PKG_VERSION } = require('../../package.json');
14
14
  // following the wavespeed-desktop convention. The WAVESPEED_CLIENT_NAME
15
15
  // environment variable overrides the name so wrapper channels can brand
16
16
  // themselves without code changes.
17
+ export const MCP_CLIENT_NAME = 'wavespeed-mcp';
17
18
  export function clientAttributionHeaders() {
18
19
  const platform = os.platform();
19
20
  return {
20
- 'X-Client-Name': process.env.WAVESPEED_CLIENT_NAME || 'wavespeed-mcp',
21
+ 'X-Client-Name': process.env.WAVESPEED_CLIENT_NAME || MCP_CLIENT_NAME,
21
22
  'X-Client-Version': PKG_VERSION,
22
23
  'X-Client-OS': platform === 'win32' ? 'windows' : platform,
23
24
  };
@@ -33,7 +34,7 @@ function readCache(baseUrl) {
33
34
  return null;
34
35
  if (Date.now() - parsed.fetched_at > CACHE_TTL_MS)
35
36
  return null;
36
- return parsed.models;
37
+ return parsed;
37
38
  }
38
39
  catch {
39
40
  return null;
@@ -63,9 +64,11 @@ async function httpError(res, method, apiPath) {
63
64
  catch {
64
65
  /* non-JSON body — fall back below */
65
66
  }
66
- if (detail)
67
- return new Error(detail);
68
- return new Error(`${method} ${apiPath} failed: ${res.status} ${res.statusText}`);
67
+ const err = detail
68
+ ? new Error(detail)
69
+ : new Error(`${method} ${apiPath} failed: ${res.status} ${res.statusText}`);
70
+ err.status = res.status;
71
+ return err;
69
72
  }
70
73
  function authHeaders() {
71
74
  const apiKey = getApiKey();
@@ -74,8 +77,14 @@ function authHeaders() {
74
77
  }
75
78
  return { Authorization: `Bearer ${apiKey}`, ...clientAttributionHeaders() };
76
79
  }
77
- async function apiGet(apiPath) {
78
- const res = await fetch(`${getBaseUrl()}${apiPath}`, { headers: authHeaders() });
80
+ // A hung socket would otherwise stall the tool call until the client gives
81
+ // up. 30s is far above a healthy read, so it only trips on a dead connection.
82
+ const GET_TIMEOUT_MS = 30_000;
83
+ async function apiGetOnce(apiPath) {
84
+ const res = await fetch(`${getBaseUrl()}${apiPath}`, {
85
+ headers: authHeaders(),
86
+ signal: AbortSignal.timeout(GET_TIMEOUT_MS),
87
+ });
79
88
  if (!res.ok)
80
89
  throw await httpError(res, 'GET', apiPath);
81
90
  const json = (await res.json());
@@ -83,6 +92,33 @@ async function apiGet(apiPath) {
83
92
  throw new Error(json.message || `API returned code ${json.code}`);
84
93
  return json.data;
85
94
  }
95
+ // Same set as wavespeed-python. An error with no HTTP status is a network
96
+ // failure or timeout; envelope/JSON errors are not transient.
97
+ const RETRYABLE_STATUS = new Set([429, 500, 502, 503, 504]);
98
+ const MAX_GET_RETRIES = 5;
99
+ function isTransient(err) {
100
+ if (typeof err?.status === 'number')
101
+ return RETRYABLE_STATUS.has(err.status);
102
+ return err instanceof TypeError || err?.name === 'TimeoutError' || err?.name === 'AbortError';
103
+ }
104
+ // Read-only GETs only — submit (POST) is deliberately never retried.
105
+ async function apiGet(apiPath, opts = {}) {
106
+ const delay = opts.retryDelayMs ?? 1000;
107
+ for (let attempt = 0;; attempt++) {
108
+ try {
109
+ return await apiGetOnce(apiPath);
110
+ }
111
+ catch (err) {
112
+ const reason = err?.message ?? String(err);
113
+ if (!isTransient(err) && !opts.describe)
114
+ throw err;
115
+ if (!isTransient(err) || attempt >= MAX_GET_RETRIES) {
116
+ throw new Error(opts.describe ? opts.describe(reason, attempt + 1) : reason);
117
+ }
118
+ await new Promise((r) => setTimeout(r, delay * (attempt + 1)));
119
+ }
120
+ }
121
+ }
86
122
  async function apiPost(apiPath, body) {
87
123
  const res = await fetch(`${getBaseUrl()}${apiPath}`, {
88
124
  method: 'POST',
@@ -96,15 +132,27 @@ async function apiPost(apiPath, body) {
96
132
  throw new Error(json.message || `API returned code ${json.code}`);
97
133
  return json.data;
98
134
  }
135
+ // The server is a long-lived process, so keep the parsed catalog in memory
136
+ // instead of re-reading the ~2MB disk cache on every tool call. Same 1h TTL;
137
+ // the disk cache still lets a fresh process start without a network fetch.
138
+ let memCache = null;
99
139
  export async function fetchModels(opts = {}) {
100
140
  const baseUrl = getBaseUrl();
101
141
  if (!opts.refresh) {
142
+ if (memCache &&
143
+ memCache.baseUrl === baseUrl &&
144
+ Date.now() - memCache.fetchedAt <= CACHE_TTL_MS) {
145
+ return memCache.models;
146
+ }
102
147
  const cached = readCache(baseUrl);
103
- if (cached)
104
- return cached;
148
+ if (cached) {
149
+ memCache = { baseUrl, fetchedAt: cached.fetched_at, models: cached.models };
150
+ return cached.models;
151
+ }
105
152
  }
106
153
  const models = await apiGet('/api/v3/models');
107
154
  writeCache(baseUrl, models);
155
+ memCache = { baseUrl, fetchedAt: Date.now(), models };
108
156
  return models;
109
157
  }
110
158
  export async function fetchBalance() {
@@ -122,24 +170,38 @@ export async function fetchPricing(modelId, inputs) {
122
170
  export async function submitPrediction(model, input) {
123
171
  return apiPost(`/api/v3/${model}`, { ...input });
124
172
  }
125
- export async function fetchPrediction(id) {
126
- return apiGet(`/api/v3/predictions/${id}/result`);
173
+ // The message always carries the task id and says the task was not declared
174
+ // failed — a status read that dies after submit must not look like a lost
175
+ // (or failed) paid task.
176
+ export async function fetchPrediction(id, opts = {}) {
177
+ return apiGet(`/api/v3/predictions/${id}/result`, {
178
+ retryDelayMs: opts.retryDelayMs,
179
+ describe: (reason, attempts) => `Could not read the status of task ${id} after ${attempts} attempt${attempts > 1 ? 's' : ''}: ${reason}. ` +
180
+ `The task was created and may still be running — call get_prediction with id "${id}" to check it (task_id: ${id}).`,
181
+ });
127
182
  }
128
- /** Poll a prediction until it reaches a terminal status. */
183
+ export const TERMINAL_FAILURES = ['failed', 'cancelled', 'timeout', 'deleted'];
184
+ /**
185
+ * Poll a prediction until it reaches a terminal status (done=true, whether
186
+ * completed or failed — the caller decides what a failure means) or the wait
187
+ * budget runs out. Running out is NOT an error: generations can take hours,
188
+ * so the caller gets the still-running prediction back (done=false) and
189
+ * reports its id instead of losing it in a thrown message. `onTick` fires
190
+ * after every non-terminal poll. Transient status-read failures are retried.
191
+ */
129
192
  export async function waitForPrediction(id, opts = {}) {
130
193
  const interval = opts.intervalMs ?? 1000;
131
- const deadline = opts.timeoutMs ? Date.now() + opts.timeoutMs : undefined;
194
+ const deadline = opts.timeoutMs !== undefined ? Date.now() + opts.timeoutMs : undefined;
132
195
  for (;;) {
133
- const item = await fetchPrediction(id);
134
- if (item.status === 'completed')
135
- return item;
136
- if (item.status === 'failed' || item.status === 'cancelled' || item.status === 'timeout' || item.status === 'deleted') {
137
- throw new Error(`Prediction ${item.status}${item.error ? `: ${item.error}` : ''} (task_id: ${id})`);
138
- }
139
- if (deadline && Date.now() > deadline) {
140
- throw new Error(`Still ${item.status} after the wait limit (task_id: ${id}). ` +
141
- `The task keeps running server-side — check it later with get_prediction.`);
196
+ const item = await fetchPrediction(id, { retryDelayMs: opts.retryDelayMs });
197
+ if (item.status === 'completed' || TERMINAL_FAILURES.includes(item.status)) {
198
+ return { item, done: true };
142
199
  }
200
+ await opts.onTick?.(item);
201
+ if (opts.signal?.aborted)
202
+ return { item, done: false };
203
+ if (deadline !== undefined && Date.now() + interval > deadline)
204
+ return { item, done: false };
143
205
  await new Promise((r) => setTimeout(r, interval));
144
206
  }
145
207
  }
@@ -1,19 +1,27 @@
1
1
  // API key + base URL resolution.
2
2
  //
3
3
  // Precedence: WAVESPEED_API_KEY / WAVESPEED_BASE_URL env vars, then the
4
- // wavespeed CLI's stored config (~/.config/wavespeed-nodejs/config.json,
5
- // written by `wavespeed login`). Reading the CLI's store means one login
4
+ // wavespeed CLI's stored config (written by `wavespeed login`; location is
5
+ // per-platform, see cliConfigPath). Reading the CLI's store means one login
6
6
  // covers both tools — the MCP server never implements its own auth flow.
7
7
  import fs from 'node:fs';
8
8
  import path from 'node:path';
9
9
  import os from 'node:os';
10
10
  export const DEFAULT_BASE_URL = 'https://api.wavespeed.ai';
11
- function cliConfigPath() {
12
- // conf@N with projectName "wavespeed" resolves to <config>/wavespeed-nodejs/config.json
13
- const base = process.env.XDG_CONFIG_HOME && process.env.XDG_CONFIG_HOME !== ''
14
- ? process.env.XDG_CONFIG_HOME
15
- : path.join(os.homedir(), '.config');
16
- return path.join(base, 'wavespeed-nodejs', 'config.json');
11
+ // Mirrors conf's env-paths resolution (projectName "wavespeed", suffix
12
+ // "nodejs"), which is per-platform — a Linux-only path would miss every
13
+ // `wavespeed login` done on macOS or Windows.
14
+ export function cliConfigPath(platform = process.platform, env = process.env, home = os.homedir()) {
15
+ const name = 'wavespeed-nodejs';
16
+ if (platform === 'darwin') {
17
+ return path.join(home, 'Library', 'Preferences', name, 'config.json');
18
+ }
19
+ if (platform === 'win32') {
20
+ const appData = env.APPDATA || path.join(home, 'AppData', 'Roaming');
21
+ return path.join(appData, name, 'Config', 'config.json');
22
+ }
23
+ const base = env.XDG_CONFIG_HOME || path.join(home, '.config');
24
+ return path.join(base, name, 'config.json');
17
25
  }
18
26
  function readCliConfig() {
19
27
  try {
@@ -0,0 +1,88 @@
1
+ // Pre-submit input validation against a model's published request schema.
2
+ //
3
+ // The API's entry whitelist silently DROPS any input key that is not in the
4
+ // model's schema — a typo like `size=2560*1440` against a model whose real
5
+ // knobs are `aspect_ratio`/`resolution` doesn't fail, it just quietly bills a
6
+ // generation with defaults. Catching that before submission is the only
7
+ // moment the mistake is still free.
8
+ //
9
+ // Validation fails OPEN: no schema (brand-new model, catalog lag, offline
10
+ // cache miss) means no check. It must never block a submission the API would
11
+ // have accepted.
12
+ export function requestSchemaProps(model) {
13
+ const props = model?.api_schema?.api_schemas?.[0]?.request_schema?.properties;
14
+ if (!props)
15
+ return null;
16
+ const keys = Object.keys(props);
17
+ return keys.length > 0 ? keys : null;
18
+ }
19
+ /**
20
+ * Legacy field names a model still accepts, as `legacy -> canonical`.
21
+ *
22
+ * The API renames public fields without breaking old callers by declaring the
23
+ * old name in the property's `x-aliases` (e.g. `last_image` carries
24
+ * `["end_image"]`) and folding it into the canonical name before validation.
25
+ * A legacy name is therefore a valid input, not a typo — it just is not listed
26
+ * as an accepted input, so schema-driven UIs show only the current name.
27
+ */
28
+ export function requestSchemaAliases(model) {
29
+ const aliases = new Map();
30
+ const props = model?.api_schema?.api_schemas?.[0]?.request_schema?.properties;
31
+ if (!props)
32
+ return aliases;
33
+ for (const [name, def] of Object.entries(props)) {
34
+ const declared = def?.['x-aliases'];
35
+ if (!Array.isArray(declared))
36
+ continue;
37
+ for (const alias of declared) {
38
+ if (typeof alias === 'string' && alias !== name)
39
+ aliases.set(alias, name);
40
+ }
41
+ }
42
+ return aliases;
43
+ }
44
+ // Small edit distance for typo hints — inputs are short snake_case tokens, so
45
+ // a full Levenshtein over ≤40 keys is nothing.
46
+ function editDistance(a, b) {
47
+ const dp = Array.from({ length: a.length + 1 }, (_, i) => [i, ...Array(b.length).fill(0)]);
48
+ for (let j = 0; j <= b.length; j++)
49
+ dp[0][j] = j;
50
+ for (let i = 1; i <= a.length; i++)
51
+ for (let j = 1; j <= b.length; j++)
52
+ dp[i][j] = Math.min(dp[i - 1][j] + 1, dp[i][j - 1] + 1, dp[i - 1][j - 1] + (a[i - 1] === b[j - 1] ? 0 : 1));
53
+ return dp[a.length][b.length];
54
+ }
55
+ function closest(key, known) {
56
+ let best;
57
+ let bestD = Infinity;
58
+ for (const k of known) {
59
+ const d = editDistance(key.toLowerCase(), k.toLowerCase());
60
+ if (d < bestD) {
61
+ bestD = d;
62
+ best = k;
63
+ }
64
+ }
65
+ // A hint is only a hint when it's plausibly the same word.
66
+ return bestD <= Math.max(2, Math.floor(key.length / 3)) ? best : undefined;
67
+ }
68
+ /**
69
+ * Compare top-level input keys with the schema's properties.
70
+ * Returns null when everything checks out or no schema is available.
71
+ */
72
+ export function findUnknownInputs(input, model) {
73
+ const known = requestSchemaProps(model);
74
+ if (!known)
75
+ return null;
76
+ const knownSet = new Set(known);
77
+ const aliases = requestSchemaAliases(model);
78
+ const unknown = Object.keys(input).filter((k) => !knownSet.has(k) && !aliases.has(k));
79
+ if (unknown.length === 0)
80
+ return null;
81
+ const suggestions = new Map();
82
+ for (const k of unknown) {
83
+ const hit = closest(k, known);
84
+ if (hit)
85
+ suggestions.set(k, hit);
86
+ }
87
+ return { unknown, known, suggestions };
88
+ }
package/dist/server.js CHANGED
@@ -15,11 +15,12 @@ import { createRequire } from 'node:module';
15
15
  import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
16
16
  import { z } from 'zod';
17
17
  import { Client } from 'wavespeed';
18
- import { fetchModels, fetchBalance, fetchPricing, fetchPrediction, submitPrediction, waitForPrediction, } from './lib/api.js';
18
+ import { fetchModels, fetchBalance, fetchPricing, submitPrediction, waitForPrediction, MCP_CLIENT_NAME, TERMINAL_FAILURES, } from './lib/api.js';
19
19
  import { getApiKey, getBaseUrl } from './lib/config.js';
20
20
  import { resolveLocalFiles } from './lib/local-files.js';
21
21
  import { uploadWithCache } from './lib/upload-cache.js';
22
22
  import { missingPriceVars, isFloorQuote } from './lib/pricing-vars.js';
23
+ import { findUnknownInputs } from './lib/validate-inputs.js';
23
24
  const require = createRequire(import.meta.url);
24
25
  const { version: PKG_VERSION } = require('../package.json');
25
26
  const PRICE_DISCLAIMER = 'Estimate only, for reference — the amount actually charged for a run is authoritative.';
@@ -39,9 +40,63 @@ async function uploadFile(filePath) {
39
40
  const apiKey = getApiKey();
40
41
  if (!apiKey)
41
42
  throw new Error('No API key configured (WAVESPEED_API_KEY or `wavespeed login`).');
42
- const client = new Client(apiKey, { baseUrl: getBaseUrl() });
43
+ const client = new Client(apiKey, { baseUrl: getBaseUrl(), clientName: MCP_CLIENT_NAME });
43
44
  return uploadWithCache(filePath, (p) => client.upload(p));
44
45
  }
46
+ // Default wait for run_model. Many MCP clients abort a tool call after 60s,
47
+ // and an aborted call never delivers the prediction id — so return before
48
+ // that with the id in hand. Longer waits are opt-in via wait_seconds.
49
+ const DEFAULT_WAIT_SECONDS = 50;
50
+ const POLL_INTERVAL_MS = 2000;
51
+ const STILL_RUNNING_HINT = 'Still running on the server — generations can take minutes to hours. ' +
52
+ 'Call get_prediction with this id (optionally with wait_seconds) to check again.';
53
+ // Emits notifications/progress while polling when the client asked for it.
54
+ // Each one carries the prediction id, and clients that reset their timeout
55
+ // on progress keep a long wait alive.
56
+ function progressTicker(extra, id, startedAt) {
57
+ const token = extra._meta?.progressToken;
58
+ if (token === undefined)
59
+ return undefined;
60
+ let n = 0;
61
+ return async (item) => {
62
+ const secs = Math.floor((Date.now() - startedAt) / 1000);
63
+ await extra
64
+ .sendNotification({
65
+ method: 'notifications/progress',
66
+ params: { progressToken: token, progress: ++n, message: `prediction ${id}: ${item.status} (${secs}s)` },
67
+ })
68
+ .catch(() => {
69
+ /* progress is advisory */
70
+ });
71
+ };
72
+ }
73
+ // Reject inputs the model's schema does not declare. The API silently drops
74
+ // unknown keys and still bills the run, so this is the last free moment to
75
+ // catch an invented parameter. Fails open when no schema is available, and
76
+ // re-fetches the catalog once before rejecting so a stale cache can't block
77
+ // a valid input.
78
+ async function assertKnownInputs(model, input) {
79
+ let report;
80
+ try {
81
+ report = findUnknownInputs(input, (await fetchModels()).find((m) => m.model_id === model));
82
+ if (report) {
83
+ const fresh = await fetchModels({ refresh: true });
84
+ report = findUnknownInputs(input, fresh.find((m) => m.model_id === model));
85
+ }
86
+ }
87
+ catch {
88
+ return; // catalog unreachable — validation is best-effort
89
+ }
90
+ if (!report)
91
+ return;
92
+ const hints = report.unknown.map((k) => {
93
+ const s = report.suggestions.get(k);
94
+ return s ? `${k} (did you mean ${s}?)` : k;
95
+ });
96
+ throw new Error(`Model ${model} does not accept: ${hints.join(', ')}. The API silently drops unknown ` +
97
+ `inputs and still bills the run, so nothing was submitted. Accepted inputs: ` +
98
+ `${report.known.join(', ')}. See get_model_schema.`);
99
+ }
45
100
  function ok(payload) {
46
101
  return { content: [{ type: 'text', text: JSON.stringify(payload, null, 2) }] };
47
102
  }
@@ -53,7 +108,9 @@ export function createServer() {
53
108
  'inputs, run_model to execute. Reference local files as "@./path"',
54
109
  'values inside input — they upload automatically; bare paths are',
55
110
  'passed through untouched and will fail model validation.',
56
- 'Use get_price before expensive runs. Never invent model IDs.',
111
+ 'Use get_price before expensive runs. Never invent model IDs or input keys.',
112
+ 'Long generations (video, some audio/3D) can run for hours: when',
113
+ 'run_model returns done=false, keep the id and poll get_prediction.',
57
114
  ].join(' '),
58
115
  });
59
116
  server.tool('list_models', 'Search the live WaveSpeed model catalog (image, video, audio, 3D). Returns model IDs usable with run_model. Do not invent model IDs — always pick one returned by this tool.', {
@@ -97,7 +154,7 @@ export function createServer() {
97
154
  property_order: schema?.['x-order-properties'] ?? [],
98
155
  });
99
156
  });
100
- server.tool('run_model', 'Run any WaveSpeed model. input keys come from get_model_schema. Local files: pass "@./path" string values — they are uploaded and replaced with hosted URLs (bare paths are NOT uploaded). Returns output URLs. If the wait limit is hit, the task keeps running; recover it with get_prediction.', {
157
+ server.tool('run_model', 'Run any WaveSpeed model. input keys come from get_model_schema; unknown keys are rejected before submission. Local files: pass "@./path" string values — they are uploaded and replaced with hosted URLs (bare paths are NOT uploaded). Returns output URLs when done=true. Some generations take minutes to hours: if the wait ends first, the result has done=false and the prediction id — the task keeps running, so check it with get_prediction.', {
101
158
  model: z.string().describe('Model ID from list_models'),
102
159
  input: z
103
160
  .record(z.unknown())
@@ -107,28 +164,36 @@ export function createServer() {
107
164
  .int()
108
165
  .min(0)
109
166
  .max(1800)
110
- .default(600)
111
- .describe('Max seconds to wait; 0 = submit only and return the prediction id'),
112
- }, { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: true }, async ({ model, input, wait_seconds }) => {
167
+ .default(DEFAULT_WAIT_SECONDS)
168
+ .describe(`Max seconds to wait before returning the still-running prediction id (default ${DEFAULT_WAIT_SECONDS}); 0 = submit only`),
169
+ }, { readOnlyHint: false, destructiveHint: false, idempotentHint: false, openWorldHint: true }, async ({ model, input, wait_seconds }, extra) => {
113
170
  const resolved = await resolveLocalFiles(input, {
114
171
  upload: async (p) => (await uploadFile(p)).url,
115
172
  });
173
+ await assertKnownInputs(model, resolved.input);
116
174
  const started = Date.now();
117
175
  const submitted = await submitPrediction(model, resolved.input);
118
176
  if (wait_seconds === 0) {
119
- return ok({ id: submitted.id, status: submitted.status, model });
177
+ return ok({ id: submitted.id, status: submitted.status, model, done: false, next: STILL_RUNNING_HINT });
120
178
  }
121
- const done = await waitForPrediction(submitted.id, {
122
- intervalMs: 2000,
179
+ const { item, done } = await waitForPrediction(submitted.id, {
180
+ intervalMs: POLL_INTERVAL_MS,
123
181
  timeoutMs: wait_seconds * 1000,
182
+ onTick: progressTicker(extra, submitted.id, started),
183
+ signal: extra.signal,
124
184
  });
185
+ if (TERMINAL_FAILURES.includes(item.status)) {
186
+ throw new Error(`Prediction ${item.status}${item.error ? `: ${item.error}` : ''} (task_id: ${submitted.id})`);
187
+ }
125
188
  return ok({
126
189
  id: submitted.id,
127
190
  model,
128
- status: done.status,
129
- outputs: done.outputs ?? [],
191
+ status: item.status,
192
+ done,
193
+ outputs: item.outputs ?? [],
130
194
  elapsed_ms: Date.now() - started,
131
195
  uploaded_files: resolved.uploaded,
196
+ ...(done ? {} : { next: STILL_RUNNING_HINT }),
132
197
  });
133
198
  });
134
199
  server.tool('get_price', 'Estimate the cost of a run before executing it (no charge). Provide the same input you would pass to run_model — pricing often depends on inputs like duration or resolution.', {
@@ -172,15 +237,32 @@ export function createServer() {
172
237
  server.tool('upload_file', 'Upload a local file to WaveSpeed and get its hosted URL (identical bytes reuse the same upload for 24h). Usually unnecessary — run_model handles "@./path" inputs itself.', { path: z.string().describe('Local file path') }, { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: true }, async ({ path: filePath }) => {
173
238
  return ok(await uploadFile(filePath));
174
239
  });
175
- server.tool('get_prediction', 'Fetch the status and outputs of a past or in-flight prediction by id — use to recover a run that hit the wait limit.', { id: z.string().describe('Prediction id returned by run_model') }, { readOnlyHint: true, openWorldHint: true }, async ({ id }) => {
176
- const item = await fetchPrediction(id);
240
+ server.tool('get_prediction', 'Fetch the status and outputs of a past or in-flight prediction by id — use to pick up a run that returned done=false. Pass wait_seconds to keep waiting for it to finish.', {
241
+ id: z.string().describe('Prediction id returned by run_model'),
242
+ wait_seconds: z
243
+ .number()
244
+ .int()
245
+ .min(0)
246
+ .max(1800)
247
+ .default(0)
248
+ .describe('Max seconds to wait for a terminal status; 0 = return the current status now'),
249
+ }, { readOnlyHint: true, openWorldHint: true }, async ({ id, wait_seconds }, extra) => {
250
+ const started = Date.now();
251
+ const { item, done } = await waitForPrediction(id, {
252
+ intervalMs: POLL_INTERVAL_MS,
253
+ timeoutMs: wait_seconds * 1000,
254
+ onTick: wait_seconds > 0 ? progressTicker(extra, id, started) : undefined,
255
+ signal: extra.signal,
256
+ });
177
257
  return ok({
178
258
  id: item.id,
179
259
  model: item.model,
180
260
  status: item.status,
261
+ done,
181
262
  outputs: item.outputs ?? [],
182
263
  error: item.error,
183
264
  created_at: item.created_at,
265
+ ...(done ? {} : { next: STILL_RUNNING_HINT }),
184
266
  });
185
267
  });
186
268
  return server;
package/package.json CHANGED
@@ -1,7 +1,7 @@
1
1
  {
2
2
  "name": "@wavespeed/mcp",
3
- "version": "1.0.7",
4
- "description": "WaveSpeed MCP server \u2014 run any model on the live WaveSpeed catalog (image, video, audio, 3D) with schema introspection, local-file upload, and price quotes.",
3
+ "version": "1.0.9",
4
+ "description": "WaveSpeed MCP server — run any model on the live WaveSpeed catalog (image, video, audio, 3D) with schema introspection, local-file upload, and price quotes.",
5
5
  "type": "module",
6
6
  "bin": {
7
7
  "wavespeed-mcp": "dist/index.js"
@@ -39,7 +39,7 @@
39
39
  ],
40
40
  "dependencies": {
41
41
  "@modelcontextprotocol/sdk": "^1.12.0",
42
- "wavespeed": "^0.2.6",
42
+ "wavespeed": "^0.3.2",
43
43
  "zod": "^3.24.0"
44
44
  },
45
45
  "devDependencies": {