From 110f47085a1d5d2f4cea54e8cb062df9589ffddf Mon Sep 17 00:00:00 2001 From: Tanish Manku Date: Mon, 24 Aug 2026 21:50:15 -0700 Subject: [PATCH] mcp: job status tools and the fake's jobs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fake-upstream/src/jobs.ts ports infra/fake-storyteller-web's job payloads — the single/batch variant and the richer session-list variant (maybe_failure_category_updated, maybe_failure_message, maybe_batch_token) — with three seeded jobs for the seeded user: finished with a CDN result, running at 40%, and failed on a content rule. Routes: single (404 envelope for unknown), batch (repeated or comma-joined tokens), session listing (session-only, newest first, include/exclude state filters). All validated against the spec. src/mcp/tools/jobs.ts normalises all three payloads into one JobSummary (category, status, is_finished, progress, model, prompt, failure, result with CDN URL) and defines get_job_status, get_jobs_status (≤50 tokens, deduplicated) and list_recent_jobs (state filters as the API's comma-joined queries). The newer failure category is the API's forward-compatible enum whose catch-all is `{ unknown: raw }` — typecheck caught that my type had it as a plain string; the normaliser now keeps the raw value. Tests (10 tool + 6 fake): normaliser cases incl. the catch-all and terminal states; finished/running/404 for a single job; batch deduplication and omission of unknown tokens; session listing order, failure detail, and filter queries. All three registered under read:jobs; the handler suite now lists all seven tools. Claude-Session: https://claude.ai/code/session_01H3gfN2NCx3hBXSrZXYHUQ3 --- mcp/fake-upstream/src/index.ts | 2 + mcp/fake-upstream/src/jobs.ts | 180 +++++++++++++++++++++ mcp/fake-upstream/src/routes/jobs.ts | 70 ++++++++ mcp/fake-upstream/src/state.ts | 4 + mcp/fake-upstream/test/jobs.test.ts | 87 ++++++++++ mcp/src/mcp/server.ts | 4 + mcp/src/mcp/tools/jobs.ts | 230 +++++++++++++++++++++++++++ mcp/test/mcp/handler.test.ts | 3 + mcp/test/mcp/jobs.test.ts | 203 +++++++++++++++++++++++ 9 files changed, 783 insertions(+) create mode 100644 mcp/fake-upstream/src/jobs.ts create mode 100644 mcp/fake-upstream/src/routes/jobs.ts create mode 100644 mcp/fake-upstream/test/jobs.test.ts create mode 100644 mcp/src/mcp/tools/jobs.ts create mode 100644 mcp/test/mcp/jobs.test.ts diff --git a/mcp/fake-upstream/src/index.ts b/mcp/fake-upstream/src/index.ts index 1e76d7103e..05d2f24ff0 100644 --- a/mcp/fake-upstream/src/index.ts +++ b/mcp/fake-upstream/src/index.ts @@ -1,6 +1,7 @@ import { Hono } from "hono"; import { billingRoutes } from "./routes/billing"; +import { jobRoutes } from "./routes/jobs"; import { omniGenRoutes } from "./routes/omni_gen"; import { sessionRoutes } from "./routes/session"; import { createStore, type FakeStore } from "./state"; @@ -20,6 +21,7 @@ export function createFakeUpstream(store: FakeStore = createStore()): Hono { app.route("/", sessionRoutes(store)); app.route("/", billingRoutes(store)); app.route("/", omniGenRoutes()); + app.route("/", jobRoutes(store)); return app; } diff --git a/mcp/fake-upstream/src/jobs.ts b/mcp/fake-upstream/src/jobs.ts new file mode 100644 index 0000000000..fd5e245b53 --- /dev/null +++ b/mcp/fake-upstream/src/jobs.ts @@ -0,0 +1,180 @@ +/** + * Inference jobs — ported from `infra/fake-storyteller-web` (wire/jobs.ts, state/entities.ts). + * The real backend returns deliberately different status objects from `/v1/jobs/job/:token` + * and `/v1/jobs/session`: only the session listing carries `maybe_failure_category_updated`, + * `maybe_failure_message` and `maybe_batch_token`. Both variants are reproduced. + * + * There is no generation in this fake yet, so jobs come from the seed: one finished, one in + * progress, one failed, all owned by the seeded user. + */ + +export type JobStatus = + | "pending" + | "started" + | "complete_success" + | "complete_failure" + | "attempt_failed" + | "dead" + | "cancelled_by_user" + | "cancelled_by_system"; + +export interface FakeJob { + readonly jobToken: string; + readonly ownerUserToken: string; + readonly inferenceCategory: string; + readonly status: JobStatus; + readonly progressPercentage: number; + readonly modelType: string | null; + readonly modelTitle: string | null; + readonly prompt: string | null; + readonly failureCategory: string | null; + readonly failureMessage: string | null; + readonly result: { mediaFileToken: string; cdnUrl: string; completedAt: string } | null; + readonly createdAt: string; + readonly updatedAt: string; +} + +/** The narrow failure enum still sent on `maybe_failure_category` (old clients). */ +const OLD_CLIENT_FAILURE_CATEGORIES = new Set([ + "face_not_detected", + "keep_alive_elapsed", + "not_yet_implemented", + "retryable_worker_error", +]); + +export function seededJobs(ownerUserToken: string): FakeJob[] { + return [ + { + jobToken: "jinf_fake_done_image", + ownerUserToken, + inferenceCategory: "image_generation", + status: "complete_success", + progressPercentage: 100, + modelType: "seedream_4", + modelTitle: "Seedream 4", + prompt: "a corgi running through a field", + failureCategory: null, + failureMessage: null, + result: { + mediaFileToken: "m_fake_corgi", + cdnUrl: "https://cdn.fake.test/media/m_fake_corgi.png", + completedAt: "2026-08-24T10:00:30Z", + }, + createdAt: "2026-08-24T10:00:00Z", + updatedAt: "2026-08-24T10:00:30Z", + }, + { + jobToken: "jinf_fake_running_video", + ownerUserToken, + inferenceCategory: "video_generation", + status: "started", + progressPercentage: 40, + modelType: "seedance_2p0", + modelTitle: "Seedance 2.0", + prompt: "the corgi, now in slow motion", + failureCategory: null, + failureMessage: null, + result: null, + createdAt: "2026-08-24T10:05:00Z", + updatedAt: "2026-08-24T10:05:20Z", + }, + { + jobToken: "jinf_fake_failed_image", + ownerUserToken, + inferenceCategory: "image_generation", + status: "complete_failure", + progressPercentage: 0, + modelType: "gpt_image_1", + modelTitle: "GPT Image 1", + prompt: "something the rules do not allow", + failureCategory: "rule_bans_user_text_prompt", + failureMessage: "The prompt was rejected by the content rules.", + result: null, + createdAt: "2026-08-24T09:00:00Z", + updatedAt: "2026-08-24T09:00:05Z", + }, + ]; +} + +/** `GET /v1/jobs/job/:token` and `GET /v1/jobs/batch`. */ +export function jobStatePayload(job: FakeJob): Record { + return { + job_token: job.jobToken, + request: requestPayload(job), + status: { + status: job.status, + maybe_extra_status_description: null, + maybe_assigned_worker: null, + maybe_assigned_cluster: null, + maybe_first_started_at: job.status === "pending" ? null : job.createdAt, + attempt_count: job.status === "pending" ? 0 : 1, + requires_keepalive: false, + maybe_failure_category: oldClientFailureCategory(job), + progress_percentage: job.progressPercentage, + }, + maybe_result: resultPayload(job, false), + created_at: job.createdAt, + updated_at: job.updatedAt, + }; +} + +/** `GET /v1/jobs/session` — the richer status the generation list consumes. */ +export function sessionJobPayload(job: FakeJob): Record { + return { + job_token: job.jobToken, + request: requestPayload(job), + status: { + status: job.status, + maybe_extra_status_description: null, + maybe_assigned_worker: null, + maybe_assigned_cluster: null, + maybe_first_started_at: job.status === "pending" ? null : job.createdAt, + maybe_current_execution_duration_seconds: job.status === "started" ? 20 : null, + attempt_count: job.status === "pending" ? 0 : 1, + requires_keepalive: false, + maybe_failure_category: oldClientFailureCategory(job), + maybe_failure_category_updated: job.failureCategory, + maybe_failure_message: job.failureMessage, + progress_percentage: job.progressPercentage, + }, + maybe_result: resultPayload(job, true), + created_at: job.createdAt, + updated_at: job.updatedAt, + }; +} + +function requestPayload(job: FakeJob): Record { + return { + inference_category: job.inferenceCategory, + maybe_prompt_token: null, + maybe_model_type: job.modelType, + maybe_model_token: null, + maybe_model_title: job.modelTitle, + maybe_raw_inference_text: job.prompt, + maybe_style_name: null, + maybe_live_portrait_details: null, + maybe_lipsync_details: null, + }; +} + +function resultPayload(job: FakeJob, includeBatchToken: boolean): Record | null { + if (job.result === null) return null; + return { + entity_type: "media_file", + entity_token: job.result.mediaFileToken, + media_links: { + cdn_url: job.result.cdnUrl, + maybe_thumbnail_template: null, + maybe_video_previews: null, + }, + maybe_public_bucket_media_path: null, + maybe_successfully_completed_at: job.result.completedAt, + ...(includeBatchToken ? { maybe_batch_token: null } : {}), + }; +} + +function oldClientFailureCategory(job: FakeJob): string | null { + return job.failureCategory !== null && OLD_CLIENT_FAILURE_CATEGORIES.has(job.failureCategory) + ? job.failureCategory + : null; +} diff --git a/mcp/fake-upstream/src/routes/jobs.ts b/mcp/fake-upstream/src/routes/jobs.ts new file mode 100644 index 0000000000..e206b2d6d9 --- /dev/null +++ b/mcp/fake-upstream/src/routes/jobs.ts @@ -0,0 +1,70 @@ +import { Hono } from "hono"; + +import { jobStatePayload, sessionJobPayload } from "../jobs"; +import { currentUser } from "../session"; +import type { FakeStore } from "../state"; + +/** + * Job status routes — ported from `infra/fake-storyteller-web` (routes/jobs.ts). Single and + * batch lookups need no user (tokens are unguessable); the session listing is session-only. + */ +export function jobRoutes(store: FakeStore): Hono { + const app = new Hono(); + + app.get("/v1/jobs/job/:token", (c) => { + const job = store.jobsByToken.get(c.req.param("token")); + if (job === undefined) { + return c.json( + { success: false, error_code: 404, error_code_str: "NotFound", message: "job not found" }, + 404, + ); + } + return c.json({ success: true, state: jobStatePayload(job) }); + }); + + app.get("/v1/jobs/batch", (c) => { + // Repeated `tokens=` params (what openapi-fetch sends) and comma-joined values both work. + const requested = c.req.queries("tokens")?.flatMap((value) => value.split(",")) ?? []; + const wanted = new Set(requested.filter((token) => token.length > 0)); + const states = [...wanted] + .map((token) => store.jobsByToken.get(token)) + .filter((job) => job !== undefined) + .map((job) => jobStatePayload(job)); + return c.json({ success: true, job_states: states }); + }); + + app.get("/v1/jobs/session", (c) => { + const user = currentUser(store, c.req.raw); + if (user === undefined) { + return c.json( + { + success: false, + error_code: 401, + error_code_str: "NotAuthorized", + message: "not authorized", + }, + 401, + ); + } + const include = statesParam(c.req.query("include_states")); + const exclude = statesParam(c.req.query("exclude_states")); + const owned = [...store.jobsByToken.values()] + .filter((job) => job.ownerUserToken === user.userToken) + .filter((job) => (include ? include.has(job.status) : true)) + .filter((job) => (exclude ? !exclude.has(job.status) : true)) + .sort((a, b) => (a.createdAt < b.createdAt ? 1 : -1)); + return c.json({ success: true, jobs: owned.map(sessionJobPayload) }); + }); + + return app; +} + +function statesParam(value: string | undefined): Set | undefined { + if (!value) return undefined; + return new Set( + value + .split(",") + .map((s) => s.trim()) + .filter((s) => s.length > 0), + ); +} diff --git a/mcp/fake-upstream/src/state.ts b/mcp/fake-upstream/src/state.ts index e6302fd135..47f94cf2fd 100644 --- a/mcp/fake-upstream/src/state.ts +++ b/mcp/fake-upstream/src/state.ts @@ -1,3 +1,5 @@ +import { type FakeJob, seededJobs } from "./jobs"; + /** * In-memory state for the fake. Seeded the same way as `infra/fake-storyteller-web` so anyone * who has used that fake finds the same account here. Everything resets when the Worker @@ -25,6 +27,7 @@ export interface FakeStore { readonly usersByToken: Map; /** signed session → user token */ readonly sessions: Map; + readonly jobsByToken: Map; } export const SEEDED_USER: FakeUser = { @@ -46,6 +49,7 @@ export function createStore(): FakeStore { return { usersByToken: new Map([[SEEDED_USER.userToken, SEEDED_USER]]), sessions: new Map(), + jobsByToken: new Map(seededJobs(SEEDED_USER.userToken).map((job) => [job.jobToken, job])), }; } diff --git a/mcp/fake-upstream/test/jobs.test.ts b/mcp/fake-upstream/test/jobs.test.ts new file mode 100644 index 0000000000..064f87ef30 --- /dev/null +++ b/mcp/fake-upstream/test/jobs.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from "vitest"; + +import { expectValid } from "../../test/helpers/contract"; +import { createFakeUpstream } from "../src/index"; +import { SEEDED_USER } from "../src/state"; + +async function signedIn() { + const app = createFakeUpstream(); + const login = await app.request("/v1/login", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + username_or_email: SEEDED_USER.username, + password: SEEDED_USER.password, + }), + }); + const { signed_session } = await login.json<{ signed_session: string }>(); + return { app, headers: { session: signed_session } }; +} + +describe("single job", () => { + it("returns a spec-shaped status with a result for a finished job", async () => { + const response = await createFakeUpstream().request("/v1/jobs/job/jinf_fake_done_image"); + expect(response.status).toBe(200); + const body = await response.json<{ + state: { status: { status: string }; maybe_result: { media_links: { cdn_url: string } } }; + }>(); + expectValid("GetInferenceJobStatusSuccessResponse", body); + expect(body.state.status.status).toBe("complete_success"); + expect(body.state.maybe_result.media_links.cdn_url).toMatch(/^https:\/\//); + }); + + it("404s an unknown token with the common envelope", async () => { + const response = await createFakeUpstream().request("/v1/jobs/job/jinf_nope"); + expect(response.status).toBe(404); + expectValid("CommonWebError", await response.json()); + }); +}); + +describe("batch", () => { + it("returns the known jobs among the requested tokens, in the spec shape", async () => { + const response = await createFakeUpstream().request( + "/v1/jobs/batch?tokens=jinf_fake_done_image&tokens=jinf_nope&tokens=jinf_fake_running_video", + ); + const body = await response.json<{ job_states: { job_token: string }[] }>(); + expectValid("BatchGetInferenceJobStatusSuccessResponse", body); + expect(body.job_states.map((j) => j.job_token).sort()).toEqual([ + "jinf_fake_done_image", + "jinf_fake_running_video", + ]); + }); +}); + +describe("session listing", () => { + it("is session-only", async () => { + expect((await createFakeUpstream().request("/v1/jobs/session")).status).toBe(401); + }); + + it("lists the user's jobs newest first, with the richer status variant", async () => { + const { app, headers } = await signedIn(); + const response = await app.request("/v1/jobs/session", { headers }); + const body = await response.json<{ + jobs: { job_token: string; status: Record }[]; + }>(); + expectValid("ListSessionJobsSuccessResponse", body); + expect(body.jobs.map((j) => j.job_token)).toEqual([ + "jinf_fake_running_video", + "jinf_fake_done_image", + "jinf_fake_failed_image", + ]); + const failed = body.jobs.find((j) => j.job_token === "jinf_fake_failed_image"); + expect(failed?.status.maybe_failure_category_updated).toBe("rule_bans_user_text_prompt"); + expect(failed?.status.maybe_failure_message).toMatch(/content rules/); + }); + + it("filters by include and exclude states", async () => { + const { app, headers } = await signedIn(); + const included = await ( + await app.request("/v1/jobs/session?include_states=started,pending", { headers }) + ).json<{ jobs: { job_token: string }[] }>(); + expect(included.jobs.map((j) => j.job_token)).toEqual(["jinf_fake_running_video"]); + const excluded = await ( + await app.request("/v1/jobs/session?exclude_states=complete_failure", { headers }) + ).json<{ jobs: { job_token: string }[] }>(); + expect(excluded.jobs.map((j) => j.job_token)).not.toContain("jinf_fake_failed_image"); + }); +}); diff --git a/mcp/src/mcp/server.ts b/mcp/src/mcp/server.ts index e7d5de1ebb..3322cc60a6 100644 --- a/mcp/src/mcp/server.ts +++ b/mcp/src/mcp/server.ts @@ -13,6 +13,7 @@ import type { UpstreamClient } from "../upstream/client"; import { estimateCost } from "./tools/estimate-cost"; import { getAccount } from "./tools/get-account"; import { getCreditBalance } from "./tools/get-credit-balance"; +import { getJobsStatus, getJobStatus, listRecentJobs } from "./tools/jobs"; import { listModels } from "./tools/list-models"; import { type ToolDefinition, ToolFailure, UpstreamSessionInvalid } from "./tools/types"; @@ -29,6 +30,9 @@ export const TOOLS: readonly ToolDefinition[] = [ getCreditBalance, listModels, estimateCost, + getJobStatus, + getJobsStatus, + listRecentJobs, ]; export interface McpServerDeps { diff --git a/mcp/src/mcp/tools/jobs.ts b/mcp/src/mcp/tools/jobs.ts new file mode 100644 index 0000000000..5893546df3 --- /dev/null +++ b/mcp/src/mcp/tools/jobs.ts @@ -0,0 +1,230 @@ +import { z } from "zod"; + +import { READ_ONLY_ANNOTATIONS, type ToolDefinition, unwrapUpstream } from "./types"; + +/** + * Job status tools. The three upstream endpoints return three slightly different payloads + * (the session listing alone carries the newer failure category and message); all are + * normalised to one `JobSummary` so a client sees the same shape wherever a job came from. + */ + +export const JOB_STATUSES = [ + "pending", + "started", + "complete_success", + "complete_failure", + "attempt_failed", + "dead", + "cancelled_by_user", + "cancelled_by_system", +] as const; + +const FINISHED = new Set([ + "complete_success", + "complete_failure", + "dead", + "cancelled_by_user", + "cancelled_by_system", +]); + +const JOB_TOKEN = z + .string() + .regex(/^jinf_[A-Za-z0-9_]+$/, "an inference job token starts with jinf_"); + +const jobSummarySchema = z.object({ + job_token: z.string(), + category: z.string().describe("What was generated, e.g. image_generation, video_generation."), + status: z.enum(JOB_STATUSES), + is_finished: z.boolean().describe("True for every terminal state, success or not."), + progress_percentage: z.number().int().min(0).max(100), + model: z.object({ type: z.string().nullable(), title: z.string().nullable() }), + prompt: z.string().nullable().describe("The text prompt, when the job had one."), + failure: z + .object({ category: z.string().nullable(), message: z.string().nullable() }) + .nullable() + .describe("Why the job failed, when it did; null otherwise."), + result: z + .object({ + entity_type: z.string(), + entity_token: z.string(), + cdn_url: z.string().describe("Where to view or download the output."), + thumbnail_template: z.string().nullable(), + completed_at: z.string().nullable(), + }) + .nullable() + .describe("The output, once there is one."), + created_at: z.string(), + updated_at: z.string(), +}); + +export type JobSummary = z.infer; + +/** The fields common to all three upstream payloads, plus the session-only extras. */ +interface RawJob { + job_token: string; + request: { + inference_category: string; + maybe_model_type?: string | null; + maybe_model_title?: string | null; + maybe_raw_inference_text?: string | null; + }; + status: { + status: (typeof JOB_STATUSES)[number]; + progress_percentage: number; + maybe_failure_category?: string | null; + /** Newer clients get the full enum; values this build predates arrive as `{ unknown }`. */ + maybe_failure_category_updated?: string | { unknown: string } | null; + maybe_failure_message?: string | null; + }; + maybe_result?: { + entity_type: string; + entity_token: string; + media_links: { cdn_url: string; maybe_thumbnail_template?: string | null }; + maybe_successfully_completed_at?: string | null; + } | null; + created_at: string; + updated_at: string; +} + +export function summarizeJob(raw: RawJob): JobSummary { + const failureCategory = + failureCategoryString(raw.status.maybe_failure_category_updated) ?? + raw.status.maybe_failure_category ?? + null; + const failureMessage = raw.status.maybe_failure_message ?? null; + const failed = + failureCategory !== null || failureMessage !== null || raw.status.status === "complete_failure"; + return { + job_token: raw.job_token, + category: raw.request.inference_category, + status: raw.status.status, + is_finished: FINISHED.has(raw.status.status), + progress_percentage: raw.status.progress_percentage, + model: { + type: raw.request.maybe_model_type ?? null, + title: raw.request.maybe_model_title ?? null, + }, + prompt: raw.request.maybe_raw_inference_text ?? null, + failure: failed ? { category: failureCategory, message: failureMessage } : null, + result: raw.maybe_result + ? { + entity_type: raw.maybe_result.entity_type, + entity_token: raw.maybe_result.entity_token, + cdn_url: raw.maybe_result.media_links.cdn_url, + thumbnail_template: raw.maybe_result.media_links.maybe_thumbnail_template ?? null, + completed_at: raw.maybe_result.maybe_successfully_completed_at ?? null, + } + : null, + created_at: raw.created_at, + updated_at: raw.updated_at, + }; +} + +/** The API's forward-compatible enum: a known literal, or `{ unknown }` holding the raw value. */ +function failureCategoryString( + value: string | { unknown: string } | null | undefined, +): string | null { + if (value === null || value === undefined) return null; + return typeof value === "string" ? value : value.unknown; +} + +export function describeJob(job: JobSummary): string { + const what = job.model.title ? `${job.category} with ${job.model.title}` : job.category; + const head = `${job.job_token}: ${what}, ${job.status.replace(/_/g, " ")}`; + if (job.status === "complete_success" && job.result) return `${head} — ${job.result.cdn_url}`; + if (job.failure) { + const why = job.failure.message ?? job.failure.category ?? "no reason given"; + return `${head} — ${why}`; + } + if (!job.is_finished) return `${head} (${String(job.progress_percentage)}%)`; + return head; +} + +const singleOutput = jobSummarySchema.shape; + +export const getJobStatus: ToolDefinition<{ job_token: typeof JOB_TOKEN }, typeof singleOutput> = { + name: "get_job_status", + title: "Get job status", + description: + "Returns the status of one Artcraft generation job by its token (jinf_…): state, progress, " + + "the model and prompt used, the failure reason if it failed, and the output's CDN URL once " + + "it succeeded. Use it to answer 'is my generation done?' for a specific job.", + requiredScope: "read:jobs", + inputSchema: { job_token: JOB_TOKEN.describe("The job token, e.g. jinf_abc123.") }, + outputSchema: singleOutput, + annotations: { ...READ_ONLY_ANNOTATIONS, title: "Get job status" }, + + async handler({ upstream }, { job_token }) { + const response = unwrapUpstream( + await upstream.GET("/v1/jobs/job/{token}", { params: { path: { token: job_token } } }), + ); + const job = summarizeJob(response.state); + return { structured: job, text: describeJob(job) }; + }, +}; + +const listOutput = { jobs: z.array(jobSummarySchema) }; + +export const getJobsStatus: ToolDefinition< + { job_tokens: z.ZodArray }, + typeof listOutput +> = { + name: "get_jobs_status", + title: "Get status of several jobs", + description: + "Returns the status of up to 50 generation jobs at once, by token. Jobs that do not exist " + + "are simply absent from the result. Prefer this over repeated get_job_status calls when " + + "checking on several generations.", + requiredScope: "read:jobs", + inputSchema: { job_tokens: z.array(JOB_TOKEN).min(1).max(50).describe("Job tokens (jinf_…).") }, + outputSchema: listOutput, + annotations: { ...READ_ONLY_ANNOTATIONS, title: "Get status of several jobs" }, + + async handler({ upstream }, { job_tokens }) { + const response = unwrapUpstream( + await upstream.GET("/v1/jobs/batch", { + params: { query: { tokens: [...new Set(job_tokens)] } }, + }), + ); + const jobs = response.job_states.map(summarizeJob); + const text = + jobs.length === 0 ? "None of those jobs were found." : jobs.map(describeJob).join("\n"); + return { structured: { jobs }, text }; + }, +}; + +const recentInput = { + include_states: z + .array(z.enum(JOB_STATUSES)) + .optional() + .describe("Only these states. Omit for all."), + exclude_states: z + .array(z.enum(JOB_STATUSES)) + .optional() + .describe("Everything except these states."), +}; + +export const listRecentJobs: ToolDefinition = { + name: "list_recent_jobs", + title: "List recent jobs", + description: + "Lists the signed-in user's recent generation jobs, newest first, with status, progress, " + + "model, prompt, failure reason, and output URL where available. Filter by state with " + + 'include_states / exclude_states (e.g. include_states: ["started", "pending"] for ' + + "what is still running). Use it for 'what did I generate today?' or 'is anything still running?'.", + requiredScope: "read:jobs", + inputSchema: recentInput, + outputSchema: listOutput, + annotations: { ...READ_ONLY_ANNOTATIONS, title: "List recent jobs" }, + + async handler({ upstream }, { include_states, exclude_states }) { + const query = { + ...(include_states?.length ? { include_states: include_states.join(",") } : {}), + ...(exclude_states?.length ? { exclude_states: exclude_states.join(",") } : {}), + }; + const response = unwrapUpstream(await upstream.GET("/v1/jobs/session", { params: { query } })); + const jobs = response.jobs.map(summarizeJob); + const text = jobs.length === 0 ? "No recent jobs." : jobs.map(describeJob).join("\n"); + return { structured: { jobs }, text }; + }, +}; diff --git a/mcp/test/mcp/handler.test.ts b/mcp/test/mcp/handler.test.ts index 54256082d0..a728de956f 100644 --- a/mcp/test/mcp/handler.test.ts +++ b/mcp/test/mcp/handler.test.ts @@ -139,7 +139,10 @@ describe("the protected MCP route", () => { "estimate_cost", "get_account", "get_credit_balance", + "get_job_status", + "get_jobs_status", "list_models", + "list_recent_jobs", ]); const result = await client.callTool({ name: "get_credit_balance", arguments: {} }); diff --git a/mcp/test/mcp/jobs.test.ts b/mcp/test/mcp/jobs.test.ts new file mode 100644 index 0000000000..55596aacfc --- /dev/null +++ b/mcp/test/mcp/jobs.test.ts @@ -0,0 +1,203 @@ +import { describe, expect, it } from "vitest"; + +import { + getJobsStatus, + getJobStatus, + listRecentJobs, + summarizeJob, +} from "../../src/mcp/tools/jobs"; +import { ToolFailure } from "../../src/mcp/tools/types"; +import { principalFromProps } from "../../src/tokens/principal"; +import { createUpstreamClient } from "../../src/upstream/client"; +import { createFakeUpstream } from "../../fake-upstream/src/index"; +import { SEEDED_USER } from "../../fake-upstream/src/state"; + +/** The tools run against the fake, signed in as the seeded user, recording the URLs sent. */ +async function fakeContext() { + const fake = createFakeUpstream(); + const login = await fake.request("/v1/login", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + username_or_email: SEEDED_USER.username, + password: SEEDED_USER.password, + }), + }); + const { signed_session } = await login.json<{ signed_session: string }>(); + const principal = principalFromProps({ + credential: { kind: "session", signedSession: signed_session }, + grantIssuedAt: 1_800_000_000_000, + userToken: SEEDED_USER.userToken, + username: SEEDED_USER.username, + displayName: SEEDED_USER.displayName, + scopes: ["read:jobs"], + }); + const urls: string[] = []; + const upstream = createUpstreamClient({ + baseUrl: "https://api.example.test", + use: "read", + credential: principal.credential, + fetch: (input, init) => { + const request = new Request(input, init); + const url = new URL(request.url); + urls.push(url.pathname + url.search); + return Promise.resolve( + fake.fetch(new Request(`http://fake${url.pathname}${url.search}`, request)), + ); + }, + }); + return { urls, context: { principal, upstream } }; +} + +describe("summarizeJob", () => { + it("prefers the newer failure category and reports a failure even without one", () => { + const base = { + job_token: "jinf_x", + request: { inference_category: "image_generation" }, + status: { status: "complete_failure" as const, progress_percentage: 0 }, + created_at: "t", + updated_at: "t", + }; + expect(summarizeJob(base).failure).toEqual({ category: null, message: null }); + expect( + summarizeJob({ + ...base, + status: { + ...base.status, + maybe_failure_category: "retryable_worker_error", + maybe_failure_category_updated: "generation_failed", + maybe_failure_message: "boom", + }, + }).failure, + ).toEqual({ category: "generation_failed", message: "boom" }); + }); + + it("keeps a failure category this build does not know, from the API's catch-all", () => { + const job = summarizeJob({ + job_token: "jinf_x", + request: { inference_category: "image_generation" }, + status: { + status: "complete_failure", + progress_percentage: 0, + maybe_failure_category_updated: { unknown: "brand_new_reason" }, + }, + created_at: "t", + updated_at: "t", + }); + expect(job.failure).toEqual({ category: "brand_new_reason", message: null }); + }); + + it("marks terminal states finished and running ones not", () => { + const of = (status: Parameters[0]["status"]["status"]) => + summarizeJob({ + job_token: "jinf_x", + request: { inference_category: "x" }, + status: { status, progress_percentage: 0 }, + created_at: "t", + updated_at: "t", + }).is_finished; + expect(of("pending")).toBe(false); + expect(of("started")).toBe(false); + expect(of("attempt_failed")).toBe(false); + expect(of("complete_success")).toBe(true); + expect(of("dead")).toBe(true); + expect(of("cancelled_by_user")).toBe(true); + }); +}); + +describe("get_job_status", () => { + it("returns a finished job with its output URL", async () => { + const { urls, context } = await fakeContext(); + const result = await getJobStatus.handler(context, { job_token: "jinf_fake_done_image" }); + expect(urls).toEqual(["/v1/jobs/job/jinf_fake_done_image"]); + expect(result.structured).toMatchObject({ + job_token: "jinf_fake_done_image", + category: "image_generation", + status: "complete_success", + is_finished: true, + model: { type: "seedream_4", title: "Seedream 4" }, + prompt: "a corgi running through a field", + failure: null, + result: { cdn_url: "https://cdn.fake.test/media/m_fake_corgi.png" }, + }); + expect(result.text).toBe( + "jinf_fake_done_image: image_generation with Seedream 4, complete success — https://cdn.fake.test/media/m_fake_corgi.png", + ); + }); + + it("describes a running job with its progress", async () => { + const { context } = await fakeContext(); + const result = await getJobStatus.handler(context, { job_token: "jinf_fake_running_video" }); + expect(result.structured.is_finished).toBe(false); + expect(result.text).toBe( + "jinf_fake_running_video: video_generation with Seedance 2.0, started (40%)", + ); + }); + + it("surfaces a 404 as a tool failure", async () => { + const { context } = await fakeContext(); + await expect(getJobStatus.handler(context, { job_token: "jinf_nope" })).rejects.toThrow( + ToolFailure, + ); + }); +}); + +describe("get_jobs_status", () => { + it("fetches several jobs in one call, deduplicating tokens, and omits unknown ones", async () => { + const { urls, context } = await fakeContext(); + const result = await getJobsStatus.handler(context, { + job_tokens: [ + "jinf_fake_done_image", + "jinf_fake_done_image", + "jinf_nope", + "jinf_fake_failed_image", + ], + }); + expect(urls[0]).toBe( + "/v1/jobs/batch?tokens=jinf_fake_done_image&tokens=jinf_nope&tokens=jinf_fake_failed_image", + ); + expect(result.structured.jobs.map((j) => j.job_token).sort()).toEqual([ + "jinf_fake_done_image", + "jinf_fake_failed_image", + ]); + expect(result.text.split("\n")).toHaveLength(2); + }); +}); + +describe("list_recent_jobs", () => { + it("lists the user's jobs newest first with the richer failure detail", async () => { + const { urls, context } = await fakeContext(); + const result = await listRecentJobs.handler(context, {}); + expect(urls).toEqual(["/v1/jobs/session"]); + expect(result.structured.jobs.map((j) => j.job_token)).toEqual([ + "jinf_fake_running_video", + "jinf_fake_done_image", + "jinf_fake_failed_image", + ]); + const failed = result.structured.jobs[2]; + expect(failed?.failure).toEqual({ + category: "rule_bans_user_text_prompt", + message: "The prompt was rejected by the content rules.", + }); + expect(result.text).toContain( + "complete failure — The prompt was rejected by the content rules.", + ); + }); + + it("passes state filters as comma-joined queries", async () => { + const { urls, context } = await fakeContext(); + const result = await listRecentJobs.handler(context, { + include_states: ["started", "pending"], + exclude_states: ["dead"], + }); + expect(urls[0]).toBe("/v1/jobs/session?include_states=started%2Cpending&exclude_states=dead"); + expect(result.structured.jobs.map((j) => j.job_token)).toEqual(["jinf_fake_running_video"]); + }); + + it("is declared read-only and scoped to read:jobs, as are its siblings", () => { + for (const tool of [getJobStatus, getJobsStatus, listRecentJobs]) { + expect(tool.requiredScope).toBe("read:jobs"); + expect(tool.annotations.readOnlyHint).toBe(true); + } + }); +});