From 9b4b9c8465ea80438ec232a7b4a612429ed882cf Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 14:01:14 +0000 Subject: [PATCH] fix(media): GPU guard for smaller cards, top-level type imports - The GPU guard waited for free VRAM at the operation's floor even on a card whose total memory is below it (an 8 GB card and a 12 GB video floor), so such jobs deferred until their budget ran out. The requirement is now capped at 90% of the card's total; a card with less than half the floor fails at once with "The GPU does not have enough memory for this job." - Tests pin that media tools reach a run only through the space's capabilities and within the bot's allowlist now that the image tool definitions exist. - Type-only imports are top-level `import type` in the media runtime, tools, RPCs and tests (AGENTS.md). - The ComfyUI settings test gives saveMediaSettings a store with nothing saved, since the save reads the previous settings before validating. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- apps/api/src/media-jobs.ts | 2 +- .../src/media-runtime.postgres.test.ts | 44 +++++++++++++++---- packages/adapters/src/media-runtime.ts | 32 ++++++++++---- packages/adapters/src/media-tools.test.ts | 30 +++++++++++++ packages/adapters/src/media-tools.ts | 3 +- .../src/media/comfyui-settings.test.ts | 5 ++- packages/testkit/src/media.postgres.test.ts | 13 +++--- 7 files changed, 101 insertions(+), 28 deletions(-) diff --git a/apps/api/src/media-jobs.ts b/apps/api/src/media-jobs.ts index 4a289e8..7cd8860 100644 --- a/apps/api/src/media-jobs.ts +++ b/apps/api/src/media-jobs.ts @@ -1,7 +1,7 @@ import { ORPCError } from "@orpc/server"; import type { MediaRuntime } from "@rakazo/adapter-kit"; +import type { EncryptedSecretStore } from "@rakazo/adapters"; import { - type EncryptedSecretStore, loadMediaSettings, MediaJobRejectedError, MediaSettingsInvalidError, diff --git a/packages/adapters/src/media-runtime.postgres.test.ts b/packages/adapters/src/media-runtime.postgres.test.ts index 8fdf406..75c597e 100644 --- a/packages/adapters/src/media-runtime.postgres.test.ts +++ b/packages/adapters/src/media-runtime.postgres.test.ts @@ -11,18 +11,14 @@ import type { } from "@rakazo/adapter-kit"; import { MediaJobDeferredError, mediaRunJobKey } from "@rakazo/adapter-kit"; import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts"; -import { clearThread, createDb, type PrismaClient } from "@rakazo/db"; +import type { PrismaClient } from "@rakazo/db"; +import { clearThread, createDb } from "@rakazo/db"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { z } from "zod"; import { storeArtifactFromFile } from "./artifact-files.js"; import { LocalArtifactStore } from "./artifacts.js"; -import { - createMediaRuntime, - MEDIA_WAKE_TRIGGER, - MediaJobRejectedError, - type MediaOperationRegistration, - type MediaRuntimeOptions, -} from "./media-runtime.js"; +import type { MediaOperationRegistration, MediaRuntimeOptions } from "./media-runtime.js"; +import { createMediaRuntime, MEDIA_WAKE_TRIGGER, MediaJobRejectedError } from "./media-runtime.js"; const describePostgres = process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip; @@ -424,6 +420,38 @@ describePostgres("media job runtime (PostgreSQL, fake operations)", () => { expect((await ctx.job(jobId)).status).toBe("done"); }); + it("runs on a card smaller than the floor once it is nearly free, and fails on far too small", async () => { + let free = 5_000; + const ctx = await setup({ + gpu: { + minVramMb: 12_000, + status: async () => ({ vramFreeMb: free, vramTotalMb: 8_192 }), + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ status: "waiting", deferrals: 1 }); + free = 7_600; + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + expect((await ctx.job(jobId)).status).toBe("done"); + + const tiny = await setup({ + gpu: { + minVramMb: 12_000, + status: async () => ({ vramFreeMb: 4_000, vramTotalMb: 4_096 }), + }, + }); + const small = await tiny.start({ prompt: "fox" }); + await tiny.runtime.runJob(small.jobId); + expect(tiny.handlerCalls).toHaveLength(0); + expect(await tiny.job(small.jobId)).toMatchObject({ + status: "failed", + error: "The GPU does not have enough memory for this job.", + deferrals: 0, + }); + }); + it("cancels a queued job at once and a running job through the lease heartbeat", async () => { const ctx = await setup({ run: async (context) => { diff --git a/packages/adapters/src/media-runtime.ts b/packages/adapters/src/media-runtime.ts index 928f08f..f36b3c8 100644 --- a/packages/adapters/src/media-runtime.ts +++ b/packages/adapters/src/media-runtime.ts @@ -27,24 +27,24 @@ import { mediaRunJobKey, runContinueJob, } from "@rakazo/adapter-kit"; +import type { + MediaJobBlock, + MediaJobStatus, + MediaOperation, + MediaOutput, + MessageBlock, +} from "@rakazo/contracts"; import { attachmentKindOf, MEDIA_JOB_ACTIVE_STATUSES, - type MediaJobBlock, - type MediaJobStatus, - type MediaOperation, MediaOperationSchema, - type MediaOutput, - type MessageBlock, } from "@rakazo/contracts"; import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; +import type { MediaJob, PrismaClient, ThreadEvents } from "@rakazo/db"; import { appendEventInTransaction, createThreadMessageInTransaction, - type MediaJob, Prisma, - type PrismaClient, - type ThreadEvents, withTransactionRetry, } from "@rakazo/db"; import { getLogger } from "@rakazo/logging"; @@ -143,6 +143,7 @@ const TIMEOUT_ERROR = "The job took too long."; const FAILED_ERROR = "The job failed."; const MISSING_SOURCE_ERROR = "A source file is no longer available."; const TOO_LARGE_ERROR = "The result is too large to keep."; +const GPU_TOO_SMALL_ERROR = "The GPU does not have enough memory for this job."; type AbortReason = "cancel" | "lease" | "abandoned" | "timeout"; @@ -191,6 +192,17 @@ function nonNegativeInt(value: number | undefined): number | undefined { return rounded >= 0 ? rounded : undefined; } +/** + * Free VRAM a job waits for: the operation's floor, capped below the card's total so a smaller + * card that can still run it (with offloading) does not wait forever. Null when the card is far + * too small to try. + */ +export function requiredVramMb(minVramMb: number, gpu: GpuStatus): number | null { + if (gpu.vramTotalMb > 0 && gpu.vramTotalMb * 2 < minVramMb) return null; + if (gpu.vramTotalMb <= 0) return minVramMb; + return Math.min(minVramMb, Math.floor(gpu.vramTotalMb * 0.9)); +} + function insideDir(dir: string, file: string): string { const resolved = path.resolve(dir, file); if (!resolved.startsWith(`${path.resolve(dir)}${path.sep}`)) { @@ -754,7 +766,9 @@ export function createMediaRuntime( const gpu = await registration.gpu .status(AbortSignal.any([signal, AbortSignal.timeout(GPU_STATUS_TIMEOUT_MS)])) .catch(() => undefined); - if (gpu && gpu.vramFreeMb < registration.gpu.minVramMb) { + const needed = gpu ? requiredVramMb(registration.gpu.minVramMb, gpu) : undefined; + if (needed === null) throw new MediaJobRejectedError(GPU_TOO_SMALL_ERROR); + if (gpu && needed !== undefined && gpu.vramFreeMb < needed) { throw new MediaJobDeferredError("Waiting for free GPU memory", GPU_RETRY_MS); } } diff --git a/packages/adapters/src/media-tools.test.ts b/packages/adapters/src/media-tools.test.ts index 9af7370..29b7a38 100644 --- a/packages/adapters/src/media-tools.test.ts +++ b/packages/adapters/src/media-tools.test.ts @@ -70,6 +70,36 @@ describe("media tools", () => { expect(selectBuiltinToolsForRun(base).some((tool) => isMediaToolName(tool.name))).toBe(false); }); + it("exposes a media tool only for operations the space can run, within the bot allowlist", () => { + const base = { + graphicalToolsAllowed: true, + groupId: null, + trigger: "user", + semanticMemoryEnabled: false, + messagingChannelRun: false, + bot: { kind: "specialist", restrictBuiltinTools: false, allowedBuiltinTools: [] as string[] }, + }; + const mediaNames = (options: Partial[0]>) => + selectBuiltinToolsForRun({ ...base, ...options }) + .map((tool) => tool.name) + .filter((name) => isMediaToolName(name)); + expect(mediaNames({})).toEqual([]); + expect(mediaNames({ mediaCapabilities: none })).toEqual([]); + expect(mediaNames({ mediaCapabilities: { ...none, generate_image: true } })).toEqual([ + "generate_image", + ]); + expect( + mediaNames({ + mediaCapabilities: { ...none, generate_image: true, edit_image: true }, + bot: { + kind: "specialist", + restrictBuiltinTools: true, + allowedBuiltinTools: ["edit_image"], + }, + }), + ).toEqual(["edit_image"]); + }); + it("adds continue_after to every media tool schema", () => { expect( mediaToolInputSchema({ diff --git a/packages/adapters/src/media-tools.ts b/packages/adapters/src/media-tools.ts index 7c0d219..0db6f26 100644 --- a/packages/adapters/src/media-tools.ts +++ b/packages/adapters/src/media-tools.ts @@ -1,5 +1,6 @@ import type { ConnectorTool, MediaCapabilities, MediaRuntime } from "@rakazo/adapter-kit"; -import { type MediaOperation, MediaOperationSchema } from "@rakazo/contracts"; +import type { MediaOperation } from "@rakazo/contracts"; +import { MediaOperationSchema } from "@rakazo/contracts"; import { MediaJobRejectedError } from "./media-runtime.js"; /** Media tools are named exactly like their operations and dispatched by one module. */ diff --git a/packages/adapters/src/media/comfyui-settings.test.ts b/packages/adapters/src/media/comfyui-settings.test.ts index 5026881..3a34c48 100644 --- a/packages/adapters/src/media/comfyui-settings.test.ts +++ b/packages/adapters/src/media/comfyui-settings.test.ts @@ -8,7 +8,10 @@ const comfyui = (workflows: { imageWorkflow?: string; videoWorkflow?: string }) describe("ComfyUI settings", () => { it("refuses custom workflows the provider could not fill", async () => { - const deps = {} as Parameters[0]; + // Nothing saved yet; validation fails before anything is stored. + const deps = { + prisma: { integrationProviderConfig: { findUnique: async () => null } }, + } as unknown as Parameters[0]; await expect( saveMediaSettings( deps, diff --git a/packages/testkit/src/media.postgres.test.ts b/packages/testkit/src/media.postgres.test.ts index f5de0c6..5c48c98 100644 --- a/packages/testkit/src/media.postgres.test.ts +++ b/packages/testkit/src/media.postgres.test.ts @@ -2,13 +2,10 @@ import { mkdtempSync, rmSync } from "node:fs"; import { writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import path from "node:path"; -import { ComposioEmulator, type MediaOperationRegistration } from "@rakazo/adapters"; -import { - type MediaJobBlock, - MediaOutputSchema, - type MediaSettingsView, - type MessageBlock, -} from "@rakazo/contracts"; +import type { MediaOperationRegistration } from "@rakazo/adapters"; +import { ComposioEmulator } from "@rakazo/adapters"; +import type { MediaJobBlock, MediaSettingsView, MessageBlock } from "@rakazo/contracts"; +import { MediaOutputSchema } from "@rakazo/contracts"; import { createThreadMessage } from "@rakazo/db"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; import type { createApp } from "../../../apps/api/src/app.ts"; @@ -154,7 +151,7 @@ describeWithDatabase("media RPCs (PostgreSQL)", () => { baseUrl: "http://127.0.0.1:8188/", apiKey: "comfy-test-key", imageModel: "sd_xl_base_1.0.safetensors", - imageWorkflow: '{"3":{"class_type":"KSampler","inputs":{}}}', + imageWorkflow: '{"6":{"class_type":"CLIPTextEncode","inputs":{"text":"{{prompt}}"}}}', videoWorkflow: "", }, transcription: { baseUrl: "http://127.0.0.1:8080", model: "whisper-1", apiKey: "" },