Merge claude/phase3-media into p3/b5-transcribe-studio
Brings in the media job RPCs, generic media tool dispatch and the image operations; transcribe joins the operation list and the tool catalog next to them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u
This commit is contained in:
Claude
commit
114103579f
@ -15,6 +15,7 @@ import type {
|
||||
ComposioProvider,
|
||||
ConnectorRegistry,
|
||||
DestinationEmulator,
|
||||
MediaOperationRegistration,
|
||||
RemoteConnectorDependencies,
|
||||
} from "@rakazo/adapters";
|
||||
import {
|
||||
@ -143,6 +144,8 @@ export async function createApp(
|
||||
email?: TransactionalEmailProvider;
|
||||
remoteConnectors?: RemoteConnectorDependencies;
|
||||
logger?: Logger;
|
||||
/** Test seam: media operations instead of the configured ones. */
|
||||
mediaOperations?: MediaOperationRegistration[];
|
||||
} = {},
|
||||
): Promise<AppHandles> {
|
||||
const {
|
||||
@ -155,6 +158,7 @@ export async function createApp(
|
||||
email: emailOverride,
|
||||
remoteConnectors,
|
||||
logger: loggerOverride,
|
||||
mediaOperations,
|
||||
...envOverrides
|
||||
} = overrides;
|
||||
const env = { ...loadEnv(process.env), ...envOverrides };
|
||||
@ -384,9 +388,9 @@ export async function createApp(
|
||||
events,
|
||||
jobs,
|
||||
artifacts,
|
||||
registrations: createMediaOperations({
|
||||
loadSettings: () => loadMediaSettings({ prisma, secrets }),
|
||||
}),
|
||||
registrations:
|
||||
mediaOperations ??
|
||||
createMediaOperations({ loadSettings: () => loadMediaSettings({ prisma, secrets }) }),
|
||||
});
|
||||
const executor = createRunExecutor({
|
||||
prisma,
|
||||
@ -468,6 +472,7 @@ export async function createApp(
|
||||
|
||||
const router = createRouter({
|
||||
cloudAgent,
|
||||
media,
|
||||
prisma,
|
||||
events,
|
||||
auth,
|
||||
|
||||
118
apps/api/src/media-jobs.ts
Normal file
118
apps/api/src/media-jobs.ts
Normal file
@ -0,0 +1,118 @@
|
||||
import { ORPCError } from "@orpc/server";
|
||||
import type { MediaRuntime } from "@rakazo/adapter-kit";
|
||||
import {
|
||||
type EncryptedSecretStore,
|
||||
loadMediaSettings,
|
||||
MediaJobRejectedError,
|
||||
MediaSettingsInvalidError,
|
||||
mediaSettingsView,
|
||||
saveMediaSettings,
|
||||
} from "@rakazo/adapters";
|
||||
import type { Actor, MediaJobStartInput, MediaSettingsUpdate } from "@rakazo/contracts";
|
||||
import type { PrismaClient } from "@rakazo/db";
|
||||
import { resolveThreadTarget } from "./thread-target.js";
|
||||
|
||||
export interface MediaRouteDeps {
|
||||
prisma: PrismaClient;
|
||||
secrets: EncryptedSecretStore;
|
||||
media?: MediaRuntime;
|
||||
}
|
||||
|
||||
/** Media backends run on the deployment owner's own server, so only they configure them. */
|
||||
function requireDeploymentOwner(actor: Actor) {
|
||||
if (!actor.isDeploymentOwner) throw new ORPCError("FORBIDDEN");
|
||||
}
|
||||
|
||||
function requireMedia(deps: MediaRouteDeps): MediaRuntime {
|
||||
if (!deps.media) throw new ORPCError("NOT_IMPLEMENTED");
|
||||
return deps.media;
|
||||
}
|
||||
|
||||
export async function getMediaSettings(deps: MediaRouteDeps, actor: Actor) {
|
||||
requireDeploymentOwner(actor);
|
||||
return mediaSettingsView(await loadMediaSettings(deps));
|
||||
}
|
||||
|
||||
export async function updateMediaSettings(
|
||||
deps: MediaRouteDeps,
|
||||
actor: Actor,
|
||||
input: MediaSettingsUpdate,
|
||||
) {
|
||||
requireDeploymentOwner(actor);
|
||||
try {
|
||||
return await saveMediaSettings(deps, input);
|
||||
} catch (error) {
|
||||
if (error instanceof MediaSettingsInvalidError) {
|
||||
throw new ORPCError("BAD_REQUEST", { message: error.message });
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Starts an operation a user picked on a card (export, transcribe, edit) in the same conversation.
|
||||
* The thread is authorized like every thread procedure; the runtime checks that the operation's
|
||||
* source artifacts belong to the actor, and the job can only ever read the actor's artifacts.
|
||||
*/
|
||||
export async function startMediaJob(deps: MediaRouteDeps, actor: Actor, input: MediaJobStartInput) {
|
||||
const media = requireMedia(deps);
|
||||
const target = await resolveThreadTarget(deps.prisma, actor, input);
|
||||
const source = await deps.prisma.message.findFirst({
|
||||
where: { id: input.messageId, threadId: target.threadId },
|
||||
select: { botId: true },
|
||||
});
|
||||
if (!source) throw new ORPCError("NOT_FOUND");
|
||||
const bots =
|
||||
target.kind === "bot"
|
||||
? [target.botId]
|
||||
: [
|
||||
...(target.leadBotId && target.memberBotIds.includes(target.leadBotId)
|
||||
? [target.leadBotId]
|
||||
: []),
|
||||
...target.memberBotIds,
|
||||
];
|
||||
const botId = source.botId && bots.includes(source.botId) ? source.botId : bots[0];
|
||||
if (!botId) throw new ORPCError("NOT_FOUND");
|
||||
try {
|
||||
return await media.start({
|
||||
operation: input.operation,
|
||||
input: input.input,
|
||||
spaceId: actor.spaceId,
|
||||
userId: actor.userId,
|
||||
botId,
|
||||
groupId: target.kind === "group" ? target.groupId : null,
|
||||
threadId: target.threadId,
|
||||
runId: null,
|
||||
});
|
||||
} catch (error) {
|
||||
if (error instanceof MediaJobRejectedError) {
|
||||
throw new ORPCError("BAD_REQUEST", { message: error.message });
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async function ownedJobId(deps: MediaRouteDeps, actor: Actor, jobId: string): Promise<string> {
|
||||
const job = await deps.prisma.mediaJob.findFirst({
|
||||
where: { id: jobId, spaceId: actor.spaceId, userId: actor.userId },
|
||||
select: { id: true },
|
||||
});
|
||||
if (!job) throw new ORPCError("NOT_FOUND");
|
||||
return job.id;
|
||||
}
|
||||
|
||||
export async function getMediaJob(deps: MediaRouteDeps, actor: Actor, jobId: string) {
|
||||
const media = requireMedia(deps);
|
||||
const snapshot = await media.get(await ownedJobId(deps, actor, jobId));
|
||||
if (!snapshot) throw new ORPCError("NOT_FOUND");
|
||||
return snapshot;
|
||||
}
|
||||
|
||||
export async function cancelMediaJob(deps: MediaRouteDeps, actor: Actor, jobId: string) {
|
||||
const media = requireMedia(deps);
|
||||
const id = await ownedJobId(deps, actor, jobId);
|
||||
await media.cancel(id);
|
||||
const snapshot = await media.get(id);
|
||||
if (!snapshot) throw new ORPCError("NOT_FOUND");
|
||||
return snapshot;
|
||||
}
|
||||
@ -6,6 +6,7 @@ import type {
|
||||
ArtifactStore,
|
||||
ConnectorCatalogItem,
|
||||
JobPublisher,
|
||||
MediaRuntime,
|
||||
MemoryStore,
|
||||
SandboxProvider,
|
||||
} from "@rakazo/adapter-kit";
|
||||
@ -166,6 +167,13 @@ import {
|
||||
renameMainConversation,
|
||||
} from "./main-conversation.js";
|
||||
import { buildMcpUpdateMaterial } from "./mcp-material.js";
|
||||
import {
|
||||
cancelMediaJob,
|
||||
getMediaJob,
|
||||
getMediaSettings,
|
||||
startMediaJob,
|
||||
updateMediaSettings,
|
||||
} from "./media-jobs.js";
|
||||
import {
|
||||
disconnectMemoryProvider,
|
||||
getMemoryProviderConfig,
|
||||
@ -447,6 +455,7 @@ function mcpAssignmentDto(row: {
|
||||
|
||||
export interface RouterDeps {
|
||||
cloudAgent?: CloudAgentConnection | null;
|
||||
media?: MediaRuntime;
|
||||
prisma: PrismaClient;
|
||||
events: ThreadEvents;
|
||||
auth: Auth;
|
||||
@ -3482,6 +3491,27 @@ export function createRouter(deps: RouterDeps) {
|
||||
return { ok: true as const };
|
||||
}),
|
||||
},
|
||||
media: {
|
||||
settings: {
|
||||
get: authed.media.settings.get.handler(({ context }) =>
|
||||
getMediaSettings(deps, context.actor),
|
||||
),
|
||||
update: authed.media.settings.update.handler(({ context, input }) =>
|
||||
updateMediaSettings(deps, context.actor, input),
|
||||
),
|
||||
},
|
||||
jobs: {
|
||||
start: authed.media.jobs.start.handler(({ context, input }) =>
|
||||
startMediaJob(deps, context.actor, input),
|
||||
),
|
||||
get: authed.media.jobs.get.handler(({ context, input }) =>
|
||||
getMediaJob(deps, context.actor, input.jobId),
|
||||
),
|
||||
cancel: authed.media.jobs.cancel.handler(({ context, input }) =>
|
||||
cancelMediaJob(deps, context.actor, input.jobId),
|
||||
),
|
||||
},
|
||||
},
|
||||
connections: {
|
||||
catalog: authed.connections.catalog.handler(async ({ context, input }) => {
|
||||
const adapterContext = connectionContext(
|
||||
|
||||
@ -2,6 +2,9 @@ import type { ConnectorTool } from "@rakazo/adapter-kit";
|
||||
import {
|
||||
BotSecretDestination,
|
||||
BotSecretName,
|
||||
EditImageInputSchema,
|
||||
GenerateImageInputSchema,
|
||||
GenerateVideoInputSchema,
|
||||
MODEL_TIERS,
|
||||
SecretAskPurpose,
|
||||
SecretHttpRequest,
|
||||
@ -11,7 +14,14 @@ import {
|
||||
import { z } from "zod";
|
||||
import { mediaToolInputSchema } from "./media-tools.js";
|
||||
|
||||
type MediaToolSchema = Parameters<typeof mediaToolInputSchema>[0];
|
||||
/** A media tool's parameters from its operation's input schema, plus `continue_after`. */
|
||||
function mediaParameters(schema: z.ZodType) {
|
||||
const json = z.toJSONSchema(schema, { io: "input" });
|
||||
return mediaToolInputSchema({
|
||||
properties: (json.properties ?? {}) as Record<string, unknown>,
|
||||
required: json.required,
|
||||
});
|
||||
}
|
||||
|
||||
export const SUBAGENT_MODEL_TIER_DESCRIPTION =
|
||||
"Optional model tier: fast for quick simple work, code for programming, standard otherwise. Uses the model this space maps to the tier, or your own model when it maps none. Ignored when model_provider and model_id are set.";
|
||||
@ -330,11 +340,29 @@ export const builtinAgentTools: ConnectorTool[] = [
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "generate_image",
|
||||
description:
|
||||
"Generate up to 4 images from a prompt on the self-hosted image server, or rework a source image (image-to-image).",
|
||||
inputSchema: mediaParameters(GenerateImageInputSchema),
|
||||
},
|
||||
{
|
||||
name: "generate_video",
|
||||
description:
|
||||
"Generate a short video clip from a prompt on the self-hosted image server, optionally starting from an image.",
|
||||
inputSchema: mediaParameters(GenerateVideoInputSchema),
|
||||
},
|
||||
{
|
||||
name: "edit_image",
|
||||
description:
|
||||
"Edit an image artifact into a new one: crop, rotate, flip, resize (crop, pad, or blur to fit a format), add borders, convert the format, or repaint it from a prompt (only when generate_image is available). Edits apply in that order.",
|
||||
inputSchema: mediaParameters(EditImageInputSchema),
|
||||
},
|
||||
{
|
||||
name: "transcribe",
|
||||
description:
|
||||
"Transcribe speech in an audio or video artifact into SRT and VTT subtitle files, shown on a media card. Set continue_after to get the transcript back, e.g. to burn the .srt in with edit_video.",
|
||||
inputSchema: mediaToolInputSchema(z.toJSONSchema(TranscribeInputSchema) as MediaToolSchema),
|
||||
inputSchema: mediaParameters(TranscribeInputSchema),
|
||||
},
|
||||
{
|
||||
name: "add_mcp_server",
|
||||
|
||||
227
packages/adapters/src/comfyui-provider.test.ts
Normal file
227
packages/adapters/src/comfyui-provider.test.ts
Normal file
@ -0,0 +1,227 @@
|
||||
import sharp from "sharp";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import type { ComfyUiProviderConfig } from "./comfyui-provider.js";
|
||||
import {
|
||||
createComfyUiProvider,
|
||||
describeExecutionError,
|
||||
parseGpuStatus,
|
||||
} from "./comfyui-provider.js";
|
||||
import { ComfyUiEmulator } from "./testing/comfyui-emulator.js";
|
||||
|
||||
const BASE_URL = "http://127.0.0.1:8188";
|
||||
const IMAGE = { prompt: "a lighthouse", width: 512, height: 512 };
|
||||
|
||||
function setup(config: Partial<ComfyUiProviderConfig> = {}, emulator = new ComfyUiEmulator()) {
|
||||
const provider = createComfyUiProvider({
|
||||
baseUrl: BASE_URL,
|
||||
fetch: emulator.fetch,
|
||||
pollIntervalMs: 1,
|
||||
...config,
|
||||
});
|
||||
return { provider, emulator };
|
||||
}
|
||||
|
||||
function submitted(emulator: ComfyUiEmulator, index = 0) {
|
||||
const graph = emulator.prompts[index]?.graph ?? {};
|
||||
const byClass = (classType: string) =>
|
||||
Object.values(graph).find((node) => node.class_type === classType)?.inputs ?? {};
|
||||
return { graph, byClass };
|
||||
}
|
||||
|
||||
describe("createComfyUiProvider", () => {
|
||||
it("refuses metadata hosts, public hosts and keys over public http", () => {
|
||||
expect(() => createComfyUiProvider({ baseUrl: "http://169.254.169.254" })).toThrow(
|
||||
"blocked metadata or link-local host",
|
||||
);
|
||||
expect(() => createComfyUiProvider({ baseUrl: "http://comfy.example.com" })).toThrow(
|
||||
"Public model endpoints are blocked",
|
||||
);
|
||||
expect(() => createComfyUiProvider({ baseUrl: "ftp://127.0.0.1" })).toThrow("http or https");
|
||||
});
|
||||
|
||||
it("validates custom workflows up front", () => {
|
||||
expect(() => createComfyUiProvider({ baseUrl: BASE_URL, imageWorkflow: "{" })).toThrow(
|
||||
"Workflow is not valid JSON",
|
||||
);
|
||||
expect(() =>
|
||||
createComfyUiProvider({ baseUrl: BASE_URL, videoWorkflow: '{"nodes": [], "links": []}' }),
|
||||
).toThrow("API format");
|
||||
});
|
||||
|
||||
it("keeps a reverse-proxy path prefix and sends the API key", async () => {
|
||||
const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" });
|
||||
const { provider } = setup({ baseUrl: `${BASE_URL}/comfy/`, apiKey: "comfy-key" }, emulator);
|
||||
await provider.generateImage(IMAGE);
|
||||
expect(emulator.requests.map((request) => request.path)).toContain("prompt");
|
||||
expect(emulator.requests.every((request) => request.authorization === "Bearer comfy-key")).toBe(
|
||||
true,
|
||||
);
|
||||
const { provider: keyless } = setup({ baseUrl: `${BASE_URL}/comfy` }, emulator);
|
||||
await expect(keyless.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 401");
|
||||
});
|
||||
|
||||
it("fills the default text-to-image workflow", async () => {
|
||||
const { provider, emulator } = setup();
|
||||
await provider.generateImage({ ...IMAGE, width: 1001, negativePrompt: "fog", seed: 42 });
|
||||
const { byClass } = submitted(emulator);
|
||||
expect(byClass("CheckpointLoaderSimple").ckpt_name).toBe("sd_xl_base_1.0.safetensors");
|
||||
expect(byClass("EmptyLatentImage")).toMatchObject({ width: 1000, height: 512 });
|
||||
expect(byClass("KSampler")).toMatchObject({ seed: 42, denoise: 1 });
|
||||
const texts = emulator.prompts[0]!.graph;
|
||||
expect(Object.values(texts).map((node) => node.inputs.text)).toEqual(
|
||||
expect.arrayContaining(["a lighthouse", "fog"]),
|
||||
);
|
||||
});
|
||||
|
||||
it("uses the configured model, else the first checkpoint the video workflows do not use", async () => {
|
||||
const emulator = new ComfyUiEmulator();
|
||||
emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors", "v1-5-pruned.safetensors"];
|
||||
emulator.checkpointListFormat = "legacy";
|
||||
await setup({}, emulator).provider.generateImage(IMAGE);
|
||||
expect(submitted(emulator, 0).byClass("CheckpointLoaderSimple").ckpt_name).toBe(
|
||||
"v1-5-pruned.safetensors",
|
||||
);
|
||||
emulator.checkpoints.push("dreamshaper.safetensors");
|
||||
await setup({ imageModel: "dreamshaper.safetensors" }, emulator).provider.generateImage(IMAGE);
|
||||
expect(submitted(emulator, 1).byClass("CheckpointLoaderSimple").ckpt_name).toBe(
|
||||
"dreamshaper.safetensors",
|
||||
);
|
||||
emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors"];
|
||||
await expect(setup({}, emulator).provider.generateImage(IMAGE)).rejects.toThrow(
|
||||
"ComfyUI has no image checkpoint installed",
|
||||
);
|
||||
});
|
||||
|
||||
it("uploads the source image for image-to-image", async () => {
|
||||
const { provider, emulator } = setup();
|
||||
const bytes = new Uint8Array(Buffer.from((await sourcePng()).bytes));
|
||||
await provider.generateImage({
|
||||
...IMAGE,
|
||||
source: { bytes, mimeType: "image/png" },
|
||||
strength: 0.3,
|
||||
});
|
||||
const { byClass } = submitted(emulator);
|
||||
const uploaded = String(byClass("LoadImage").image);
|
||||
expect(uploaded).toMatch(/^evoo-.+\.png$/);
|
||||
expect(Buffer.from(emulator.uploads.get(uploaded)!).equals(Buffer.from(bytes))).toBe(true);
|
||||
expect(byClass("KSampler").denoise).toBe(0.3);
|
||||
});
|
||||
|
||||
it("lets a custom workflow replace only the default it matches", async () => {
|
||||
const custom = {
|
||||
"1": {
|
||||
class_type: "CustomTextToImage",
|
||||
inputs: { text: "style: {{prompt}}", seed: "{{seed}}" },
|
||||
},
|
||||
"2": { class_type: "SaveImage", inputs: { images: ["1", 0], filename_prefix: "x" } },
|
||||
};
|
||||
const { provider, emulator } = setup({ imageWorkflow: JSON.stringify(custom) });
|
||||
await provider.generateImage(IMAGE);
|
||||
expect(submitted(emulator, 0).byClass("CustomTextToImage").text).toBe("style: a lighthouse");
|
||||
await provider.generateImage({ ...IMAGE, source: await sourcePng() });
|
||||
expect(submitted(emulator, 1).byClass("LoadImage").image).toBeDefined();
|
||||
});
|
||||
|
||||
it("fills the LTX-Video workflow with an 8n+1 frame count and reports the clip", async () => {
|
||||
const { provider, emulator } = setup();
|
||||
const [clip] = await provider.generateVideo({
|
||||
prompt: "a paper boat",
|
||||
width: 470,
|
||||
height: 836,
|
||||
durationSeconds: 5,
|
||||
});
|
||||
const { byClass } = submitted(emulator);
|
||||
expect(byClass("EmptyLTXVLatentVideo")).toMatchObject({ width: 480, height: 832, length: 121 });
|
||||
expect(byClass("CreateVideo").fps).toBe(24);
|
||||
expect(clip).toMatchObject({
|
||||
mimeType: "video/mp4",
|
||||
width: 480,
|
||||
height: 832,
|
||||
durationMs: 5042,
|
||||
});
|
||||
});
|
||||
|
||||
it("reads video from a custom workflow's gifs output and skips frame previews", async () => {
|
||||
const custom = {
|
||||
"1": { class_type: "TextToFrames", inputs: { text: "{{prompt}}" } },
|
||||
"2": { class_type: "PreviewImage", inputs: { images: ["1", 0] } },
|
||||
"3": { class_type: "VHS_VideoCombine", inputs: { images: ["1", 0], frame_rate: 16 } },
|
||||
};
|
||||
const { provider } = setup({ videoWorkflow: JSON.stringify(custom) });
|
||||
const clips = await provider.generateVideo({ ...IMAGE, durationSeconds: 2 });
|
||||
expect(clips.map((clip) => [clip.mimeType, clip.durationMs])).toEqual([
|
||||
["video/mp4", undefined],
|
||||
]);
|
||||
});
|
||||
|
||||
it("explains rejected workflows and failed runs", async () => {
|
||||
const emulator = new ComfyUiEmulator();
|
||||
const { provider } = setup({ imageModel: "missing.safetensors" }, emulator);
|
||||
await expect(provider.generateImage(IMAGE)).rejects.toThrow(
|
||||
"ComfyUI rejected the workflow: CheckpointLoaderSimple: Value not in list: ckpt_name: 'missing.safetensors' not in list",
|
||||
);
|
||||
const working = setup({}, emulator).provider;
|
||||
emulator.failNextExecution = "CUDA out of memory. Tried to allocate 2.00 GiB\nTraceback...";
|
||||
await expect(working.generateImage(IMAGE)).rejects.toThrow(
|
||||
/^ComfyUI failed in KSampler: CUDA out of memory\. Tried to allocate 2\.00 GiB$/,
|
||||
);
|
||||
emulator.failNextStatus = 500;
|
||||
await expect(working.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 500");
|
||||
});
|
||||
|
||||
it("reports an unreachable server", async () => {
|
||||
const provider = createComfyUiProvider({
|
||||
baseUrl: BASE_URL,
|
||||
fetch: async () => {
|
||||
throw new TypeError("fetch failed");
|
||||
},
|
||||
});
|
||||
await expect(provider.generateImage(IMAGE)).rejects.toThrow("ComfyUI is unreachable");
|
||||
await expect(provider.gpuStatus()).rejects.toThrow("ComfyUI is unreachable");
|
||||
});
|
||||
|
||||
it("interrupts its own running prompt when the caller cancels", async () => {
|
||||
const emulator = new ComfyUiEmulator();
|
||||
emulator.pendingPolls = 1_000_000;
|
||||
const { provider } = setup({ pollIntervalMs: 5 }, emulator);
|
||||
const controller = new AbortController();
|
||||
const run = provider.generateImage({ ...IMAGE, signal: controller.signal });
|
||||
while (!emulator.requests.some((request) => request.path.startsWith("history/"))) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 1));
|
||||
}
|
||||
controller.abort(new Error("cancelled by user"));
|
||||
await expect(run).rejects.toThrow("cancelled by user");
|
||||
expect(emulator.deleted).toEqual(["prompt-1"]);
|
||||
expect(emulator.interrupted).toEqual(["prompt-1"]);
|
||||
});
|
||||
|
||||
it("reports GPU memory other processes leave to ComfyUI", async () => {
|
||||
const emulator = new ComfyUiEmulator();
|
||||
emulator.gpu = { totalMb: 12_288, freeMb: 3_000, torchReservedMb: 7_000, torchFreeMb: 1_000 };
|
||||
await expect(setup({}, emulator).provider.gpuStatus()).resolves.toEqual({
|
||||
vramFreeMb: 9_000,
|
||||
vramTotalMb: 12_288,
|
||||
});
|
||||
emulator.gpu = null;
|
||||
await expect(setup({}, emulator).provider.gpuStatus()).resolves.toBeUndefined();
|
||||
expect(parseGpuStatus({ devices: [{ type: "cuda", vram_total: 0, vram_free: 0 }] })).toBe(
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
it("names interrupted runs", () => {
|
||||
expect(describeExecutionError([["execution_interrupted", {}]])).toBe(
|
||||
"ComfyUI interrupted the run",
|
||||
);
|
||||
expect(describeExecutionError(undefined)).toBe("ComfyUI failed");
|
||||
});
|
||||
});
|
||||
|
||||
async function sourcePng() {
|
||||
const png = await sharp({
|
||||
create: { width: 32, height: 32, channels: 3, background: { r: 1, g: 2, b: 3 } },
|
||||
})
|
||||
.png()
|
||||
.toBuffer();
|
||||
return { bytes: new Uint8Array(png), mimeType: "image/png" };
|
||||
}
|
||||
593
packages/adapters/src/comfyui-provider.ts
Normal file
593
packages/adapters/src/comfyui-provider.ts
Normal file
@ -0,0 +1,593 @@
|
||||
import { randomInt, randomUUID } from "node:crypto";
|
||||
import { setTimeout as sleep } from "node:timers/promises";
|
||||
import type {
|
||||
AdapterDescriptor,
|
||||
GeneratedMedia,
|
||||
GpuStatus,
|
||||
ImageGenerationCapabilities,
|
||||
ImageGenerationProvider,
|
||||
ImageGenerationRequest,
|
||||
MediaFile,
|
||||
MediaProgressListener,
|
||||
VideoGenerationCapabilities,
|
||||
VideoGenerationProvider,
|
||||
VideoGenerationRequest,
|
||||
} from "@rakazo/adapter-kit";
|
||||
import type { AttachmentKind } from "@rakazo/contracts";
|
||||
import {
|
||||
ARTIFACT_MAX_BYTES_BY_KIND,
|
||||
attachmentKindOf,
|
||||
MEDIA_VIDEO_MAX_SECONDS,
|
||||
} from "@rakazo/contracts";
|
||||
import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core";
|
||||
import type { ComfyWorkflow, WorkflowTemplate, WorkflowValues } from "@rakazo/core/media/workflows";
|
||||
import {
|
||||
DEFAULT_WORKFLOWS,
|
||||
fillWorkflow,
|
||||
parseWorkflowTemplate,
|
||||
} from "@rakazo/core/media/workflows";
|
||||
import {
|
||||
assertAllowedOpenAiCompatibleRequestUrl,
|
||||
assertHttpsForKeyedOpenAiCompatibleUrl,
|
||||
} from "./openai-compatible-url.js";
|
||||
import { createOpenAiCompatibleFetch } from "./pi-openai-compatible-provider.js";
|
||||
import { readBodyCapped } from "./web-ssrf.js";
|
||||
|
||||
export const COMFYUI_PROVIDER_ID = "comfyui";
|
||||
/** Frame rate of generated video; frame counts follow LTX-Video's 8n+1 rule. */
|
||||
export const COMFYUI_VIDEO_FPS = 24;
|
||||
const MAX_IMAGE_SIDE = 2048;
|
||||
const REQUEST_TIMEOUT_MS = 30_000;
|
||||
const TRANSFER_TIMEOUT_MS = 120_000;
|
||||
const CANCEL_TIMEOUT_MS = 5_000;
|
||||
const IMAGE_PROMPT_TIMEOUT_MS = 10 * 60_000;
|
||||
const VIDEO_PROMPT_TIMEOUT_MS = 30 * 60_000;
|
||||
const POLL_INTERVAL_MS = 1_000;
|
||||
const MAX_JSON_BYTES = 8 * 1024 * 1024;
|
||||
const MAX_ERROR_CHARS = 400;
|
||||
const MIB = 1024 * 1024;
|
||||
const IMAGE_STEPS = 25;
|
||||
const VIDEO_STEPS = 30;
|
||||
const DEFAULT_STRENGTH = 0.6;
|
||||
const DEFAULT_NEGATIVE_PROMPTS = {
|
||||
image: "blurry, low quality, distorted, watermark, text",
|
||||
video: "worst quality, inconsistent motion, blurry, jittery, distorted",
|
||||
};
|
||||
const JSON_HEADERS = { Accept: "application/json", "Content-Type": "application/json" };
|
||||
|
||||
export interface ComfyUiProviderConfig {
|
||||
baseUrl: string;
|
||||
apiKey?: string;
|
||||
/** Checkpoint for `{{model}}`; the first installed image checkpoint when omitted. */
|
||||
imageModel?: string;
|
||||
/** API-format workflow JSON that replaces the matching default image workflow. */
|
||||
imageWorkflow?: string;
|
||||
/** API-format workflow JSON that replaces the matching default video workflow. */
|
||||
videoWorkflow?: string;
|
||||
fetch?: typeof globalThis.fetch;
|
||||
/** How long to wait between checks for a finished prompt; tests shorten it. */
|
||||
pollIntervalMs?: number;
|
||||
}
|
||||
|
||||
export interface ComfyUiProvider extends ImageGenerationProvider, VideoGenerationProvider {
|
||||
describe(): AdapterDescriptor<ImageGenerationCapabilities & VideoGenerationCapabilities>;
|
||||
gpuStatus(signal?: AbortSignal): Promise<GpuStatus | undefined>;
|
||||
}
|
||||
|
||||
type MediaKind = "image" | "video";
|
||||
type ReadBody<T> = (response: Response, signal: AbortSignal) => Promise<T>;
|
||||
|
||||
interface OutputFile {
|
||||
filename: string;
|
||||
subfolder: string;
|
||||
type: string;
|
||||
mimeType: string;
|
||||
kind: AttachmentKind;
|
||||
}
|
||||
|
||||
interface GenerationJob {
|
||||
kind: MediaKind;
|
||||
template: WorkflowTemplate;
|
||||
prompt: string;
|
||||
negativePrompt?: string;
|
||||
width: number;
|
||||
height: number;
|
||||
seed?: number;
|
||||
count: number;
|
||||
source?: MediaFile;
|
||||
strength?: number;
|
||||
frames?: number;
|
||||
signal?: AbortSignal;
|
||||
onProgress?: MediaProgressListener;
|
||||
}
|
||||
|
||||
/**
|
||||
* Image and video generation on a self-hosted ComfyUI server through its HTTP API: submit an
|
||||
* API-format workflow, poll its history, download the outputs. Requests go through the same
|
||||
* address-checked fetch as OpenAI-compatible model servers.
|
||||
*/
|
||||
export function createComfyUiProvider(config: ComfyUiProviderConfig): ComfyUiProvider {
|
||||
return new ComfyUiClient(config);
|
||||
}
|
||||
|
||||
class ComfyUiClient implements ComfyUiProvider {
|
||||
private readonly baseUrl: URL;
|
||||
private readonly apiKey: string | undefined;
|
||||
private readonly imageModel: string | undefined;
|
||||
private readonly imageWorkflow: WorkflowTemplate | undefined;
|
||||
private readonly videoWorkflow: WorkflowTemplate | undefined;
|
||||
private readonly fetch: typeof globalThis.fetch;
|
||||
private readonly pollIntervalMs: number;
|
||||
private readonly clientId = randomUUID();
|
||||
|
||||
constructor(config: ComfyUiProviderConfig) {
|
||||
this.baseUrl = comfyUiBaseUrl(config.baseUrl);
|
||||
this.apiKey = config.apiKey?.trim() || undefined;
|
||||
assertHttpsForKeyedOpenAiCompatibleUrl(this.baseUrl, this.apiKey);
|
||||
this.imageModel = config.imageModel?.trim() || undefined;
|
||||
this.imageWorkflow = config.imageWorkflow?.trim()
|
||||
? parseWorkflowTemplate(config.imageWorkflow)
|
||||
: undefined;
|
||||
this.videoWorkflow = config.videoWorkflow?.trim()
|
||||
? parseWorkflowTemplate(config.videoWorkflow)
|
||||
: undefined;
|
||||
this.fetch = createOpenAiCompatibleFetch(config.fetch);
|
||||
this.pollIntervalMs = config.pollIntervalMs ?? POLL_INTERVAL_MS;
|
||||
}
|
||||
|
||||
describe() {
|
||||
return {
|
||||
id: COMFYUI_PROVIDER_ID,
|
||||
contractVersion: "1",
|
||||
adapterVersion: "0.1.0",
|
||||
capabilities: {
|
||||
textToImage: true,
|
||||
imageToImage: true,
|
||||
inpaint: false,
|
||||
maxWidth: MAX_IMAGE_SIDE,
|
||||
maxHeight: MAX_IMAGE_SIDE,
|
||||
textToVideo: true,
|
||||
imageToVideo: true,
|
||||
maxDurationSeconds: MEDIA_VIDEO_MAX_SECONDS,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async gpuStatus(signal?: AbortSignal): Promise<GpuStatus | undefined> {
|
||||
return parseGpuStatus(await this.getJson("system_stats", signal));
|
||||
}
|
||||
|
||||
async generateImage(request: ImageGenerationRequest): Promise<GeneratedMedia[]> {
|
||||
if (request.mask) throw new Error("ComfyUI inpainting is not supported");
|
||||
return this.generate({
|
||||
...request,
|
||||
kind: "image",
|
||||
template: pickTemplate(
|
||||
this.imageWorkflow,
|
||||
Boolean(request.source),
|
||||
DEFAULT_WORKFLOWS.textToImage,
|
||||
DEFAULT_WORKFLOWS.imageToImage,
|
||||
),
|
||||
width: roundSide(request.width, 8),
|
||||
height: roundSide(request.height, 8),
|
||||
count: request.count ?? 1,
|
||||
});
|
||||
}
|
||||
|
||||
async generateVideo(request: VideoGenerationRequest): Promise<GeneratedMedia[]> {
|
||||
const seconds = Math.min(Math.max(request.durationSeconds, 1), MEDIA_VIDEO_MAX_SECONDS);
|
||||
return this.generate({
|
||||
...request,
|
||||
kind: "video",
|
||||
template: pickTemplate(
|
||||
this.videoWorkflow,
|
||||
Boolean(request.source),
|
||||
DEFAULT_WORKFLOWS.textToVideo,
|
||||
DEFAULT_WORKFLOWS.imageToVideo,
|
||||
),
|
||||
width: roundSide(request.width, 32),
|
||||
height: roundSide(request.height, 32),
|
||||
frames: 8 * Math.round((seconds * COMFYUI_VIDEO_FPS) / 8) + 1,
|
||||
count: 1,
|
||||
});
|
||||
}
|
||||
|
||||
private async generate(job: GenerationJob): Promise<GeneratedMedia[]> {
|
||||
const { template, signal } = job;
|
||||
const values: WorkflowValues = {
|
||||
prompt: job.prompt,
|
||||
negative_prompt: job.negativePrompt ?? DEFAULT_NEGATIVE_PROMPTS[job.kind],
|
||||
width: job.width,
|
||||
height: job.height,
|
||||
steps: job.kind === "image" ? IMAGE_STEPS : VIDEO_STEPS,
|
||||
strength: job.strength ?? DEFAULT_STRENGTH,
|
||||
frames: job.frames ?? 1,
|
||||
fps: COMFYUI_VIDEO_FPS,
|
||||
};
|
||||
if (template.placeholders.has("model")) values.model = await this.imageCheckpoint(signal);
|
||||
if (template.placeholders.has("image")) {
|
||||
if (!job.source) throw new Error("ComfyUI workflow needs an input image");
|
||||
values.image = await this.upload(job.source, signal);
|
||||
}
|
||||
const sized = template.placeholders.has("width") && template.placeholders.has("height");
|
||||
const timed = template.placeholders.has("frames") && template.placeholders.has("fps");
|
||||
const firstSeed = job.seed ?? randomInt(2 ** 48 - 1);
|
||||
const results: GeneratedMedia[] = [];
|
||||
for (let index = 0; index < job.count; index += 1) {
|
||||
const graph = fillWorkflow(template, { ...values, seed: firstSeed + index });
|
||||
const outputs = await this.runPrompt(graph, job.kind, signal);
|
||||
const files = outputFiles(outputs, job.kind);
|
||||
if (files.length === 0) throw new Error(`ComfyUI produced no ${job.kind}`);
|
||||
for (const file of files) {
|
||||
results.push({
|
||||
bytes: await this.download(file, signal),
|
||||
mimeType: file.mimeType,
|
||||
name: file.filename,
|
||||
...(sized && file.kind === job.kind ? { width: job.width, height: job.height } : {}),
|
||||
...(timed && file.kind === "video"
|
||||
? { durationMs: Math.round(((values.frames ?? 1) / COMFYUI_VIDEO_FPS) * 1000) }
|
||||
: {}),
|
||||
});
|
||||
}
|
||||
job.onProgress?.((index + 1) / job.count);
|
||||
}
|
||||
return results;
|
||||
}
|
||||
|
||||
/** Submits a prompt and waits for its outputs; a cancelled or timed-out prompt is dequeued. */
|
||||
private async runPrompt(
|
||||
graph: ComfyWorkflow,
|
||||
kind: MediaKind,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<Record<string, unknown>> {
|
||||
const deadline = AbortSignal.timeout(
|
||||
kind === "image" ? IMAGE_PROMPT_TIMEOUT_MS : VIDEO_PROMPT_TIMEOUT_MS,
|
||||
);
|
||||
const runSignal = signal ? AbortSignal.any([signal, deadline]) : deadline;
|
||||
const promptId = await this.submit(graph, runSignal);
|
||||
try {
|
||||
for (;;) {
|
||||
const outputs = await this.finishedOutputs(promptId, runSignal);
|
||||
if (outputs) return outputs;
|
||||
await sleep(this.pollIntervalMs, undefined, { signal: runSignal });
|
||||
}
|
||||
} catch (error) {
|
||||
if (!runSignal.aborted) throw error;
|
||||
await this.cancelPrompt(promptId);
|
||||
if (signal?.aborted) throw signal.reason ?? error;
|
||||
throw new Error("ComfyUI timed out");
|
||||
}
|
||||
}
|
||||
|
||||
private async submit(graph: ComfyWorkflow, signal: AbortSignal): Promise<string> {
|
||||
const payload = { prompt: graph, client_id: this.clientId };
|
||||
return this.postJson("prompt", payload, signal, async (response, readSignal) => {
|
||||
if (response.status === 400) {
|
||||
const rejection = await readJson(response, readSignal).catch(() => undefined);
|
||||
throw new Error(describeWorkflowRejection(rejection));
|
||||
}
|
||||
const result = await readOkJson(response, readSignal);
|
||||
const promptId = isRecord(result) ? result.prompt_id : undefined;
|
||||
if (typeof promptId !== "string" || !promptId) {
|
||||
throw new Error("ComfyUI did not accept the workflow");
|
||||
}
|
||||
return promptId;
|
||||
});
|
||||
}
|
||||
|
||||
private async finishedOutputs(
|
||||
promptId: string,
|
||||
signal: AbortSignal,
|
||||
): Promise<Record<string, unknown> | undefined> {
|
||||
const history = await this.getJson(`history/${encodeURIComponent(promptId)}`, signal);
|
||||
const entry = isRecord(history) ? history[promptId] : undefined;
|
||||
if (!isRecord(entry)) return undefined;
|
||||
const status = isRecord(entry.status) ? entry.status : {};
|
||||
if (status.status_str === "error") throw new Error(describeExecutionError(status.messages));
|
||||
if (status.completed === false) return undefined;
|
||||
return isRecord(entry.outputs) ? entry.outputs : {};
|
||||
}
|
||||
|
||||
/** Best effort: drops the prompt from the queue, or interrupts it when it is ours and running. */
|
||||
private async cancelPrompt(promptId: string): Promise<void> {
|
||||
try {
|
||||
await this.postJson(
|
||||
"queue",
|
||||
{ delete: [promptId] },
|
||||
undefined,
|
||||
discardBody,
|
||||
CANCEL_TIMEOUT_MS,
|
||||
);
|
||||
const queue = await this.getJson("queue", undefined, CANCEL_TIMEOUT_MS);
|
||||
const running =
|
||||
isRecord(queue) && Array.isArray(queue.queue_running) ? queue.queue_running : [];
|
||||
if (running.some((entry) => Array.isArray(entry) && entry[1] === promptId)) {
|
||||
const interrupt = { prompt_id: promptId };
|
||||
await this.postJson("interrupt", interrupt, undefined, discardBody, CANCEL_TIMEOUT_MS);
|
||||
}
|
||||
} catch {
|
||||
// The prompt then finishes on its own; nothing waits for it.
|
||||
}
|
||||
}
|
||||
|
||||
private async imageCheckpoint(signal: AbortSignal | undefined): Promise<string> {
|
||||
if (this.imageModel) return this.imageModel;
|
||||
const info = await this.getJson("object_info/CheckpointLoaderSimple", signal);
|
||||
const videoCheckpoints = new Set(
|
||||
[DEFAULT_WORKFLOWS.textToVideo, DEFAULT_WORKFLOWS.imageToVideo, this.videoWorkflow].flatMap(
|
||||
(template) => (template ? checkpointNames(template.graph) : []),
|
||||
),
|
||||
);
|
||||
const checkpoint = parseCheckpointList(info).find((name) => !videoCheckpoints.has(name));
|
||||
if (!checkpoint) throw new Error("ComfyUI has no image checkpoint installed");
|
||||
return checkpoint;
|
||||
}
|
||||
|
||||
private async upload(file: MediaFile, signal: AbortSignal | undefined): Promise<string> {
|
||||
const form = new FormData();
|
||||
const extension = attachmentExtensionForMimeType(file.mimeType) || ".png";
|
||||
form.append(
|
||||
"image",
|
||||
new Blob([Buffer.from(file.bytes)], { type: file.mimeType }),
|
||||
`evoo-${randomUUID()}${extension}`,
|
||||
);
|
||||
form.append("type", "input");
|
||||
form.append("overwrite", "true");
|
||||
// Serialize here: the address-checked fetch runs on the undici package, whose FormData
|
||||
// class differs from the global one.
|
||||
const encoded = new Response(form);
|
||||
const headers = {
|
||||
Accept: "application/json",
|
||||
"Content-Type": encoded.headers.get("content-type") ?? "",
|
||||
};
|
||||
const body = new Uint8Array(await encoded.arrayBuffer());
|
||||
return this.send("upload/image", { method: "POST", headers, body }, signal, {
|
||||
timeoutMs: TRANSFER_TIMEOUT_MS,
|
||||
read: async (response, readSignal) => {
|
||||
const result = await readOkJson(response, readSignal);
|
||||
const name = isRecord(result) && typeof result.name === "string" ? result.name : "";
|
||||
if (!name) throw new Error("ComfyUI did not accept the image");
|
||||
const subfolder =
|
||||
isRecord(result) && typeof result.subfolder === "string" ? result.subfolder : "";
|
||||
return subfolder ? `${subfolder}/${name}` : name;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private download(file: OutputFile, signal: AbortSignal | undefined): Promise<Uint8Array> {
|
||||
const query = new URLSearchParams({
|
||||
filename: file.filename,
|
||||
subfolder: file.subfolder,
|
||||
type: file.type,
|
||||
});
|
||||
return this.send(`view?${query}`, {}, signal, {
|
||||
timeoutMs: TRANSFER_TIMEOUT_MS,
|
||||
read: async (response, readSignal) => {
|
||||
await assertOk(response);
|
||||
return readBodyCapped(response, ARTIFACT_MAX_BYTES_BY_KIND[file.kind], readSignal);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private getJson(path: string, signal: AbortSignal | undefined, timeoutMs?: number) {
|
||||
const init = { headers: { Accept: "application/json" } };
|
||||
return this.send(path, init, signal, { timeoutMs, read: readOkJson });
|
||||
}
|
||||
|
||||
private postJson<T>(
|
||||
path: string,
|
||||
payload: unknown,
|
||||
signal: AbortSignal | undefined,
|
||||
read: ReadBody<T>,
|
||||
timeoutMs?: number,
|
||||
): Promise<T> {
|
||||
const init = { method: "POST", headers: JSON_HEADERS, body: JSON.stringify(payload) };
|
||||
return this.send(path, init, signal, { timeoutMs, read });
|
||||
}
|
||||
|
||||
/** One request, with the timeout covering the response body too. */
|
||||
private async send<T>(
|
||||
path: string,
|
||||
init: {
|
||||
method?: string;
|
||||
headers?: Record<string, string>;
|
||||
body?: string | Uint8Array<ArrayBuffer>;
|
||||
},
|
||||
signal: AbortSignal | undefined,
|
||||
options: { timeoutMs?: number; read: ReadBody<T> },
|
||||
): Promise<T> {
|
||||
const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS);
|
||||
const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout;
|
||||
const headers: Record<string, string> = { ...init.headers };
|
||||
if (this.apiKey) headers.Authorization = `Bearer ${this.apiKey}`;
|
||||
try {
|
||||
const response = await this.fetch(new URL(path, this.baseUrl).href, {
|
||||
method: init.method ?? "GET",
|
||||
headers,
|
||||
body: init.body,
|
||||
redirect: "error",
|
||||
signal: requestSignal,
|
||||
});
|
||||
return await options.read(response, requestSignal);
|
||||
} catch (error) {
|
||||
if (signal?.aborted) throw error;
|
||||
if (timeout.aborted) throw new Error("ComfyUI timed out");
|
||||
// undici reports refused or reset connections as a bare "fetch failed".
|
||||
if (error instanceof TypeError && error.message === "fetch failed") {
|
||||
throw new Error("ComfyUI is unreachable");
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function comfyUiBaseUrl(raw: string): URL {
|
||||
const url = assertAllowedOpenAiCompatibleRequestUrl(raw.trim());
|
||||
url.search = "";
|
||||
url.hash = "";
|
||||
url.pathname = `${url.pathname.replace(/\/+$/, "")}/`;
|
||||
return url;
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return Boolean(value) && typeof value === "object" && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function truncate(text: string): string {
|
||||
return text.length > MAX_ERROR_CHARS ? `${text.slice(0, MAX_ERROR_CHARS - 1)}…` : text;
|
||||
}
|
||||
|
||||
function roundSide(side: number, multiple: number): number {
|
||||
const rounded = Math.round(side / multiple) * multiple;
|
||||
return Math.min(Math.max(rounded, multiple * Math.ceil(64 / multiple)), MAX_IMAGE_SIDE);
|
||||
}
|
||||
|
||||
/** A custom workflow replaces the default that matches whether it takes an input image. */
|
||||
function pickTemplate(
|
||||
custom: WorkflowTemplate | undefined,
|
||||
withImage: boolean,
|
||||
textDefault: WorkflowTemplate,
|
||||
imageDefault: WorkflowTemplate,
|
||||
): WorkflowTemplate {
|
||||
if (custom && custom.placeholders.has("image") === withImage) return custom;
|
||||
return withImage ? imageDefault : textDefault;
|
||||
}
|
||||
|
||||
function checkpointNames(graph: ComfyWorkflow): string[] {
|
||||
return Object.values(graph).flatMap((node) =>
|
||||
typeof node.inputs.ckpt_name === "string" ? [node.inputs.ckpt_name] : [],
|
||||
);
|
||||
}
|
||||
|
||||
async function discardBody(response: Response): Promise<void> {
|
||||
await response.body?.cancel().catch(() => undefined);
|
||||
}
|
||||
|
||||
async function assertOk(response: Response): Promise<void> {
|
||||
if (response.ok) return;
|
||||
await discardBody(response);
|
||||
throw new Error(`ComfyUI returned ${response.status}`);
|
||||
}
|
||||
|
||||
async function readJson(response: Response, signal: AbortSignal): Promise<unknown> {
|
||||
const bytes = await readBodyCapped(response, MAX_JSON_BYTES, signal);
|
||||
try {
|
||||
return JSON.parse(new TextDecoder().decode(bytes));
|
||||
} catch {
|
||||
throw new Error("ComfyUI returned invalid JSON");
|
||||
}
|
||||
}
|
||||
|
||||
async function readOkJson(response: Response, signal: AbortSignal): Promise<unknown> {
|
||||
await assertOk(response);
|
||||
return readJson(response, signal);
|
||||
}
|
||||
|
||||
/**
|
||||
* ComfyUI can evict its own models, so memory its PyTorch process holds counts as available;
|
||||
* only memory other processes use (another model server on the same GPU) does not.
|
||||
*/
|
||||
export function parseGpuStatus(body: unknown): GpuStatus | undefined {
|
||||
const devices = isRecord(body) && Array.isArray(body.devices) ? body.devices : [];
|
||||
for (const device of devices) {
|
||||
if (!isRecord(device) || device.type === "cpu") continue;
|
||||
const total = finiteNumber(device.vram_total);
|
||||
const free = finiteNumber(device.vram_free);
|
||||
if (!total || free === undefined) continue;
|
||||
const reserved = finiteNumber(device.torch_vram_total) ?? 0;
|
||||
const reservedFree = finiteNumber(device.torch_vram_free) ?? 0;
|
||||
const available = Math.min(total, Math.max(0, free - reservedFree + reserved));
|
||||
return { vramFreeMb: Math.floor(available / MIB), vramTotalMb: Math.floor(total / MIB) };
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function finiteNumber(value: unknown): number | undefined {
|
||||
return typeof value === "number" && Number.isFinite(value) ? value : undefined;
|
||||
}
|
||||
|
||||
/** Older ComfyUI lists combo values as `[[...names], {...}]`, newer as `["COMBO", {options}]`. */
|
||||
export function parseCheckpointList(body: unknown): string[] {
|
||||
const loader = isRecord(body) ? body.CheckpointLoaderSimple : undefined;
|
||||
const input = isRecord(loader) ? loader.input : undefined;
|
||||
const required = isRecord(input) ? input.required : undefined;
|
||||
const spec = isRecord(required) ? required.ckpt_name : undefined;
|
||||
if (!Array.isArray(spec)) return [];
|
||||
const names = Array.isArray(spec[0])
|
||||
? spec[0]
|
||||
: isRecord(spec[1]) && Array.isArray(spec[1].options)
|
||||
? spec[1].options
|
||||
: [];
|
||||
return names.filter((name): name is string => typeof name === "string" && name.length > 0);
|
||||
}
|
||||
|
||||
/** Saved files win over previews; video runs return clips (or animations), not frames. */
|
||||
function outputFiles(outputs: Record<string, unknown>, kind: MediaKind): OutputFile[] {
|
||||
const files: OutputFile[] = [];
|
||||
for (const node of Object.values(outputs)) {
|
||||
if (!isRecord(node)) continue;
|
||||
for (const key of ["images", "gifs", "videos"]) {
|
||||
const entries = node[key];
|
||||
if (!Array.isArray(entries)) continue;
|
||||
for (const entry of entries) {
|
||||
if (!isRecord(entry) || typeof entry.filename !== "string" || !entry.filename) continue;
|
||||
const mimeType = inferAttachmentMimeType(entry.filename);
|
||||
if (!mimeType) continue;
|
||||
files.push({
|
||||
filename: entry.filename,
|
||||
subfolder: typeof entry.subfolder === "string" ? entry.subfolder : "",
|
||||
type: typeof entry.type === "string" ? entry.type : "output",
|
||||
mimeType,
|
||||
kind: attachmentKindOf(mimeType),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
const videos = files.filter((file) => file.kind === "video");
|
||||
const wanted =
|
||||
kind === "image"
|
||||
? files.filter((file) => file.kind === "image")
|
||||
: videos.length > 0
|
||||
? videos
|
||||
: files.filter((file) => file.mimeType === "image/gif" || file.mimeType === "image/webp");
|
||||
const saved = wanted.filter((file) => file.type === "output");
|
||||
return saved.length > 0 ? saved : wanted;
|
||||
}
|
||||
|
||||
export function describeWorkflowRejection(body: unknown): string {
|
||||
const problems: string[] = [];
|
||||
const nodeErrors = isRecord(body) && isRecord(body.node_errors) ? body.node_errors : {};
|
||||
for (const node of Object.values(nodeErrors)) {
|
||||
if (!isRecord(node) || !Array.isArray(node.errors)) continue;
|
||||
for (const error of node.errors) {
|
||||
if (!isRecord(error)) continue;
|
||||
const parts = [node.class_type, error.message, error.details];
|
||||
problems.push(parts.filter((part) => typeof part === "string" && part).join(": "));
|
||||
}
|
||||
}
|
||||
if (problems.length === 0 && isRecord(body) && isRecord(body.error)) {
|
||||
if (typeof body.error.message === "string") problems.push(body.error.message);
|
||||
}
|
||||
return truncate(
|
||||
problems.length > 0
|
||||
? `ComfyUI rejected the workflow: ${problems.join("; ")}`
|
||||
: "ComfyUI rejected the workflow",
|
||||
);
|
||||
}
|
||||
|
||||
export function describeExecutionError(messages: unknown): string {
|
||||
const events = Array.isArray(messages) ? messages : [];
|
||||
for (const event of events) {
|
||||
if (!Array.isArray(event)) continue;
|
||||
const [name, data] = event;
|
||||
if (name === "execution_interrupted") return "ComfyUI interrupted the run";
|
||||
if (name !== "execution_error" || !isRecord(data)) continue;
|
||||
const node = typeof data.node_type === "string" ? data.node_type : "a node";
|
||||
const message =
|
||||
typeof data.exception_message === "string"
|
||||
? data.exception_message.trim().split("\n")[0]
|
||||
: "";
|
||||
return truncate(`ComfyUI failed in ${node}${message ? `: ${message}` : ""}`);
|
||||
}
|
||||
return "ComfyUI failed";
|
||||
}
|
||||
@ -1,6 +1,7 @@
|
||||
import type { ConnectorTool } from "@rakazo/adapter-kit";
|
||||
import type { ConnectorTool, MediaCapabilities } from "@rakazo/adapter-kit";
|
||||
import { agentConnectionTools, builtinAgentTools } from "./builtin-tools.js";
|
||||
import { selectCloudAgentTools } from "./cloud-agent-tools-select.js";
|
||||
import { selectMediaTools } from "./media-tools.js";
|
||||
import { selectMemoryTools } from "./memory-tools.js";
|
||||
import { filterImageReturningComputerTools } from "./model-vision.js";
|
||||
import {
|
||||
@ -81,6 +82,8 @@ export function selectBuiltinToolsForRun(options: {
|
||||
trigger: string;
|
||||
semanticMemoryEnabled: boolean;
|
||||
cloudAgentEnabled?: boolean;
|
||||
/** Media tools appear only for operations the space can run. */
|
||||
mediaCapabilities?: MediaCapabilities;
|
||||
messagingChannelRun: boolean;
|
||||
/** Cross-owner agent connections only exist for chat-linked bots. */
|
||||
agentConnectionsEnabled?: boolean;
|
||||
@ -88,21 +91,24 @@ export function selectBuiltinToolsForRun(options: {
|
||||
foreignAgentRun?: boolean;
|
||||
bot: BuiltinToolPolicy;
|
||||
}) {
|
||||
const selected = selectCloudAgentTools(
|
||||
selectMemoryTools(
|
||||
filterBuiltinToolsForRun(
|
||||
filterBuiltinToolsForThread(
|
||||
filterPageBrowserTools(
|
||||
filterImageReturningComputerTools(builtinAgentTools, options.graphicalToolsAllowed),
|
||||
options.pageBrowserAllowed ?? options.graphicalToolsAllowed,
|
||||
const selected = selectMediaTools(
|
||||
selectCloudAgentTools(
|
||||
selectMemoryTools(
|
||||
filterBuiltinToolsForRun(
|
||||
filterBuiltinToolsForThread(
|
||||
filterPageBrowserTools(
|
||||
filterImageReturningComputerTools(builtinAgentTools, options.graphicalToolsAllowed),
|
||||
options.pageBrowserAllowed ?? options.graphicalToolsAllowed,
|
||||
),
|
||||
options.groupId,
|
||||
),
|
||||
options.groupId,
|
||||
options.trigger,
|
||||
),
|
||||
options.trigger,
|
||||
options.semanticMemoryEnabled,
|
||||
),
|
||||
options.semanticMemoryEnabled,
|
||||
Boolean(options.cloudAgentEnabled),
|
||||
),
|
||||
Boolean(options.cloudAgentEnabled),
|
||||
options.mediaCapabilities,
|
||||
).filter(
|
||||
(tool) =>
|
||||
!(options.foreignAgentRun && tool.name === "update_user_profile") &&
|
||||
|
||||
@ -283,6 +283,7 @@ import {
|
||||
needsOAuthProbe,
|
||||
parseMcpServerToolArgs,
|
||||
} from "./mcp-server-tool.js";
|
||||
import { executeMediaTool, isMediaToolName } from "./media-tools.js";
|
||||
import { loadAgentMemoryContext } from "./memory-context.js";
|
||||
import type { MemoryProviderResolver } from "./memory-provider-factory.js";
|
||||
import { selectMemoryTools } from "./memory-tools.js";
|
||||
@ -1478,6 +1479,7 @@ export function createRunExecutor(deps: ExecutorDeps) {
|
||||
trigger: run.trigger,
|
||||
semanticMemoryEnabled,
|
||||
cloudAgentEnabled: cloudAgentsEnabled(cloudAgent, run.spaceId),
|
||||
mediaCapabilities: await deps.media?.capabilities(run.spaceId).catch(() => undefined),
|
||||
messagingChannelRun,
|
||||
agentConnectionsEnabled: hasMessagingIdentity,
|
||||
foreignAgentRun,
|
||||
@ -2599,6 +2601,24 @@ export function createRunExecutor(deps: ExecutorDeps) {
|
||||
return computerScreenToolResult(() => tool(browser, computer, context, args), finish);
|
||||
}
|
||||
|
||||
if (isMediaToolName(name)) {
|
||||
return finish(
|
||||
await executeMediaTool(
|
||||
deps.media,
|
||||
{
|
||||
spaceId: run.spaceId,
|
||||
userId: run.userId,
|
||||
botId: bot.id,
|
||||
groupId: thread.groupId,
|
||||
threadId: run.threadId,
|
||||
runId: run.id,
|
||||
operationKey: effectKey,
|
||||
},
|
||||
name,
|
||||
args,
|
||||
),
|
||||
);
|
||||
}
|
||||
if (name.startsWith("cloud_agent_")) {
|
||||
return finish(
|
||||
await executeCloudAgentTool(
|
||||
|
||||
111
packages/adapters/src/image-generation-conformance.test.ts
Normal file
111
packages/adapters/src/image-generation-conformance.test.ts
Normal file
@ -0,0 +1,111 @@
|
||||
import type {
|
||||
ImageGenerationProvider,
|
||||
MediaFile,
|
||||
VideoGenerationProvider,
|
||||
} from "@rakazo/adapter-kit";
|
||||
import { attachmentKindOf } from "@rakazo/contracts";
|
||||
import sharp from "sharp";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { createComfyUiProvider } from "./comfyui-provider.js";
|
||||
import { ComfyUiEmulator } from "./testing/comfyui-emulator.js";
|
||||
|
||||
type MediaGenerationProvider = ImageGenerationProvider & VideoGenerationProvider;
|
||||
|
||||
const factories: Record<string, () => MediaGenerationProvider> = {
|
||||
comfyui: () =>
|
||||
createComfyUiProvider({
|
||||
baseUrl: "http://127.0.0.1:8188",
|
||||
fetch: new ComfyUiEmulator().fetch,
|
||||
pollIntervalMs: 1,
|
||||
}),
|
||||
};
|
||||
|
||||
async function sourceImage(): Promise<MediaFile> {
|
||||
const png = await sharp({
|
||||
create: { width: 96, height: 64, channels: 3, background: { r: 10, g: 20, b: 30 } },
|
||||
})
|
||||
.png()
|
||||
.toBuffer();
|
||||
return { bytes: new Uint8Array(png), mimeType: "image/png", name: "source.png" };
|
||||
}
|
||||
|
||||
/**
|
||||
* Offline conformance: every image/video generation provider describes its limits, returns
|
||||
* the requested number of decodable images, is reproducible by seed, takes a source image,
|
||||
* reports progress, returns video, and honors cancellation.
|
||||
*/
|
||||
for (const [name, create] of Object.entries(factories)) {
|
||||
describe(`${name} media generation conformance (offline)`, () => {
|
||||
it("describes the contract and its limits", () => {
|
||||
const image: ImageGenerationProvider = create();
|
||||
const video: VideoGenerationProvider = create();
|
||||
expect(image.describe().contractVersion).toBe("1");
|
||||
expect(image.describe().capabilities.maxWidth).toBeGreaterThanOrEqual(1024);
|
||||
expect(image.describe().capabilities.maxHeight).toBeGreaterThanOrEqual(1024);
|
||||
expect(video.describe().contractVersion).toBe("1");
|
||||
expect(video.describe().capabilities.maxDurationSeconds).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("returns the requested number of images, reproducible by seed", async () => {
|
||||
const provider = create();
|
||||
const progress: number[] = [];
|
||||
const request = { prompt: "a red kite", width: 256, height: 192, count: 2, seed: 7 };
|
||||
const images = await provider.generateImage({
|
||||
...request,
|
||||
onProgress: (fraction) => progress.push(fraction),
|
||||
});
|
||||
expect(images).toHaveLength(2);
|
||||
for (const image of images) {
|
||||
expect(attachmentKindOf(image.mimeType)).toBe("image");
|
||||
const metadata = await sharp(image.bytes).metadata();
|
||||
expect([metadata.width, metadata.height]).toEqual([256, 192]);
|
||||
}
|
||||
expect(Buffer.from(images[0]!.bytes).equals(Buffer.from(images[1]!.bytes))).toBe(false);
|
||||
expect(progress.at(-1)).toBe(1);
|
||||
const again = await provider.generateImage({ ...request, count: 1 });
|
||||
expect(Buffer.from(again[0]!.bytes).equals(Buffer.from(images[0]!.bytes))).toBe(true);
|
||||
});
|
||||
|
||||
it("generates from a source image", async () => {
|
||||
const images = await create().generateImage({
|
||||
prompt: "the same scene at night",
|
||||
width: 96,
|
||||
height: 64,
|
||||
source: await sourceImage(),
|
||||
strength: 0.5,
|
||||
});
|
||||
expect(images).toHaveLength(1);
|
||||
expect(attachmentKindOf(images[0]!.mimeType)).toBe("image");
|
||||
});
|
||||
|
||||
it("returns video from text and from a start frame", async () => {
|
||||
const provider = create();
|
||||
for (const source of [undefined, await sourceImage()]) {
|
||||
const [clip, ...rest] = await provider.generateVideo({
|
||||
prompt: "waves rolling onto a beach",
|
||||
width: 768,
|
||||
height: 512,
|
||||
durationSeconds: 2,
|
||||
source,
|
||||
});
|
||||
expect(rest).toHaveLength(0);
|
||||
expect(attachmentKindOf(clip!.mimeType)).toBe("video");
|
||||
expect(clip!.bytes.byteLength).toBeGreaterThan(0);
|
||||
}
|
||||
});
|
||||
|
||||
it("reports GPU memory with free not above total", async () => {
|
||||
const status = await create().gpuStatus?.();
|
||||
if (!status) return;
|
||||
expect(status.vramTotalMb).toBeGreaterThan(0);
|
||||
expect(status.vramFreeMb).toBeGreaterThanOrEqual(0);
|
||||
expect(status.vramFreeMb).toBeLessThanOrEqual(status.vramTotalMb);
|
||||
});
|
||||
|
||||
it("rejects when the caller cancels", async () => {
|
||||
const request = { prompt: "x", width: 64, height: 64, signal: AbortSignal.abort() };
|
||||
await expect(create().generateImage(request)).rejects.toThrow();
|
||||
await expect(create().generateVideo({ ...request, durationSeconds: 1 })).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
}
|
||||
@ -24,6 +24,7 @@ export * from "./cloud-agent-provider-env.js";
|
||||
export * from "./cloud-agent-service.js";
|
||||
export * from "./cloud-agent-tools.js";
|
||||
export * from "./cloud-agent-tools-select.js";
|
||||
export * from "./comfyui-provider.js";
|
||||
export * from "./composio-catalog-cache.js";
|
||||
export * from "./composio-connector.js";
|
||||
export * from "./composio-emulator.js";
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import type { MediaSettings } from "@rakazo/contracts";
|
||||
import { createImageOperations } from "./media/image-operations.js";
|
||||
import { createTranscriptionOperations } from "./media/transcription.js";
|
||||
import type { MediaOperationRegistration } from "./media-runtime.js";
|
||||
|
||||
@ -12,5 +13,8 @@ export interface MediaOperationDeps {
|
||||
* from this one list, so both processes run the same operations.
|
||||
*/
|
||||
export function createMediaOperations(deps: MediaOperationDeps): MediaOperationRegistration[] {
|
||||
return [...createTranscriptionOperations({ loadSettings: deps.loadSettings })];
|
||||
return [
|
||||
...createImageOperations(deps),
|
||||
...createTranscriptionOperations({ loadSettings: deps.loadSettings }),
|
||||
];
|
||||
}
|
||||
|
||||
615
packages/adapters/src/media-runtime.postgres.test.ts
Normal file
615
packages/adapters/src/media-runtime.postgres.test.ts
Normal file
@ -0,0 +1,615 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { existsSync } from "node:fs";
|
||||
import { mkdtemp, readdir, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import path from "node:path";
|
||||
import type {
|
||||
BackgroundJob,
|
||||
JobPublisher,
|
||||
MediaJobContext,
|
||||
MediaOperationResult,
|
||||
} 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 { 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";
|
||||
|
||||
const describePostgres =
|
||||
process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip;
|
||||
|
||||
const PNG = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3, 4]);
|
||||
|
||||
type FakeInput = { prompt: string; source?: string };
|
||||
type Run = (context: MediaJobContext) => Promise<MediaOperationResult>;
|
||||
|
||||
function deferred<T = void>() {
|
||||
let resolve!: (value: T) => void;
|
||||
const promise = new Promise<T>((done) => {
|
||||
resolve = done;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
describePostgres("media job runtime (PostgreSQL, fake operations)", () => {
|
||||
let prisma: PrismaClient;
|
||||
let db: ReturnType<typeof createDb>;
|
||||
let root: string;
|
||||
const users: string[] = [];
|
||||
|
||||
beforeAll(async () => {
|
||||
db = createDb(process.env.DATABASE_URL!);
|
||||
prisma = db.prisma;
|
||||
root = await mkdtemp(path.join(tmpdir(), "media-runtime-test-"));
|
||||
});
|
||||
afterAll(async () => {
|
||||
if (!db) return;
|
||||
await prisma.organization.deleteMany({ where: { id: { in: users } } });
|
||||
await prisma.user.deleteMany({ where: { id: { in: users } } });
|
||||
await prisma.$disconnect();
|
||||
await db.pool.end();
|
||||
await rm(root, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
async function setup(
|
||||
options: {
|
||||
run?: Run;
|
||||
available?: () => Promise<boolean>;
|
||||
gpu?: MediaOperationRegistration["gpu"];
|
||||
runtime?: MediaRuntimeOptions;
|
||||
} = {},
|
||||
) {
|
||||
const id = randomUUID();
|
||||
users.push(id);
|
||||
await prisma.user.create({
|
||||
data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true },
|
||||
});
|
||||
await prisma.organization.create({
|
||||
data: { id, name: "Media test", slug: id, createdAt: new Date() },
|
||||
});
|
||||
await prisma.space.create({
|
||||
data: { id, organizationId: id, name: "Media test", isDefault: true },
|
||||
});
|
||||
await prisma.member.create({
|
||||
data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() },
|
||||
});
|
||||
const bot = await prisma.bot.create({
|
||||
data: { spaceId: id, userId: id, name: "Studio", color: "test-color" },
|
||||
});
|
||||
const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } });
|
||||
const task = await prisma.task.create({
|
||||
data: {
|
||||
spaceId: id,
|
||||
userId: id,
|
||||
botId: bot.id,
|
||||
threadId: thread.id,
|
||||
prompt: "Make a picture",
|
||||
status: "running",
|
||||
},
|
||||
});
|
||||
const run = await prisma.run.create({
|
||||
data: {
|
||||
spaceId: id,
|
||||
userId: id,
|
||||
botId: bot.id,
|
||||
threadId: thread.id,
|
||||
taskId: task.id,
|
||||
status: "running",
|
||||
trigger: "user",
|
||||
},
|
||||
});
|
||||
const artifacts = new LocalArtifactStore(path.join(root, id));
|
||||
const enqueued: BackgroundJob[] = [];
|
||||
const cancelled: string[] = [];
|
||||
const jobs: JobPublisher = {
|
||||
enqueue: async (job) => {
|
||||
enqueued.push(job);
|
||||
},
|
||||
cancel: async (key) => {
|
||||
cancelled.push(key);
|
||||
},
|
||||
close: async () => undefined,
|
||||
};
|
||||
const notify = vi.fn(async () => undefined);
|
||||
const handlerCalls: string[] = [];
|
||||
const defaultRun: Run = async (context) => {
|
||||
const input = context.input as FakeInput;
|
||||
if (input.source) await context.fetchArtifact(input.source);
|
||||
await context.progress(0.5);
|
||||
const file = path.join(context.tempDir, "out.png");
|
||||
const poster = path.join(context.tempDir, "poster.jpg");
|
||||
await writeFile(file, PNG);
|
||||
await writeFile(poster, PNG);
|
||||
await context.addOutput({
|
||||
path: file,
|
||||
name: "picture.png",
|
||||
mimeType: "image/png",
|
||||
width: 64,
|
||||
height: 32,
|
||||
posterPath: poster,
|
||||
parentArtifactId: input.source,
|
||||
});
|
||||
return { summary: `Drew ${input.prompt}.` };
|
||||
};
|
||||
const registration: MediaOperationRegistration<FakeInput> = {
|
||||
operation: "generate_image",
|
||||
inputSchema: z.object({ prompt: z.string().min(1), source: z.string().optional() }),
|
||||
title: (input) => input.prompt,
|
||||
available: options.available ?? (async () => true),
|
||||
artifactIds: (input) => (input.source ? [input.source] : []),
|
||||
handler: {
|
||||
operation: "generate_image",
|
||||
run: async (context) => {
|
||||
handlerCalls.push(context.jobId);
|
||||
return (options.run ?? defaultRun)(context);
|
||||
},
|
||||
},
|
||||
...(options.gpu ? { gpu: options.gpu } : {}),
|
||||
};
|
||||
const tempRoot = path.join(root, id, "tmp");
|
||||
const runtimeDeps = {
|
||||
prisma,
|
||||
events: { notify },
|
||||
jobs,
|
||||
artifacts,
|
||||
registrations: [registration as MediaOperationRegistration],
|
||||
};
|
||||
const runtimeOptions: MediaRuntimeOptions = {
|
||||
tempRoot,
|
||||
heartbeatMs: 50,
|
||||
progressIntervalMs: 0,
|
||||
...options.runtime,
|
||||
};
|
||||
const runtime = createMediaRuntime(runtimeDeps, runtimeOptions);
|
||||
const start = (input: Record<string, unknown>, extra: Record<string, unknown> = {}) =>
|
||||
runtime.start({
|
||||
operation: "generate_image",
|
||||
input,
|
||||
spaceId: id,
|
||||
userId: id,
|
||||
botId: bot.id,
|
||||
groupId: null,
|
||||
threadId: thread.id,
|
||||
runId: run.id,
|
||||
...extra,
|
||||
});
|
||||
const job = (jobId: string) => prisma.mediaJob.findUniqueOrThrow({ where: { id: jobId } });
|
||||
const card = async (jobId: string): Promise<MediaJobBlock> => {
|
||||
const row = await job(jobId);
|
||||
const message = await prisma.message.findUniqueOrThrow({ where: { id: row.messageId! } });
|
||||
const block = (message.blocks as MessageBlock[]).find(
|
||||
(entry): entry is MediaJobBlock => entry.kind === "media",
|
||||
);
|
||||
if (!block) throw new Error("media card missing");
|
||||
return block;
|
||||
};
|
||||
const sourceArtifact = async () => {
|
||||
const file = path.join(root, `${randomUUID()}.png`);
|
||||
await writeFile(file, PNG);
|
||||
return storeArtifactFromFile(
|
||||
{ prisma, artifacts },
|
||||
{
|
||||
path: file,
|
||||
name: "source.png",
|
||||
mimeType: "image/png",
|
||||
spaceId: id,
|
||||
userId: id,
|
||||
botId: bot.id,
|
||||
groupId: null,
|
||||
runId: null,
|
||||
},
|
||||
);
|
||||
};
|
||||
return {
|
||||
id,
|
||||
bot,
|
||||
thread,
|
||||
run,
|
||||
runtime,
|
||||
runtimeDeps,
|
||||
runtimeOptions,
|
||||
enqueued,
|
||||
cancelled,
|
||||
notify,
|
||||
handlerCalls,
|
||||
tempRoot,
|
||||
start,
|
||||
job,
|
||||
card,
|
||||
sourceArtifact,
|
||||
};
|
||||
}
|
||||
|
||||
it("runs a job to done, stores outputs with posters and versions, and cleans its temp dir", async () => {
|
||||
const ctx = await setup();
|
||||
const source = await ctx.sourceArtifact();
|
||||
const started = await ctx.start({ prompt: "a red fox", source: source.id });
|
||||
|
||||
const queued = await ctx.card(started.jobId);
|
||||
expect(queued).toMatchObject({ status: "queued", title: "a red fox", outputs: [] });
|
||||
expect(ctx.enqueued.map((job) => job.name)).toEqual(["media.run"]);
|
||||
expect(ctx.enqueued[0]).toMatchObject({ payload: { jobId: started.jobId }, queue: "media" });
|
||||
const message = await prisma.message.findUniqueOrThrow({ where: { id: started.messageId } });
|
||||
expect(message).toMatchObject({ role: "bot", botId: ctx.bot.id, runId: ctx.run.id });
|
||||
|
||||
await ctx.runtime.runJob(started.jobId);
|
||||
|
||||
const done = await ctx.job(started.jobId);
|
||||
expect(done).toMatchObject({ status: "done", progress: 1, attempts: 1, leaseToken: null });
|
||||
expect(done.finishedAt).toBeInstanceOf(Date);
|
||||
const block = await ctx.card(started.jobId);
|
||||
expect(block.status).toBe("done");
|
||||
expect(block.outputs).toHaveLength(1);
|
||||
const [output] = block.outputs;
|
||||
expect(output).toMatchObject({
|
||||
kind: "image",
|
||||
mimeType: "image/png",
|
||||
name: "picture.png",
|
||||
size: PNG.byteLength,
|
||||
width: 64,
|
||||
height: 32,
|
||||
parentArtifactId: source.id,
|
||||
});
|
||||
expect(done.outputArtifactIds).toEqual([output!.posterArtifactId, output!.artifactId]);
|
||||
const stored = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } });
|
||||
expect(stored).toMatchObject({
|
||||
spaceId: ctx.id,
|
||||
userId: ctx.id,
|
||||
botId: ctx.bot.id,
|
||||
runId: ctx.run.id,
|
||||
parentArtifactId: source.id,
|
||||
width: 64,
|
||||
height: 32,
|
||||
});
|
||||
const events = await prisma.event.findMany({
|
||||
where: { threadId: ctx.thread.id, type: "thread.message.updated" },
|
||||
orderBy: { seq: "asc" },
|
||||
});
|
||||
const statuses = events.map(
|
||||
(event) => ((event.payload as { blocks: MessageBlock[] }).blocks[0] as MediaJobBlock).status,
|
||||
);
|
||||
expect(statuses[0]).toBe("running");
|
||||
expect(statuses.at(-1)).toBe("done");
|
||||
expect(ctx.notify).toHaveBeenCalled();
|
||||
expect(existsSync(path.join(ctx.tempRoot, started.jobId))).toBe(false);
|
||||
// No wake without continue_after.
|
||||
expect(
|
||||
await prisma.run.count({ where: { trigger: MEDIA_WAKE_TRIGGER, spaceId: ctx.id } }),
|
||||
).toBe(0);
|
||||
expect(await ctx.runtime.get(started.jobId)).toMatchObject({
|
||||
id: started.jobId,
|
||||
status: "done",
|
||||
outputs: block.outputs,
|
||||
});
|
||||
});
|
||||
|
||||
it("validates input, availability and readable artifacts before creating a job", async () => {
|
||||
const ctx = await setup();
|
||||
await expect(ctx.start({ prompt: "" })).rejects.toBeInstanceOf(MediaJobRejectedError);
|
||||
await expect(ctx.start({ prompt: "x", source: "not-an-artifact" })).rejects.toThrow(
|
||||
"Unknown artifact.",
|
||||
);
|
||||
const other = await setup();
|
||||
const foreign = await other.sourceArtifact();
|
||||
await expect(ctx.start({ prompt: "x", source: foreign.id })).rejects.toThrow(
|
||||
"Unknown artifact.",
|
||||
);
|
||||
await expect(ctx.start({ prompt: "x" }, { threadId: other.thread.id })).rejects.toThrow(
|
||||
"This conversation is no longer available.",
|
||||
);
|
||||
const unavailable = await setup({ available: async () => false });
|
||||
await expect(unavailable.start({ prompt: "x" })).rejects.toThrow("not available");
|
||||
expect(
|
||||
await prisma.mediaJob.count({ where: { spaceId: { in: [ctx.id, unavailable.id] } } }),
|
||||
).toBe(0);
|
||||
});
|
||||
|
||||
it("returns the first job for a replayed tool call", async () => {
|
||||
const ctx = await setup();
|
||||
const first = await ctx.start({ prompt: "fox" }, { operationKey: `effect:${ctx.id}` });
|
||||
const replay = await ctx.start({ prompt: "fox" }, { operationKey: `effect:${ctx.id}` });
|
||||
expect(replay).toEqual(first);
|
||||
expect(await prisma.mediaJob.count({ where: { spaceId: ctx.id } })).toBe(1);
|
||||
});
|
||||
|
||||
it("fails a job with the operation's error and still cleans up", async () => {
|
||||
let tempDir = "";
|
||||
const ctx = await setup({
|
||||
run: async (context) => {
|
||||
tempDir = context.tempDir;
|
||||
await writeFile(path.join(context.tempDir, "partial.bin"), PNG);
|
||||
throw new Error(`ComfyUI rejected the workflow (${context.tempDir}/x)`);
|
||||
},
|
||||
});
|
||||
const { jobId } = await ctx.start({ prompt: "fox" });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
const failed = await ctx.job(jobId);
|
||||
expect(failed.status).toBe("failed");
|
||||
expect(failed.error).toBe("ComfyUI rejected the workflow (/x)");
|
||||
expect((await ctx.card(jobId)).error).toBe(failed.error);
|
||||
expect(existsSync(tempDir)).toBe(false);
|
||||
});
|
||||
|
||||
it("refuses outputs outside the temp dir and sources of another user", async () => {
|
||||
const other = await setup();
|
||||
const foreign = await other.sourceArtifact();
|
||||
const ctx = await setup({
|
||||
run: async (context) => {
|
||||
await context.fetchArtifact(foreign.id);
|
||||
return { summary: "" };
|
||||
},
|
||||
});
|
||||
const { jobId } = await ctx.start({ prompt: "fox" });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(await ctx.job(jobId)).toMatchObject({
|
||||
status: "failed",
|
||||
error: "A source file is no longer available.",
|
||||
});
|
||||
|
||||
const outside = await setup({
|
||||
run: async (context) => {
|
||||
await context.addOutput({ path: "/etc/hostname", name: "x.png", mimeType: "image/png" });
|
||||
return { summary: "" };
|
||||
},
|
||||
});
|
||||
const escaped = await outside.start({ prompt: "fox" });
|
||||
await outside.runtime.runJob(escaped.jobId);
|
||||
expect((await outside.job(escaped.jobId)).status).toBe("failed");
|
||||
expect(await prisma.artifact.count({ where: { spaceId: outside.id } })).toBe(0);
|
||||
});
|
||||
|
||||
it("defers with backoff, then fails once the deferral budget is spent", async () => {
|
||||
const ctx = await setup({
|
||||
run: async () => {
|
||||
throw new MediaJobDeferredError("ComfyUI queue is full", 2_000);
|
||||
},
|
||||
runtime: { maxDeferrals: 2 },
|
||||
});
|
||||
const { jobId } = await ctx.start({ prompt: "fox" });
|
||||
const before = Date.now();
|
||||
await ctx.runtime.runJob(jobId);
|
||||
const waiting = await ctx.job(jobId);
|
||||
expect(waiting).toMatchObject({ status: "waiting", deferrals: 1, leaseToken: null });
|
||||
expect(waiting.nextAttemptAt!.getTime()).toBeGreaterThanOrEqual(before + 2_000);
|
||||
expect((await ctx.card(jobId)).status).toBe("waiting");
|
||||
const retry = ctx.enqueued.at(-1)!;
|
||||
expect(retry).toMatchObject({ name: "media.run", payload: { jobId }, queue: "media" });
|
||||
expect(retry.availableAt?.getTime()).toBe(waiting.nextAttemptAt!.getTime());
|
||||
|
||||
// Too early: the delayed job owns the next attempt.
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(ctx.handlerCalls).toHaveLength(1);
|
||||
|
||||
await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
const second = await ctx.job(jobId);
|
||||
expect(second).toMatchObject({ status: "waiting", deferrals: 2 });
|
||||
// Backoff doubles.
|
||||
expect(second.nextAttemptAt!.getTime() - Date.now()).toBeGreaterThan(3_000);
|
||||
|
||||
await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(await ctx.job(jobId)).toMatchObject({
|
||||
status: "failed",
|
||||
error: "The media server stayed busy. Try again later.",
|
||||
});
|
||||
});
|
||||
|
||||
it("waits for free GPU memory before running the handler", async () => {
|
||||
let free = 1_000;
|
||||
const ctx = await setup({
|
||||
gpu: {
|
||||
minVramMb: 6_000,
|
||||
status: async () => ({ vramFreeMb: free, vramTotalMb: 24_000 }),
|
||||
},
|
||||
});
|
||||
const { jobId } = await ctx.start({ prompt: "fox" });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(ctx.handlerCalls).toHaveLength(0);
|
||||
expect(await ctx.job(jobId)).toMatchObject({ status: "waiting", deferrals: 1 });
|
||||
|
||||
free = 20_000;
|
||||
await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(ctx.handlerCalls).toHaveLength(1);
|
||||
expect((await ctx.job(jobId)).status).toBe("done");
|
||||
});
|
||||
|
||||
it("cancels a queued job at once and a running job through the lease heartbeat", async () => {
|
||||
const ctx = await setup({
|
||||
run: async (context) => {
|
||||
await new Promise<void>((resolve) => {
|
||||
context.signal.addEventListener("abort", () => resolve(), { once: true });
|
||||
});
|
||||
context.signal.throwIfAborted();
|
||||
return { summary: "" };
|
||||
},
|
||||
});
|
||||
const queued = await ctx.start({ prompt: "queued" });
|
||||
await ctx.runtime.cancel(queued.jobId);
|
||||
expect(await ctx.job(queued.jobId)).toMatchObject({ status: "cancelled" });
|
||||
expect((await ctx.card(queued.jobId)).status).toBe("cancelled");
|
||||
expect(ctx.cancelled).toEqual([mediaRunJobKey(queued.jobId)]);
|
||||
await ctx.runtime.runJob(queued.jobId);
|
||||
expect(ctx.handlerCalls).toHaveLength(0);
|
||||
|
||||
// Another process (a second runtime) asks to cancel a job this one runs.
|
||||
const running = await ctx.start({ prompt: "running" });
|
||||
const runningJob = ctx.runtime.runJob(running.jobId);
|
||||
await vi.waitFor(() => expect(ctx.handlerCalls).toEqual([running.jobId]));
|
||||
const otherProcess = createMediaRuntime(ctx.runtimeDeps, ctx.runtimeOptions);
|
||||
await otherProcess.cancel(running.jobId);
|
||||
expect((await ctx.job(running.jobId)).status).toBe("running");
|
||||
await runningJob;
|
||||
expect(await ctx.job(running.jobId)).toMatchObject({ status: "cancelled", leaseToken: null });
|
||||
expect((await ctx.card(running.jobId)).status).toBe("cancelled");
|
||||
});
|
||||
|
||||
it("never runs one job twice and keeps one job on the GPU at a time", async () => {
|
||||
const release = deferred();
|
||||
const entered = deferred();
|
||||
const ctx = await setup({
|
||||
run: async () => {
|
||||
entered.resolve();
|
||||
await release.promise;
|
||||
return { summary: "ok" };
|
||||
},
|
||||
});
|
||||
const otherProcess = createMediaRuntime(ctx.runtimeDeps, ctx.runtimeOptions);
|
||||
const first = await ctx.start({ prompt: "one" });
|
||||
const second = await ctx.start({ prompt: "two" });
|
||||
|
||||
const a = ctx.runtime.runJob(first.jobId);
|
||||
await entered.promise;
|
||||
await otherProcess.runJob(first.jobId);
|
||||
await otherProcess.runJob(second.jobId);
|
||||
expect(ctx.handlerCalls).toEqual([first.jobId]);
|
||||
expect(await ctx.job(second.jobId)).toMatchObject({ status: "waiting", attempts: 0 });
|
||||
|
||||
release.resolve();
|
||||
await a;
|
||||
await prisma.mediaJob.update({
|
||||
where: { id: second.jobId },
|
||||
data: { nextAttemptAt: new Date() },
|
||||
});
|
||||
await otherProcess.runJob(second.jobId);
|
||||
expect(ctx.handlerCalls).toEqual([first.jobId, second.jobId]);
|
||||
expect((await ctx.job(second.jobId)).status).toBe("done");
|
||||
});
|
||||
|
||||
it("recovers a job whose runner died and fails it after too many restarts", async () => {
|
||||
const ctx = await setup();
|
||||
const { jobId } = await ctx.start({ prompt: "fox" });
|
||||
const crashedAt = new Date(Date.now() - 1_000);
|
||||
await prisma.mediaJob.update({
|
||||
where: { id: jobId },
|
||||
data: {
|
||||
status: "running",
|
||||
attempts: 1,
|
||||
leaseToken: "dead-runner",
|
||||
leaseExpiresAt: crashedAt,
|
||||
updatedAt: crashedAt,
|
||||
},
|
||||
});
|
||||
ctx.enqueued.length = 0;
|
||||
await ctx.runtime.reconcile();
|
||||
expect(ctx.enqueued).toEqual([
|
||||
expect.objectContaining({ name: "media.run", payload: { jobId } }),
|
||||
]);
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(await ctx.job(jobId)).toMatchObject({ status: "done", attempts: 2 });
|
||||
|
||||
const doomed = await ctx.start({ prompt: "crashes" });
|
||||
await prisma.mediaJob.update({
|
||||
where: { id: doomed.jobId },
|
||||
data: {
|
||||
status: "running",
|
||||
attempts: 3,
|
||||
leaseToken: "dead-runner",
|
||||
leaseExpiresAt: crashedAt,
|
||||
},
|
||||
});
|
||||
await ctx.runtime.runJob(doomed.jobId);
|
||||
expect(await ctx.job(doomed.jobId)).toMatchObject({
|
||||
status: "failed",
|
||||
error: "The job was interrupted. Try again.",
|
||||
});
|
||||
expect(ctx.handlerCalls).toEqual([jobId]);
|
||||
});
|
||||
|
||||
it("re-enqueues lost and overdue jobs, outside the queue once they look stuck", async () => {
|
||||
const ctx = await setup();
|
||||
const fresh = await ctx.start({ prompt: "fresh" });
|
||||
const lost = await ctx.start({ prompt: "lost" });
|
||||
const stuck = await ctx.start({ prompt: "stuck" });
|
||||
const minuteAgo = new Date(Date.now() - 60_000);
|
||||
const hourAgo = new Date(Date.now() - 3_600_000);
|
||||
await prisma.mediaJob.update({
|
||||
where: { id: lost.jobId },
|
||||
data: { updatedAt: minuteAgo, createdAt: minuteAgo },
|
||||
});
|
||||
await prisma.mediaJob.update({
|
||||
where: { id: stuck.jobId },
|
||||
data: { updatedAt: hourAgo, createdAt: hourAgo },
|
||||
});
|
||||
ctx.enqueued.length = 0;
|
||||
await ctx.runtime.reconcile();
|
||||
const byJob = new Map(
|
||||
ctx.enqueued.map((job) => [(job.payload as { jobId: string }).jobId, job]),
|
||||
);
|
||||
expect(byJob.has(fresh.jobId)).toBe(false);
|
||||
expect(byJob.get(lost.jobId)).toMatchObject({ queue: "media", availableAt: minuteAgo });
|
||||
expect(byJob.get(stuck.jobId)?.queue).toBeUndefined();
|
||||
});
|
||||
|
||||
it("wakes the starting bot with the result when continue_after is set", async () => {
|
||||
const ctx = await setup();
|
||||
const { jobId } = await ctx.start({ prompt: "fox" }, { continueAfter: true });
|
||||
ctx.enqueued.length = 0;
|
||||
await ctx.runtime.runJob(jobId);
|
||||
const wake = await prisma.run.findFirstOrThrow({
|
||||
where: { spaceId: ctx.id, trigger: MEDIA_WAKE_TRIGGER },
|
||||
include: { task: true },
|
||||
});
|
||||
expect(wake).toMatchObject({ botId: ctx.bot.id, threadId: ctx.thread.id, status: "queued" });
|
||||
expect(wake.task.prompt).toContain(`Media job ${jobId}`);
|
||||
expect(wake.task.prompt).toContain("Drew fox.");
|
||||
expect(wake.task.prompt).toContain("picture.png");
|
||||
expect(ctx.enqueued).toEqual([
|
||||
expect.objectContaining({ name: "run.continue", payload: { runId: wake.id } }),
|
||||
]);
|
||||
});
|
||||
|
||||
it("cancels a job whose conversation was cleared without running it", async () => {
|
||||
const ctx = await setup();
|
||||
const { jobId } = await ctx.start({ prompt: "fox" }, { continueAfter: true });
|
||||
await prisma.run.update({ where: { id: ctx.run.id }, data: { status: "completed" } });
|
||||
await clearThread(prisma, { spaceId: ctx.id, botId: ctx.bot.id, threadId: ctx.thread.id });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(ctx.handlerCalls).toHaveLength(0);
|
||||
expect(await ctx.job(jobId)).toMatchObject({ status: "cancelled" });
|
||||
expect(
|
||||
await prisma.run.count({ where: { spaceId: ctx.id, trigger: MEDIA_WAKE_TRIGGER } }),
|
||||
).toBe(0);
|
||||
});
|
||||
|
||||
it("fails a job that outlives its run limit", async () => {
|
||||
const ctx = await setup({
|
||||
run: async (context) => {
|
||||
await new Promise((resolve) => setTimeout(resolve, 5_000));
|
||||
context.signal.throwIfAborted();
|
||||
return { summary: "" };
|
||||
},
|
||||
runtime: { maxRunMs: 50, abortGraceMs: 20 },
|
||||
});
|
||||
const { jobId } = await ctx.start({ prompt: "slow" });
|
||||
await ctx.runtime.runJob(jobId);
|
||||
expect(await ctx.job(jobId)).toMatchObject({
|
||||
status: "failed",
|
||||
error: "The job took too long.",
|
||||
});
|
||||
expect(await readdir(ctx.tempRoot)).toEqual([]);
|
||||
});
|
||||
|
||||
it("reports which operations a space can run", async () => {
|
||||
const ctx = await setup();
|
||||
const capabilities = await ctx.runtime.capabilities(ctx.id);
|
||||
expect(capabilities.generate_image).toBe(true);
|
||||
expect(capabilities.transcribe).toBe(false);
|
||||
const broken = await setup({
|
||||
available: async () => {
|
||||
throw new Error("probe failed");
|
||||
},
|
||||
});
|
||||
expect((await broken.runtime.capabilities(broken.id)).generate_image).toBe(false);
|
||||
});
|
||||
});
|
||||
@ -1,7 +1,11 @@
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { mkdir, readFile, rm, stat, writeFile } from "node:fs/promises";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { constants, createWriteStream } from "node:fs";
|
||||
import { copyFile, mkdir, rm } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import path from "node:path";
|
||||
import { Readable } from "node:stream";
|
||||
import { pipeline } from "node:stream/promises";
|
||||
import type { ReadableStream as NodeReadableStream } from "node:stream/web";
|
||||
import type {
|
||||
AdapterContext,
|
||||
ArtifactStore,
|
||||
@ -17,13 +21,13 @@ import type {
|
||||
StartMediaJobInput,
|
||||
} from "@rakazo/adapter-kit";
|
||||
import {
|
||||
ArtifactTooLargeError,
|
||||
MediaJobDeferredError,
|
||||
mediaRunJob,
|
||||
mediaRunJobKey,
|
||||
runContinueJob,
|
||||
} from "@rakazo/adapter-kit";
|
||||
import {
|
||||
ARTIFACT_MAX_BYTES_BY_KIND,
|
||||
attachmentKindOf,
|
||||
MEDIA_JOB_ACTIVE_STATUSES,
|
||||
type MediaJobBlock,
|
||||
@ -33,11 +37,7 @@ import {
|
||||
type MediaOutput,
|
||||
type MessageBlock,
|
||||
} from "@rakazo/contracts";
|
||||
import {
|
||||
attachmentExtensionForMimeType,
|
||||
inferAttachmentMimeType,
|
||||
validateAttachmentMimeType,
|
||||
} from "@rakazo/core";
|
||||
import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core";
|
||||
import {
|
||||
appendEventInTransaction,
|
||||
createThreadMessageInTransaction,
|
||||
@ -49,6 +49,7 @@ import {
|
||||
} from "@rakazo/db";
|
||||
import { getLogger } from "@rakazo/logging";
|
||||
import type { z } from "zod";
|
||||
import { storeArtifactFromFile } from "./artifact-files.js";
|
||||
|
||||
/** Run trigger of the turn that wakes a bot when its media job ends (`continue_after`). */
|
||||
export const MEDIA_WAKE_TRIGGER = "media";
|
||||
@ -505,6 +506,7 @@ export function createMediaRuntime(
|
||||
|
||||
const storeArtifactFile = async (
|
||||
job: MediaJob,
|
||||
signal: AbortSignal,
|
||||
file: {
|
||||
path: string;
|
||||
name: string;
|
||||
@ -515,38 +517,37 @@ export function createMediaRuntime(
|
||||
parentArtifactId?: string;
|
||||
},
|
||||
) => {
|
||||
// Swap point for the streaming artifact store: until it lands, files go through put().
|
||||
const bytes = new Uint8Array(await readFile(file.path));
|
||||
const context = artifactContext(job, `media-output:${job.id}`);
|
||||
const stored = await deps.artifacts.put(
|
||||
{ name: file.name, mimeType: file.mimeType, bytes },
|
||||
context,
|
||||
);
|
||||
try {
|
||||
return await deps.prisma.artifact.create({
|
||||
data: {
|
||||
spaceId: job.spaceId,
|
||||
userId: job.userId,
|
||||
botId: job.botId,
|
||||
groupId: job.groupId,
|
||||
runId: job.runId,
|
||||
name: file.name,
|
||||
mimeType: file.mimeType,
|
||||
size: bytes.byteLength,
|
||||
hash: createHash("sha256").update(bytes).digest("hex"),
|
||||
storageKey: stored.id,
|
||||
width: file.width,
|
||||
height: file.height,
|
||||
durationMs: file.durationMs,
|
||||
parentArtifactId: file.parentArtifactId,
|
||||
},
|
||||
return await storeArtifactFromFile(deps, {
|
||||
...file,
|
||||
spaceId: job.spaceId,
|
||||
userId: job.userId,
|
||||
botId: job.botId,
|
||||
groupId: job.groupId,
|
||||
runId: job.runId,
|
||||
signal,
|
||||
});
|
||||
} catch (error) {
|
||||
await deps.artifacts.remove(stored.id, context).catch(() => undefined);
|
||||
if (error instanceof ArtifactTooLargeError) throw new MediaJobRejectedError(TOO_LARGE_ERROR);
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
/** Copies a stored artifact to `file`: a clone of the local file when possible, else a stream. */
|
||||
const copyArtifact = async (job: MediaJob, storageKey: string, file: string) => {
|
||||
const context = artifactContext(job, `media-input:${job.id}`);
|
||||
const local = await deps.artifacts.localPath?.(storageKey, context).catch(() => undefined);
|
||||
if (local) {
|
||||
await copyFile(local, file, constants.COPYFILE_EXCL | constants.COPYFILE_FICLONE);
|
||||
return;
|
||||
}
|
||||
const stream = await deps.artifacts.createReadStream(storageKey, context);
|
||||
await pipeline(
|
||||
Readable.fromWeb(stream as NodeReadableStream<Uint8Array>),
|
||||
createWriteStream(file, { flags: "wx", mode: 0o600 }),
|
||||
);
|
||||
};
|
||||
|
||||
const removeArtifacts = async (job: MediaJob, ids: string[]) => {
|
||||
const rows = await deps.prisma.artifact.findMany({
|
||||
where: { id: { in: ids }, spaceId: job.spaceId },
|
||||
@ -657,12 +658,7 @@ export function createMediaRuntime(
|
||||
attachmentExtensionForMimeType(row.mimeType) ||
|
||||
(/^\.[a-z0-9]{1,8}$/i.test(path.extname(row.name)) ? path.extname(row.name) : "");
|
||||
const file = path.join(tempDir, `input-${row.id}${extension.toLowerCase()}`);
|
||||
// Swap point for the streaming artifact store (localPath / createReadStream).
|
||||
const bytes = await deps.artifacts.get(
|
||||
row.storageKey,
|
||||
artifactContext(job, `media-input:${job.id}`),
|
||||
);
|
||||
await writeFile(file, bytes, { mode: 0o600, flag: "wx" });
|
||||
await copyArtifact(job, row.storageKey, file);
|
||||
const input: MediaJobInputFile = {
|
||||
artifactId: row.id,
|
||||
path: file,
|
||||
@ -679,12 +675,8 @@ export function createMediaRuntime(
|
||||
},
|
||||
addOutput: async (file: MediaJobOutputFile) => {
|
||||
signal.throwIfAborted();
|
||||
validateAttachmentMimeType(file.mimeType);
|
||||
const kind = attachmentKindOf(file.mimeType);
|
||||
const filePath = insideDir(tempDir, file.path);
|
||||
const size = (await stat(filePath)).size;
|
||||
if (size > ARTIFACT_MAX_BYTES_BY_KIND[kind])
|
||||
throw new MediaJobRejectedError(TOO_LARGE_ERROR);
|
||||
const parent = file.parentArtifactId
|
||||
? await deps.prisma.artifact.findFirst({
|
||||
where: { id: file.parentArtifactId, spaceId: job.spaceId, userId: job.userId },
|
||||
@ -702,7 +694,7 @@ export function createMediaRuntime(
|
||||
const posterName = `${path.parse(name).name}-poster${path.extname(posterPath).toLowerCase()}`;
|
||||
const posterMime = inferAttachmentMimeType(posterName);
|
||||
if (posterMime && attachmentKindOf(posterMime) === "image") {
|
||||
const poster = await storeArtifactFile(job, {
|
||||
const poster = await storeArtifactFile(job, signal, {
|
||||
path: posterPath,
|
||||
name: posterName,
|
||||
mimeType: posterMime,
|
||||
@ -716,7 +708,7 @@ export function createMediaRuntime(
|
||||
const width = positiveInt(file.width);
|
||||
const height = positiveInt(file.height);
|
||||
const durationMs = nonNegativeInt(file.durationMs);
|
||||
const artifact = await storeArtifactFile(job, {
|
||||
const artifact = await storeArtifactFile(job, signal, {
|
||||
path: filePath,
|
||||
name,
|
||||
mimeType: file.mimeType,
|
||||
|
||||
@ -1,10 +1,6 @@
|
||||
import {
|
||||
type MediaSettings,
|
||||
MediaSettingsSchema,
|
||||
type MediaSettingsUpdate,
|
||||
MediaSettingsUpdateSchema,
|
||||
type MediaSettingsView,
|
||||
} from "@rakazo/contracts";
|
||||
import type { MediaSettings, MediaSettingsUpdate, MediaSettingsView } from "@rakazo/contracts";
|
||||
import { MediaSettingsSchema, MediaSettingsUpdateSchema } from "@rakazo/contracts";
|
||||
import { parseWorkflowTemplate } from "@rakazo/core/media/workflows";
|
||||
import type { PrismaClient } from "@rakazo/db";
|
||||
import {
|
||||
assertAllowedOpenAiCompatibleRequestUrl,
|
||||
@ -82,31 +78,59 @@ function keptKey(
|
||||
return new URL(previous.baseUrl).origin === new URL(baseUrl).origin ? previous.apiKey : undefined;
|
||||
}
|
||||
|
||||
/** A media setting the deployment owner has to correct; the message says what is wrong. */
|
||||
export class MediaSettingsInvalidError extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = "MediaSettingsInvalidError";
|
||||
}
|
||||
}
|
||||
|
||||
function assertValidWorkflow(label: string, workflow: string | undefined): void {
|
||||
if (!workflow) return;
|
||||
try {
|
||||
parseWorkflowTemplate(workflow);
|
||||
} catch (error) {
|
||||
throw new Error(`${label}: ${error instanceof Error ? error.message : String(error)}`);
|
||||
}
|
||||
}
|
||||
|
||||
function resolveSettings(previous: MediaSettings, update: MediaSettingsUpdate): MediaSettings {
|
||||
const parsed = MediaSettingsUpdateSchema.parse(update);
|
||||
try {
|
||||
assertValidWorkflow("Image workflow", parsed.comfyui?.imageWorkflow);
|
||||
assertValidWorkflow("Video workflow", parsed.comfyui?.videoWorkflow);
|
||||
let comfyui: MediaSettings["comfyui"] = null;
|
||||
if (parsed.comfyui) {
|
||||
const baseUrl = normalizeComfyUiBaseUrl(parsed.comfyui.baseUrl);
|
||||
const apiKey = keptKey(parsed.comfyui.apiKey, previous.comfyui, baseUrl);
|
||||
assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey);
|
||||
comfyui = { ...parsed.comfyui, baseUrl, apiKey };
|
||||
}
|
||||
let transcription: MediaSettings["transcription"] = null;
|
||||
if (parsed.transcription) {
|
||||
const baseUrl = normalizeTranscriptionBaseUrl(parsed.transcription.baseUrl);
|
||||
const apiKey = keptKey(parsed.transcription.apiKey, previous.transcription, baseUrl);
|
||||
assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey);
|
||||
transcription = { ...parsed.transcription, baseUrl, apiKey };
|
||||
}
|
||||
return MediaSettingsSchema.parse({ comfyui, transcription });
|
||||
} catch (error) {
|
||||
throw new MediaSettingsInvalidError(
|
||||
error instanceof Error ? error.message : "These media settings are not valid",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates and stores the media backends (deployment owner only; the caller authorizes).
|
||||
* Throws an Error with a user-facing message when a URL is not allowed.
|
||||
* Throws MediaSettingsInvalidError when a URL is not allowed.
|
||||
*/
|
||||
export async function saveMediaSettings(
|
||||
deps: MediaSettingsDeps,
|
||||
update: MediaSettingsUpdate,
|
||||
): Promise<MediaSettingsView> {
|
||||
const parsed = MediaSettingsUpdateSchema.parse(update);
|
||||
const previous = await loadMediaSettings(deps);
|
||||
let comfyui: MediaSettings["comfyui"] = null;
|
||||
if (parsed.comfyui) {
|
||||
const baseUrl = normalizeComfyUiBaseUrl(parsed.comfyui.baseUrl);
|
||||
const apiKey = keptKey(parsed.comfyui.apiKey, previous.comfyui, baseUrl);
|
||||
assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey);
|
||||
comfyui = { ...parsed.comfyui, baseUrl, apiKey };
|
||||
}
|
||||
let transcription: MediaSettings["transcription"] = null;
|
||||
if (parsed.transcription) {
|
||||
const baseUrl = normalizeTranscriptionBaseUrl(parsed.transcription.baseUrl);
|
||||
const apiKey = keptKey(parsed.transcription.apiKey, previous.transcription, baseUrl);
|
||||
assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey);
|
||||
transcription = { ...parsed.transcription, baseUrl, apiKey };
|
||||
}
|
||||
const settings = MediaSettingsSchema.parse({ comfyui, transcription });
|
||||
const settings = resolveSettings(await loadMediaSettings(deps), update);
|
||||
if (!settings.comfyui && !settings.transcription) {
|
||||
await deps.prisma.integrationProviderConfig.deleteMany({ where: { id: MEDIA_SETTINGS_ID } });
|
||||
return mediaSettingsView(settings);
|
||||
|
||||
125
packages/adapters/src/media-tools.test.ts
Normal file
125
packages/adapters/src/media-tools.test.ts
Normal file
@ -0,0 +1,125 @@
|
||||
import type { MediaCapabilities, MediaRuntime, StartMediaJobInput } from "@rakazo/adapter-kit";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { selectBuiltinToolsForRun } from "./executor-tools.js";
|
||||
import { MediaJobRejectedError } from "./media-runtime.js";
|
||||
import {
|
||||
executeMediaTool,
|
||||
isMediaToolName,
|
||||
MEDIA_CONTINUE_AFTER_PROPERTY,
|
||||
mediaToolInputSchema,
|
||||
selectMediaTools,
|
||||
} from "./media-tools.js";
|
||||
|
||||
const none: MediaCapabilities = {
|
||||
generate_image: false,
|
||||
generate_video: false,
|
||||
edit_image: false,
|
||||
edit_video: false,
|
||||
export_video: false,
|
||||
transcribe: false,
|
||||
};
|
||||
|
||||
const scope = {
|
||||
spaceId: "space-1",
|
||||
userId: "user-1",
|
||||
botId: "bot-1",
|
||||
groupId: null,
|
||||
threadId: "thread-1",
|
||||
runId: "run-1",
|
||||
operationKey: "effect-1",
|
||||
};
|
||||
|
||||
function runtime(start: MediaRuntime["start"]): MediaRuntime {
|
||||
return {
|
||||
capabilities: async () => none,
|
||||
start,
|
||||
get: async () => null,
|
||||
cancel: async () => undefined,
|
||||
};
|
||||
}
|
||||
|
||||
describe("media tools", () => {
|
||||
it("names tools after their operations", () => {
|
||||
expect(isMediaToolName("generate_image")).toBe(true);
|
||||
expect(isMediaToolName("transcribe")).toBe(true);
|
||||
expect(isMediaToolName("render_plot")).toBe(false);
|
||||
});
|
||||
|
||||
it("shows a media tool only when the space can run its operation", () => {
|
||||
const tools = [{ name: "generate_image" }, { name: "transcribe" }, { name: "web_search" }];
|
||||
expect(selectMediaTools(tools, undefined).map((tool) => tool.name)).toEqual(["web_search"]);
|
||||
expect(selectMediaTools(tools, { ...none, transcribe: true }).map((tool) => tool.name)).toEqual(
|
||||
["transcribe", "web_search"],
|
||||
);
|
||||
});
|
||||
|
||||
it("leaves non-media tools of a run untouched", () => {
|
||||
const base = {
|
||||
graphicalToolsAllowed: true,
|
||||
groupId: null,
|
||||
trigger: "user",
|
||||
semanticMemoryEnabled: false,
|
||||
messagingChannelRun: false,
|
||||
bot: { kind: "specialist", restrictBuiltinTools: false, allowedBuiltinTools: [] },
|
||||
};
|
||||
const names = (mediaCapabilities?: MediaCapabilities) =>
|
||||
selectBuiltinToolsForRun({ ...base, mediaCapabilities })
|
||||
.map((tool) => tool.name)
|
||||
.filter((name) => !isMediaToolName(name));
|
||||
expect(names({ ...none, generate_image: true, transcribe: true })).toEqual(names(undefined));
|
||||
expect(selectBuiltinToolsForRun(base).some((tool) => isMediaToolName(tool.name))).toBe(false);
|
||||
});
|
||||
|
||||
it("adds continue_after to every media tool schema", () => {
|
||||
expect(
|
||||
mediaToolInputSchema({
|
||||
properties: { prompt: { type: "string" } },
|
||||
required: ["prompt"],
|
||||
}),
|
||||
).toEqual({
|
||||
type: "object",
|
||||
properties: { prompt: { type: "string" }, ...MEDIA_CONTINUE_AFTER_PROPERTY },
|
||||
required: ["prompt"],
|
||||
});
|
||||
});
|
||||
|
||||
it("starts the job with the tool scope and returns at once", async () => {
|
||||
const start = vi.fn(async (_input: StartMediaJobInput) => ({
|
||||
jobId: "job-1",
|
||||
messageId: "message-1",
|
||||
}));
|
||||
const result = await executeMediaTool(runtime(start), scope, "generate_image", {
|
||||
prompt: "a fox",
|
||||
continue_after: true,
|
||||
});
|
||||
expect(start).toHaveBeenCalledWith({
|
||||
operation: "generate_image",
|
||||
input: { prompt: "a fox" },
|
||||
...scope,
|
||||
continueAfter: true,
|
||||
});
|
||||
expect(result).toMatchObject({ jobId: "job-1", status: "queued" });
|
||||
expect((result as { note: string }).note).toContain("woken");
|
||||
|
||||
await executeMediaTool(runtime(start), scope, "generate_image", { prompt: "a fox" });
|
||||
expect(start.mock.calls[1]?.[0]).toMatchObject({ continueAfter: false });
|
||||
});
|
||||
|
||||
it("returns refusals to the model and rethrows other failures", async () => {
|
||||
const refused = runtime(async () => {
|
||||
throw new MediaJobRejectedError("Invalid generate_image input: prompt: Required");
|
||||
});
|
||||
expect(await executeMediaTool(refused, scope, "generate_image", {})).toEqual({
|
||||
error: "Invalid generate_image input: prompt: Required",
|
||||
});
|
||||
expect(await executeMediaTool(undefined, scope, "transcribe", {})).toEqual({
|
||||
error: "transcribe is not available.",
|
||||
});
|
||||
const broken = runtime(async () => {
|
||||
throw new Error("database down");
|
||||
});
|
||||
await expect(executeMediaTool(broken, scope, "transcribe", {})).rejects.toThrow(
|
||||
"database down",
|
||||
);
|
||||
});
|
||||
});
|
||||
28
packages/adapters/src/media/comfyui-settings.test.ts
Normal file
28
packages/adapters/src/media/comfyui-settings.test.ts
Normal file
@ -0,0 +1,28 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { saveMediaSettings } from "../media-settings.js";
|
||||
|
||||
const comfyui = (workflows: { imageWorkflow?: string; videoWorkflow?: string }) => ({
|
||||
comfyui: { baseUrl: "http://127.0.0.1:8188", ...workflows },
|
||||
transcription: null,
|
||||
});
|
||||
|
||||
describe("ComfyUI settings", () => {
|
||||
it("refuses custom workflows the provider could not fill", async () => {
|
||||
const deps = {} as Parameters<typeof saveMediaSettings>[0];
|
||||
await expect(
|
||||
saveMediaSettings(
|
||||
deps,
|
||||
comfyui({ imageWorkflow: '{"1": {"class_type": "X", "inputs": {}}}' }),
|
||||
),
|
||||
).rejects.toThrow("Image workflow: Workflow needs a {{prompt}} placeholder");
|
||||
await expect(
|
||||
saveMediaSettings(
|
||||
deps,
|
||||
comfyui({
|
||||
videoWorkflow:
|
||||
'{"1": {"class_type": "X", "inputs": {"t": "{{prompt}}", "c": "{{cfg}}"}}}',
|
||||
}),
|
||||
),
|
||||
).rejects.toThrow("Video workflow: Unknown workflow placeholders: {{cfg}}");
|
||||
});
|
||||
});
|
||||
110
packages/adapters/src/media/image-edit.test.ts
Normal file
110
packages/adapters/src/media/image-edit.test.ts
Normal file
@ -0,0 +1,110 @@
|
||||
import sharp from "sharp";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { editImageGeometry, encodeImage, imageSize } from "./image-edit.js";
|
||||
|
||||
/** Left half red, right half blue. */
|
||||
async function halves(width = 40, height = 20): Promise<Buffer> {
|
||||
const half = Math.floor(width / 2);
|
||||
const blue = await sharp({
|
||||
create: { width: width - half, height, channels: 3, background: { r: 0, g: 0, b: 255 } },
|
||||
})
|
||||
.png()
|
||||
.toBuffer();
|
||||
return sharp({ create: { width, height, channels: 3, background: { r: 255, g: 0, b: 0 } } })
|
||||
.composite([{ input: blue, left: half, top: 0 }])
|
||||
.png()
|
||||
.toBuffer();
|
||||
}
|
||||
|
||||
async function pixel(image: Uint8Array, x: number, y: number) {
|
||||
const { data, info } = await sharp(image).raw().toBuffer({ resolveWithObject: true });
|
||||
const offset = (y * info.width + x) * info.channels;
|
||||
return [...data.subarray(offset, offset + info.channels)];
|
||||
}
|
||||
|
||||
async function size(image: Uint8Array) {
|
||||
const { width, height } = await sharp(image).metadata();
|
||||
return [width, height];
|
||||
}
|
||||
|
||||
describe("editImageGeometry", () => {
|
||||
it("applies crop before rotate, in source coordinates", async () => {
|
||||
const result = await editImageGeometry(await halves(), {
|
||||
crop: { left: 20, top: 0, width: 20, height: 10 },
|
||||
rotate: 90,
|
||||
});
|
||||
expect(await size(result)).toEqual([10, 20]);
|
||||
expect((await pixel(result, 5, 5)).slice(0, 3)).toEqual([0, 0, 255]);
|
||||
});
|
||||
|
||||
it("flips horizontally and vertically", async () => {
|
||||
const flopped = await editImageGeometry(await halves(), { flip: "horizontal" });
|
||||
expect((await pixel(flopped, 0, 0)).slice(0, 3)).toEqual([0, 0, 255]);
|
||||
const flipped = await editImageGeometry(await halves(), { flip: "vertical" });
|
||||
expect((await pixel(flipped, 0, 0)).slice(0, 3)).toEqual([255, 0, 0]);
|
||||
});
|
||||
|
||||
it("resizes to a preset with crop, pad or blur, and keeps the aspect for one side", async () => {
|
||||
const source = await halves(160, 90);
|
||||
for (const fit of ["crop", "pad", "blur"] as const) {
|
||||
const result = await editImageGeometry(source, { resize: { aspect: "square", fit } });
|
||||
expect(await size(result)).toEqual([1080, 1080]);
|
||||
}
|
||||
const padded = await editImageGeometry(source, { resize: { aspect: "vertical", fit: "pad" } });
|
||||
expect((await pixel(padded, 0, 0)).at(-1)).toBe(0);
|
||||
const scaled = await editImageGeometry(source, { resize: { width: 80, fit: "crop" } });
|
||||
expect(await size(scaled)).toEqual([80, 45]);
|
||||
});
|
||||
|
||||
it("extends with a color or transparency", async () => {
|
||||
const colored = await editImageGeometry(await halves(), {
|
||||
extend: { top: 5, right: 0, bottom: 0, left: 3, color: "#00ff00" },
|
||||
});
|
||||
expect(await size(colored)).toEqual([43, 25]);
|
||||
expect((await pixel(colored, 0, 0)).slice(0, 3)).toEqual([0, 255, 0]);
|
||||
const clear = await editImageGeometry(await halves(), {
|
||||
extend: { top: 2, right: 2, bottom: 2, left: 2 },
|
||||
});
|
||||
expect((await pixel(clear, 0, 0)).at(-1)).toBe(0);
|
||||
});
|
||||
|
||||
it("applies EXIF orientation first", async () => {
|
||||
const rotatedJpeg = await sharp(await halves())
|
||||
.jpeg()
|
||||
.withMetadata({ orientation: 6 })
|
||||
.toBuffer();
|
||||
expect(await imageSize(rotatedJpeg)).toEqual({ width: 20, height: 40 });
|
||||
expect(await size(await editImageGeometry(rotatedJpeg, { flip: "vertical" }))).toEqual([
|
||||
20, 40,
|
||||
]);
|
||||
});
|
||||
|
||||
it("refuses crops outside the image and oversized results", async () => {
|
||||
await expect(
|
||||
editImageGeometry(await halves(), { crop: { left: 30, top: 0, width: 20, height: 10 } }),
|
||||
).rejects.toThrow("The crop area is outside the 40×20 image");
|
||||
await expect(
|
||||
editImageGeometry(await halves(), {
|
||||
extend: { top: 9000, right: 0, bottom: 0, left: 9000 },
|
||||
}),
|
||||
).rejects.toThrow("The edited image would be too large");
|
||||
await expect(
|
||||
editImageGeometry(await halves(), { resize: { width: 8192, height: 8193, fit: "crop" } }),
|
||||
).rejects.toThrow("The edited image would be too large");
|
||||
});
|
||||
});
|
||||
|
||||
describe("encodeImage", () => {
|
||||
it("encodes each format and flattens transparency for JPEG", async () => {
|
||||
const clear = await editImageGeometry(await halves(), {
|
||||
extend: { top: 2, right: 0, bottom: 0, left: 0 },
|
||||
});
|
||||
for (const format of ["png", "jpeg", "webp"] as const) {
|
||||
const encoded = await encodeImage(clear, format);
|
||||
expect((await sharp(encoded.bytes).metadata()).format).toBe(format);
|
||||
expect([encoded.width, encoded.height]).toEqual([40, 22]);
|
||||
}
|
||||
const jpeg = await encodeImage(clear, "jpeg");
|
||||
expect((await pixel(jpeg.bytes, 0, 0)).every((channel) => channel < 16)).toBe(true);
|
||||
});
|
||||
});
|
||||
156
packages/adapters/src/media/image-edit.ts
Normal file
156
packages/adapters/src/media/image-edit.ts
Normal file
@ -0,0 +1,156 @@
|
||||
import type { EditImageInput, ImageFormat } from "@rakazo/contracts";
|
||||
import { MEDIA_PRESETS } from "@rakazo/contracts";
|
||||
import type sharpFactory from "sharp";
|
||||
import type { Sharp } from "sharp";
|
||||
|
||||
type SharpFactory = typeof sharpFactory;
|
||||
export type ImageGeometryEdits = Pick<
|
||||
EditImageInput,
|
||||
"crop" | "rotate" | "flip" | "resize" | "extend"
|
||||
>;
|
||||
|
||||
/** Edited images stay below this many pixels (8192²). */
|
||||
const MAX_EDIT_PIXELS = 8192 * 8192;
|
||||
const TRANSPARENT = { r: 0, g: 0, b: 0, alpha: 0 };
|
||||
const BLACK = { r: 0, g: 0, b: 0 };
|
||||
|
||||
let sharpModule: Promise<SharpFactory> | undefined;
|
||||
|
||||
/** sharp is a native module; load it on first use so a platform without it only loses edits. */
|
||||
export function loadSharp(): Promise<SharpFactory> {
|
||||
sharpModule ??= import("sharp").then((module) => module.default);
|
||||
return sharpModule;
|
||||
}
|
||||
|
||||
export async function sharpAvailable(): Promise<boolean> {
|
||||
return loadSharp().then(
|
||||
() => true,
|
||||
() => false,
|
||||
);
|
||||
}
|
||||
|
||||
/** Displayed size, with EXIF orientation applied. */
|
||||
export async function imageSize(
|
||||
input: string | Uint8Array,
|
||||
): Promise<{ width: number; height: number }> {
|
||||
const sharp = await loadSharp();
|
||||
const { autoOrient } = await sharp(input).metadata();
|
||||
return { width: autoOrient.width, height: autoOrient.height };
|
||||
}
|
||||
|
||||
/** Covers exactly width×height (center crop) as PNG: generation inputs at the size asked for. */
|
||||
export async function coverImage(
|
||||
input: string | Uint8Array,
|
||||
width: number,
|
||||
height: number,
|
||||
): Promise<Uint8Array> {
|
||||
const sharp = await loadSharp();
|
||||
const png = await sharp(input, { autoOrient: true })
|
||||
.resize(width, height, { fit: "cover" })
|
||||
.png()
|
||||
.toBuffer();
|
||||
return new Uint8Array(png);
|
||||
}
|
||||
|
||||
function parseColor(hex: string | undefined) {
|
||||
if (!hex) return TRANSPARENT;
|
||||
const value = Number.parseInt(hex.slice(1), 16);
|
||||
return { r: (value >> 16) & 255, g: (value >> 8) & 255, b: value & 255, alpha: 1 };
|
||||
}
|
||||
|
||||
function assertSize(width: number, height: number): void {
|
||||
if (width * height > MAX_EDIT_PIXELS) throw new Error("The edited image would be too large");
|
||||
}
|
||||
|
||||
/** Runs each step on decoded pixels, so steps apply in the order given, not sharp's own order. */
|
||||
async function materialize(sharp: SharpFactory, image: Sharp) {
|
||||
const { data, info } = await image.raw().toBuffer({ resolveWithObject: true });
|
||||
const { width, height, channels } = info;
|
||||
return { image: sharp(data, { raw: { width, height, channels } }), width, height };
|
||||
}
|
||||
|
||||
async function resizeTo(
|
||||
sharp: SharpFactory,
|
||||
image: Sharp,
|
||||
width: number,
|
||||
height: number,
|
||||
fit: "crop" | "pad" | "blur",
|
||||
): Promise<Sharp> {
|
||||
if (fit === "crop") return image.resize(width, height, { fit: "cover" });
|
||||
if (fit === "pad")
|
||||
return image.resize(width, height, { fit: "contain", background: TRANSPARENT });
|
||||
const source = await image.png().toBuffer();
|
||||
const background = await sharp(source)
|
||||
.resize(width, height, { fit: "cover" })
|
||||
.blur(Math.max(width, height) / 40)
|
||||
.toBuffer();
|
||||
const foreground = await sharp(source).resize(width, height, { fit: "inside" }).toBuffer();
|
||||
return sharp(background).composite([{ input: foreground, gravity: "centre" }]);
|
||||
}
|
||||
|
||||
/**
|
||||
* Applies crop, rotate, flip, resize and extend in that order and returns a lossless PNG.
|
||||
* EXIF orientation is applied first, so coordinates match what viewers show.
|
||||
*/
|
||||
export async function editImageGeometry(
|
||||
input: string | Uint8Array,
|
||||
edits: ImageGeometryEdits,
|
||||
): Promise<Buffer> {
|
||||
const sharp = await loadSharp();
|
||||
let current = await materialize(sharp, sharp(input, { autoOrient: true }));
|
||||
const step = async (apply: (image: Sharp) => Sharp | Promise<Sharp>) => {
|
||||
current = await materialize(sharp, await apply(current.image));
|
||||
};
|
||||
const { crop, rotate, flip, resize, extend } = edits;
|
||||
if (crop) {
|
||||
if (crop.left + crop.width > current.width || crop.top + crop.height > current.height) {
|
||||
throw new Error(`The crop area is outside the ${current.width}×${current.height} image`);
|
||||
}
|
||||
await step((image) => image.extract(crop));
|
||||
}
|
||||
if (rotate) await step((image) => image.rotate(rotate));
|
||||
if (flip) await step((image) => (flip === "vertical" ? image.flip() : image.flop()));
|
||||
if (resize) {
|
||||
const preset = resize.aspect ? MEDIA_PRESETS[resize.aspect] : undefined;
|
||||
const width = preset?.width ?? resize.width;
|
||||
const height = preset?.height ?? resize.height;
|
||||
const scale = width ? width / current.width : height! / current.height;
|
||||
assertSize(width ?? current.width * scale, height ?? current.height * scale);
|
||||
await step((image) =>
|
||||
width && height
|
||||
? resizeTo(sharp, image, width, height, resize.fit)
|
||||
: image.resize(width ?? null, height ?? null),
|
||||
);
|
||||
}
|
||||
if (extend) {
|
||||
const { top, right, bottom, left, color } = extend;
|
||||
assertSize(current.width + left + right, current.height + top + bottom);
|
||||
await step((image) =>
|
||||
image.extend({ top, right, bottom, left, background: parseColor(color) }),
|
||||
);
|
||||
}
|
||||
return current.image.png().toBuffer();
|
||||
}
|
||||
|
||||
export const IMAGE_FORMAT_MIME_TYPES: Record<ImageFormat, string> = {
|
||||
png: "image/png",
|
||||
jpeg: "image/jpeg",
|
||||
webp: "image/webp",
|
||||
};
|
||||
|
||||
/** Encodes the final image; JPEG has no transparency, so transparent areas turn black. */
|
||||
export async function encodeImage(
|
||||
input: Uint8Array,
|
||||
format: ImageFormat,
|
||||
): Promise<{ bytes: Buffer; width: number; height: number }> {
|
||||
const sharp = await loadSharp();
|
||||
const image = sharp(input);
|
||||
const encoded =
|
||||
format === "jpeg"
|
||||
? image.flatten({ background: BLACK }).jpeg({ quality: 90, mozjpeg: true })
|
||||
: format === "webp"
|
||||
? image.webp({ quality: 90 })
|
||||
: image.png();
|
||||
const { data, info } = await encoded.toBuffer({ resolveWithObject: true });
|
||||
return { bytes: data, width: info.width, height: info.height };
|
||||
}
|
||||
172
packages/adapters/src/media/image-operations.postgres.test.ts
Normal file
172
packages/adapters/src/media/image-operations.postgres.test.ts
Normal file
@ -0,0 +1,172 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { mkdtemp, rm } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import type { JobPublisher } from "@rakazo/adapter-kit";
|
||||
import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts";
|
||||
import { ComfyUiSettingsSchema } from "@rakazo/contracts";
|
||||
import type { PrismaClient } from "@rakazo/db";
|
||||
import { createDb } from "@rakazo/db";
|
||||
import { afterAll, beforeAll, describe, expect, it } from "vitest";
|
||||
import { LocalArtifactStore } from "../artifacts.js";
|
||||
import { createMediaRuntime, MediaJobRejectedError } from "../media-runtime.js";
|
||||
import { ComfyUiEmulator } from "../testing/comfyui-emulator.js";
|
||||
import { createImageOperations } from "./image-operations.js";
|
||||
|
||||
const describePostgres =
|
||||
process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip;
|
||||
|
||||
describePostgres("ComfyUI operations in the media runtime (PostgreSQL + ComfyUI emulator)", () => {
|
||||
let db: ReturnType<typeof createDb>;
|
||||
let prisma: PrismaClient;
|
||||
let root: string;
|
||||
const users: string[] = [];
|
||||
|
||||
beforeAll(async () => {
|
||||
db = createDb(process.env.DATABASE_URL!);
|
||||
prisma = db.prisma;
|
||||
root = await mkdtemp(join(tmpdir(), "image-operations-pg-"));
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
if (!db) return;
|
||||
await prisma.mediaJob.deleteMany({ where: { userId: { in: users } } });
|
||||
await prisma.artifact.deleteMany({ where: { userId: { in: users } } });
|
||||
await prisma.organization.deleteMany({ where: { id: { in: users } } });
|
||||
await prisma.user.deleteMany({ where: { id: { in: users } } });
|
||||
await prisma.$disconnect();
|
||||
await db.pool.end();
|
||||
await rm(root, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
async function setup() {
|
||||
const id = randomUUID();
|
||||
users.push(id);
|
||||
await prisma.user.create({
|
||||
data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true },
|
||||
});
|
||||
await prisma.organization.create({
|
||||
data: { id, name: "Media test", slug: id, createdAt: new Date() },
|
||||
});
|
||||
await prisma.space.create({ data: { id, organizationId: id, name: "Media", isDefault: true } });
|
||||
await prisma.member.create({
|
||||
data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() },
|
||||
});
|
||||
const bot = await prisma.bot.create({
|
||||
data: { spaceId: id, userId: id, name: "Studio", color: "test-color" },
|
||||
});
|
||||
const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } });
|
||||
const emulator = new ComfyUiEmulator();
|
||||
emulator.pendingPolls = 0;
|
||||
const jobs: JobPublisher = {
|
||||
enqueue: async () => undefined,
|
||||
cancel: async () => undefined,
|
||||
close: async () => undefined,
|
||||
};
|
||||
const runtime = createMediaRuntime(
|
||||
{
|
||||
prisma,
|
||||
events: { notify: async () => undefined },
|
||||
jobs,
|
||||
artifacts: new LocalArtifactStore(join(root, id, "artifacts")),
|
||||
registrations: createImageOperations({
|
||||
loadSettings: async () => ({
|
||||
comfyui: ComfyUiSettingsSchema.parse({ baseUrl: "http://127.0.0.1:8188" }),
|
||||
transcription: null,
|
||||
}),
|
||||
fetch: emulator.fetch,
|
||||
}),
|
||||
},
|
||||
{ tempRoot: join(root, id, "jobs") },
|
||||
);
|
||||
const scope = {
|
||||
spaceId: id,
|
||||
userId: id,
|
||||
botId: bot.id,
|
||||
groupId: null,
|
||||
threadId: thread.id,
|
||||
runId: null,
|
||||
};
|
||||
return { emulator, runtime, scope };
|
||||
}
|
||||
|
||||
async function card(messageId: string, jobId: string): Promise<MediaJobBlock | undefined> {
|
||||
const message = await prisma.message.findUniqueOrThrow({ where: { id: messageId } });
|
||||
return (message.blocks as MessageBlock[]).find(
|
||||
(block): block is MediaJobBlock => block.kind === "media" && block.jobId === jobId,
|
||||
);
|
||||
}
|
||||
|
||||
it("generates images onto the card and edits one into a linked version", async () => {
|
||||
const { runtime, scope } = await setup();
|
||||
const generated = await runtime.start({
|
||||
...scope,
|
||||
operation: "generate_image",
|
||||
input: { prompt: "a lighthouse at dusk", aspect: "landscape", count: 2, seed: 3 },
|
||||
});
|
||||
await runtime.runJob(generated.jobId);
|
||||
const job = await runtime.get(generated.jobId);
|
||||
expect(job).toMatchObject({ status: "done", title: "a lighthouse at dusk" });
|
||||
expect(job?.outputs.map((output) => [output.mimeType, output.width, output.height])).toEqual([
|
||||
["image/png", 1344, 768],
|
||||
["image/png", 1344, 768],
|
||||
]);
|
||||
expect(await card(generated.messageId, generated.jobId)).toMatchObject({
|
||||
status: "done",
|
||||
outputs: job?.outputs,
|
||||
});
|
||||
|
||||
const source = job!.outputs[0]!.artifactId;
|
||||
const edited = await runtime.start({
|
||||
...scope,
|
||||
operation: "edit_image",
|
||||
input: { artifactId: source, resize: { aspect: "vertical", fit: "blur" }, format: "jpeg" },
|
||||
});
|
||||
await runtime.runJob(edited.jobId);
|
||||
const edit = await runtime.get(edited.jobId);
|
||||
expect(edit).toMatchObject({ status: "done", title: "9:16 · JPEG" });
|
||||
const [output] = edit!.outputs;
|
||||
expect(output).toMatchObject({ mimeType: "image/jpeg", width: 1080, height: 1920 });
|
||||
const artifact = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } });
|
||||
expect(artifact).toMatchObject({ parentArtifactId: source, width: 1080, height: 1920 });
|
||||
});
|
||||
|
||||
it("refuses source artifacts the user cannot read", async () => {
|
||||
const { runtime, scope } = await setup();
|
||||
const other = await setup();
|
||||
const generated = await other.runtime.start({
|
||||
...other.scope,
|
||||
operation: "generate_image",
|
||||
input: { prompt: "private" },
|
||||
});
|
||||
await other.runtime.runJob(generated.jobId);
|
||||
const foreign = (await other.runtime.get(generated.jobId))!.outputs[0]!.artifactId;
|
||||
await expect(
|
||||
runtime.start({
|
||||
...scope,
|
||||
operation: "generate_video",
|
||||
input: { prompt: "animate it", sourceArtifactId: foreign },
|
||||
}),
|
||||
).rejects.toThrow(MediaJobRejectedError);
|
||||
await expect(
|
||||
runtime.start({
|
||||
...scope,
|
||||
operation: "edit_image",
|
||||
input: { artifactId: foreign, rotate: 90 },
|
||||
}),
|
||||
).rejects.toThrow(MediaJobRejectedError);
|
||||
});
|
||||
|
||||
it("waits while the GPU lacks free memory", async () => {
|
||||
const { emulator, runtime, scope } = await setup();
|
||||
emulator.gpu = { totalMb: 24_576, freeMb: 4_000, torchReservedMb: 0, torchFreeMb: 0 };
|
||||
const video = await runtime.start({
|
||||
...scope,
|
||||
operation: "generate_video",
|
||||
input: { prompt: "rain on a window" },
|
||||
});
|
||||
await runtime.runJob(video.jobId);
|
||||
expect(await runtime.get(video.jobId)).toMatchObject({ status: "waiting", outputs: [] });
|
||||
expect(emulator.prompts).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
274
packages/adapters/src/media/image-operations.test.ts
Normal file
274
packages/adapters/src/media/image-operations.test.ts
Normal file
@ -0,0 +1,274 @@
|
||||
import { copyFile, mkdtemp, rm, stat, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import type { MediaJobContext, MediaJobOutputFile } from "@rakazo/adapter-kit";
|
||||
import { MediaJobDeferredError } from "@rakazo/adapter-kit";
|
||||
import type { ComfyUiSettings, MediaOutput } from "@rakazo/contracts";
|
||||
import { attachmentKindOf, ComfyUiSettingsSchema } from "@rakazo/contracts";
|
||||
import sharp from "sharp";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { ComfyUiEmulator } from "../testing/comfyui-emulator.js";
|
||||
import {
|
||||
createImageOperations,
|
||||
editImageTitle,
|
||||
fileStem,
|
||||
generationSize,
|
||||
IMAGE_GENERATION_MIN_VRAM_MB,
|
||||
VIDEO_GENERATION_MIN_VRAM_MB,
|
||||
} from "./image-operations.js";
|
||||
|
||||
const dirs: string[] = [];
|
||||
afterEach(async () => {
|
||||
await Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
async function tempDir(): Promise<string> {
|
||||
const dir = await mkdtemp(join(tmpdir(), "image-operations-"));
|
||||
dirs.push(dir);
|
||||
return dir;
|
||||
}
|
||||
|
||||
async function fixture(name: string, width: number, height: number): Promise<string> {
|
||||
const path = join(await tempDir(), name);
|
||||
const image = sharp({
|
||||
create: { width, height, channels: 3, background: { r: 200, g: 120, b: 40 } },
|
||||
});
|
||||
await writeFile(path, await (name.endsWith(".jpg") ? image.jpeg() : image.png()).toBuffer());
|
||||
return path;
|
||||
}
|
||||
|
||||
function setup(
|
||||
settings: ComfyUiSettings | null = ComfyUiSettingsSchema.parse({
|
||||
baseUrl: "http://127.0.0.1:8188",
|
||||
}),
|
||||
) {
|
||||
const emulator = new ComfyUiEmulator();
|
||||
emulator.pendingPolls = 0;
|
||||
const operations = createImageOperations({
|
||||
loadSettings: async () => ({ comfyui: settings, transcription: null }),
|
||||
fetch: emulator.fetch,
|
||||
});
|
||||
const operation = (name: string) => operations.find((entry) => entry.operation === name)!;
|
||||
return { emulator, operations, operation };
|
||||
}
|
||||
|
||||
async function run(
|
||||
operation: ReturnType<typeof setup>["operation"],
|
||||
name: string,
|
||||
input: unknown,
|
||||
artifacts: Record<string, string> = {},
|
||||
) {
|
||||
const dir = await tempDir();
|
||||
const outputs: MediaJobOutputFile[] = [];
|
||||
const progress: number[] = [];
|
||||
const context: MediaJobContext = {
|
||||
jobId: "job-1",
|
||||
spaceId: "space-1",
|
||||
userId: "user-1",
|
||||
botId: null,
|
||||
input,
|
||||
tempDir: dir,
|
||||
signal: new AbortController().signal,
|
||||
progress: async (fraction) => {
|
||||
progress.push(fraction);
|
||||
},
|
||||
fetchArtifact: async (artifactId) => {
|
||||
const source = artifacts[artifactId];
|
||||
if (!source) throw new Error("Artifact not found");
|
||||
const name = source.split("/").at(-1)!;
|
||||
const path = join(dir, `input-${name}`);
|
||||
await copyFile(source, path);
|
||||
const mimeType = name.endsWith(".jpg")
|
||||
? "image/jpeg"
|
||||
: name.endsWith(".mp4")
|
||||
? "video/mp4"
|
||||
: "image/png";
|
||||
return {
|
||||
artifactId,
|
||||
path,
|
||||
name,
|
||||
mimeType,
|
||||
kind: attachmentKindOf(mimeType),
|
||||
size: (await stat(path)).size,
|
||||
};
|
||||
},
|
||||
addOutput: async (file) => {
|
||||
outputs.push(file);
|
||||
const output: MediaOutput = {
|
||||
artifactId: `output-${outputs.length}`,
|
||||
kind: attachmentKindOf(file.mimeType),
|
||||
mimeType: file.mimeType,
|
||||
name: file.name,
|
||||
size: (await stat(file.path)).size,
|
||||
};
|
||||
return output;
|
||||
},
|
||||
};
|
||||
const result = await operation(name).handler.run(context);
|
||||
return { result, outputs, progress };
|
||||
}
|
||||
|
||||
describe("image operations", () => {
|
||||
it("are available by configuration and guard the GPU", async () => {
|
||||
const configured = setup();
|
||||
expect(configured.operations.map((entry) => entry.operation)).toEqual([
|
||||
"generate_image",
|
||||
"generate_video",
|
||||
"edit_image",
|
||||
]);
|
||||
for (const entry of configured.operations) {
|
||||
await expect(entry.available({ spaceId: "space-1" })).resolves.toBe(true);
|
||||
}
|
||||
const generateImage = configured.operation("generate_image");
|
||||
expect(generateImage.gpu?.minVramMb).toBe(IMAGE_GENERATION_MIN_VRAM_MB);
|
||||
expect(configured.operation("generate_video").gpu?.minVramMb).toBe(
|
||||
VIDEO_GENERATION_MIN_VRAM_MB,
|
||||
);
|
||||
await expect(generateImage.gpu?.status()).resolves.toEqual({
|
||||
vramFreeMb: 20_000,
|
||||
vramTotalMb: 24_576,
|
||||
});
|
||||
const unconfigured = setup(null);
|
||||
await expect(
|
||||
unconfigured.operation("generate_image").available({ spaceId: "s" }),
|
||||
).resolves.toBe(false);
|
||||
await expect(unconfigured.operation("edit_image").available({ spaceId: "s" })).resolves.toBe(
|
||||
true,
|
||||
);
|
||||
await expect(run(unconfigured.operation, "generate_image", { prompt: "x" })).rejects.toThrow(
|
||||
"ComfyUI is not configured",
|
||||
);
|
||||
});
|
||||
|
||||
it("generates images at the aspect asked for", async () => {
|
||||
const { operation } = setup();
|
||||
const { result, outputs, progress } = await run(operation, "generate_image", {
|
||||
prompt: "A red kite, über the sea!",
|
||||
aspect: "vertical",
|
||||
count: 2,
|
||||
seed: 11,
|
||||
continue_after: true,
|
||||
});
|
||||
expect(
|
||||
outputs.map(({ name, mimeType, width, height }) => [name, mimeType, width, height]),
|
||||
).toEqual([
|
||||
["a-red-kite-uber-the-sea-1.png", "image/png", 768, 1344],
|
||||
["a-red-kite-uber-the-sea-2.png", "image/png", 768, 1344],
|
||||
]);
|
||||
expect(progress.at(-1)).toBe(1);
|
||||
expect(result.summary).toBe("Generated 2 images with seed 11.");
|
||||
});
|
||||
|
||||
it("starts from a source image at the source's aspect and links the result to it", async () => {
|
||||
const { operation, emulator } = setup();
|
||||
const photo = await fixture("photo.jpg", 300, 200);
|
||||
const { outputs } = await run(
|
||||
operation,
|
||||
"generate_image",
|
||||
{ prompt: "night", sourceArtifactId: "photo-1" },
|
||||
{ "photo-1": photo },
|
||||
);
|
||||
const [upload] = emulator.uploads.values();
|
||||
expect((await sharp(upload!).metadata()).width).toBe(1280);
|
||||
expect(outputs[0]).toMatchObject({ width: 1280, height: 832, parentArtifactId: "photo-1" });
|
||||
const clip = await fixture("clip.mp4", 8, 8);
|
||||
await expect(
|
||||
run(operation, "generate_image", { prompt: "x", sourceArtifactId: "c" }, { c: clip }),
|
||||
).rejects.toThrow("The source must be an image");
|
||||
});
|
||||
|
||||
it("generates video in the preset's aspect", async () => {
|
||||
const { operation, emulator } = setup();
|
||||
const { outputs, result } = await run(operation, "generate_video", {
|
||||
prompt: "paper boat on a river",
|
||||
aspect: "vertical",
|
||||
});
|
||||
expect(outputs).toEqual([
|
||||
expect.objectContaining({
|
||||
name: "paper-boat-on-a-river.mp4",
|
||||
mimeType: "video/mp4",
|
||||
width: 480,
|
||||
height: 832,
|
||||
durationMs: 5042,
|
||||
}),
|
||||
]);
|
||||
expect(emulator.prompts[0]?.graph["6"]?.inputs).toMatchObject({ width: 480, height: 832 });
|
||||
expect(result.summary).toMatch(/^Generated a 5 s video with seed \d+\.$/);
|
||||
});
|
||||
|
||||
it("edits locally in order and keeps the source as parent", async () => {
|
||||
const { operation, emulator } = setup(null);
|
||||
const photo = await fixture("photo.jpg", 400, 300);
|
||||
const { outputs, result } = await run(
|
||||
operation,
|
||||
"edit_image",
|
||||
{
|
||||
artifactId: "photo-1",
|
||||
crop: { left: 0, top: 0, width: 300, height: 300 },
|
||||
rotate: 90,
|
||||
resize: { aspect: "vertical", fit: "blur" },
|
||||
format: "webp",
|
||||
},
|
||||
{ "photo-1": photo },
|
||||
);
|
||||
expect(outputs).toEqual([
|
||||
expect.objectContaining({
|
||||
name: "photo-edit.webp",
|
||||
mimeType: "image/webp",
|
||||
width: 1080,
|
||||
height: 1920,
|
||||
parentArtifactId: "photo-1",
|
||||
}),
|
||||
]);
|
||||
expect((await sharp(outputs[0]!.path).metadata()).format).toBe("webp");
|
||||
expect(result.summary).toBe("Edited photo.jpg into photo-edit.webp (1080×1920).");
|
||||
expect(emulator.requests).toHaveLength(0);
|
||||
await expect(
|
||||
run(operation, "edit_image", { artifactId: "photo-1", prompt: "snow" }, { "photo-1": photo }),
|
||||
).rejects.toThrow("ComfyUI is not configured");
|
||||
});
|
||||
|
||||
it("repaints an edit on ComfyUI at the edited size, deferring while the GPU is busy", async () => {
|
||||
const { operation, emulator } = setup();
|
||||
const photo = await fixture("photo.png", 300, 200);
|
||||
const input = {
|
||||
artifactId: "photo-1",
|
||||
crop: { left: 0, top: 0, width: 250, height: 200 },
|
||||
prompt: "snow",
|
||||
};
|
||||
const { outputs } = await run(operation, "edit_image", input, { "photo-1": photo });
|
||||
expect(outputs[0]).toMatchObject({ name: "photo-edit.png", width: 250, height: 200 });
|
||||
const [upload] = emulator.uploads.values();
|
||||
expect(await sharp(upload!).metadata()).toMatchObject({ width: 1152, height: 896 });
|
||||
emulator.gpu = { totalMb: 24_576, freeMb: 2_000, torchReservedMb: 0, torchFreeMb: 0 };
|
||||
await expect(run(operation, "edit_image", input, { "photo-1": photo })).rejects.toBeInstanceOf(
|
||||
MediaJobDeferredError,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("helpers", () => {
|
||||
it("sizes generations near the pixel budget", () => {
|
||||
expect(generationSize(1, 1024 * 1024, 64)).toEqual({ width: 1024, height: 1024 });
|
||||
expect(generationSize(16 / 9, 768 * 512, 32)).toEqual({ width: 832, height: 480 });
|
||||
expect(generationSize(21 / 9, 1024 * 1024, 64)).toEqual({ width: 1536, height: 640 });
|
||||
});
|
||||
|
||||
it("names files and titles edits briefly", () => {
|
||||
expect(fileStem(" Ünïcode — Straße 42 at night, very long prompt here ", "image")).toBe(
|
||||
"unicode-strasse-42-at-night-very",
|
||||
);
|
||||
expect(fileStem("夜の海", "video")).toBe("video");
|
||||
expect(
|
||||
editImageTitle({
|
||||
artifactId: "a",
|
||||
crop: { left: 0, top: 0, width: 1, height: 1 },
|
||||
rotate: 90,
|
||||
resize: { aspect: "vertical", fit: "crop" },
|
||||
format: "png",
|
||||
}),
|
||||
).toBe("crop · 90° · 9:16 · PNG");
|
||||
expect(editImageTitle({ artifactId: "a", resize: { width: 800, fit: "crop" } })).toBe("800×…");
|
||||
expect(editImageTitle({ artifactId: "a", prompt: "make it snow" })).toBe("make it snow");
|
||||
});
|
||||
});
|
||||
315
packages/adapters/src/media/image-operations.ts
Normal file
315
packages/adapters/src/media/image-operations.ts
Normal file
@ -0,0 +1,315 @@
|
||||
import { randomInt } from "node:crypto";
|
||||
import { writeFile } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import type {
|
||||
GeneratedMedia,
|
||||
GpuStatus,
|
||||
MediaJobContext,
|
||||
MediaJobInputFile,
|
||||
MediaOperationResult,
|
||||
} from "@rakazo/adapter-kit";
|
||||
import { MediaJobDeferredError } from "@rakazo/adapter-kit";
|
||||
import type {
|
||||
EditImageInput,
|
||||
GenerateImageInput,
|
||||
GenerateVideoInput,
|
||||
ImageFormat,
|
||||
MediaPresetId,
|
||||
} from "@rakazo/contracts";
|
||||
import {
|
||||
attachmentKindOf,
|
||||
EditImageInputSchema,
|
||||
GenerateImageInputSchema,
|
||||
GenerateVideoInputSchema,
|
||||
MEDIA_PRESETS,
|
||||
} from "@rakazo/contracts";
|
||||
import { attachmentExtensionForMimeType } from "@rakazo/core";
|
||||
import type { ComfyUiProvider } from "../comfyui-provider.js";
|
||||
import { createComfyUiProvider } from "../comfyui-provider.js";
|
||||
import type { MediaOperationDeps } from "../media-operations.js";
|
||||
import type { MediaOperationRegistration } from "../media-runtime.js";
|
||||
import {
|
||||
coverImage,
|
||||
editImageGeometry,
|
||||
encodeImage,
|
||||
IMAGE_FORMAT_MIME_TYPES,
|
||||
imageSize,
|
||||
sharpAvailable,
|
||||
} from "./image-edit.js";
|
||||
|
||||
/** Free GPU memory a generation needs before it starts; fixed, so a custom workflow cannot lower it. */
|
||||
export const IMAGE_GENERATION_MIN_VRAM_MB = 6_000;
|
||||
export const VIDEO_GENERATION_MIN_VRAM_MB = 12_000;
|
||||
/** Pixel budgets the default workflows are tuned for (SDXL-class images, LTX-Video clips). */
|
||||
const IMAGE_PIXELS = 1024 * 1024;
|
||||
const VIDEO_PIXELS = 768 * 512;
|
||||
const GPU_RETRY_MS = 60_000;
|
||||
|
||||
export interface ImageOperationsDeps extends MediaOperationDeps {
|
||||
/** Tests point the provider at an emulator. */
|
||||
fetch?: typeof globalThis.fetch;
|
||||
}
|
||||
|
||||
type ProviderSource = () => Promise<ComfyUiProvider | null>;
|
||||
|
||||
/** A size close to `pixels` with this aspect ratio, in multiples of `multiple`. */
|
||||
export function generationSize(aspect: number, pixels: number, multiple: number) {
|
||||
return {
|
||||
width: Math.max(multiple, Math.round(Math.sqrt(pixels * aspect) / multiple) * multiple),
|
||||
height: Math.max(multiple, Math.round(Math.sqrt(pixels / aspect) / multiple) * multiple),
|
||||
};
|
||||
}
|
||||
|
||||
function presetAspect(preset: MediaPresetId): number {
|
||||
return MEDIA_PRESETS[preset].width / MEDIA_PRESETS[preset].height;
|
||||
}
|
||||
|
||||
/** File name stem from a prompt: its first words in lowercase ASCII. */
|
||||
export function fileStem(text: string, fallback: string): string {
|
||||
const words = text
|
||||
.normalize("NFKD")
|
||||
.replace(/\p{M}/gu, "")
|
||||
.toLowerCase()
|
||||
.replace(/ß/g, "ss")
|
||||
.match(/[a-z0-9]+/g);
|
||||
return words?.slice(0, 6).join("-").slice(0, 48).replace(/-+$/, "") || fallback;
|
||||
}
|
||||
|
||||
/** Short card title for an edit without a prompt, such as "crop · 9:16 · PNG". */
|
||||
export function editImageTitle(input: EditImageInput): string {
|
||||
if (input.prompt) return input.prompt;
|
||||
const { resize } = input;
|
||||
const parts = [
|
||||
input.crop && "crop",
|
||||
input.rotate && `${input.rotate}°`,
|
||||
input.flip && "flip",
|
||||
resize &&
|
||||
(resize.aspect
|
||||
? MEDIA_PRESETS[resize.aspect].aspect
|
||||
: `${resize.width ?? "…"}×${resize.height ?? "…"}`),
|
||||
input.extend && "extend",
|
||||
input.format?.toUpperCase(),
|
||||
];
|
||||
return parts.filter(Boolean).join(" · ");
|
||||
}
|
||||
|
||||
function reportProgress(context: MediaJobContext) {
|
||||
return (fraction: number) => {
|
||||
context.progress(fraction).catch(() => undefined);
|
||||
};
|
||||
}
|
||||
|
||||
async function requireProvider(source: ProviderSource): Promise<ComfyUiProvider> {
|
||||
const provider = await source();
|
||||
if (!provider) throw new Error("ComfyUI is not configured");
|
||||
return provider;
|
||||
}
|
||||
|
||||
/** Defers instead of starting a generation the GPU has no room for right now. */
|
||||
async function assertGpuReady(
|
||||
provider: ComfyUiProvider,
|
||||
minVramMb: number,
|
||||
signal: AbortSignal,
|
||||
): Promise<void> {
|
||||
const status: GpuStatus | undefined = await provider.gpuStatus(signal);
|
||||
if (status && status.vramFreeMb < Math.min(minVramMb, status.vramTotalMb)) {
|
||||
throw new MediaJobDeferredError("Waiting for free GPU memory", GPU_RETRY_MS);
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchImage(
|
||||
context: MediaJobContext,
|
||||
artifactId: string,
|
||||
): Promise<MediaJobInputFile & { width: number; height: number }> {
|
||||
const file = await context.fetchArtifact(artifactId);
|
||||
if (file.kind !== "image") throw new Error("The source must be an image");
|
||||
return { ...file, ...(await imageSize(file.path)) };
|
||||
}
|
||||
|
||||
async function addGeneratedOutputs(
|
||||
context: MediaJobContext,
|
||||
results: GeneratedMedia[],
|
||||
stem: string,
|
||||
parentArtifactId: string | undefined,
|
||||
): Promise<void> {
|
||||
for (const [index, result] of results.entries()) {
|
||||
const extension = attachmentExtensionForMimeType(result.mimeType);
|
||||
const name = `${stem}${results.length > 1 ? `-${index + 1}` : ""}${extension}`;
|
||||
const path = join(context.tempDir, name);
|
||||
await writeFile(path, result.bytes);
|
||||
const size =
|
||||
attachmentKindOf(result.mimeType) === "image"
|
||||
? await imageSize(path)
|
||||
: { width: result.width, height: result.height };
|
||||
await context.addOutput({
|
||||
path,
|
||||
name,
|
||||
mimeType: result.mimeType,
|
||||
...size,
|
||||
...(result.durationMs !== undefined ? { durationMs: result.durationMs } : {}),
|
||||
...(parentArtifactId ? { parentArtifactId } : {}),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async function generateImage(
|
||||
context: MediaJobContext,
|
||||
providers: ProviderSource,
|
||||
): Promise<MediaOperationResult> {
|
||||
const input = GenerateImageInputSchema.parse(context.input);
|
||||
const provider = await requireProvider(providers);
|
||||
const source = input.sourceArtifactId
|
||||
? await fetchImage(context, input.sourceArtifactId)
|
||||
: undefined;
|
||||
const aspect = input.aspect
|
||||
? presetAspect(input.aspect)
|
||||
: source
|
||||
? source.width / source.height
|
||||
: 1;
|
||||
const size =
|
||||
input.width && input.height
|
||||
? { width: input.width, height: input.height }
|
||||
: generationSize(aspect, IMAGE_PIXELS, 64);
|
||||
const seed = input.seed ?? randomInt(2 ** 48 - 1);
|
||||
const images = await provider.generateImage({
|
||||
prompt: input.prompt,
|
||||
negativePrompt: input.negativePrompt,
|
||||
...size,
|
||||
count: input.count,
|
||||
seed,
|
||||
source: source && {
|
||||
bytes: await coverImage(source.path, size.width, size.height),
|
||||
mimeType: "image/png",
|
||||
},
|
||||
strength: input.strength,
|
||||
signal: context.signal,
|
||||
onProgress: reportProgress(context),
|
||||
});
|
||||
await addGeneratedOutputs(context, images, fileStem(input.prompt, "image"), source?.artifactId);
|
||||
const count = images.length === 1 ? "1 image" : `${images.length} images`;
|
||||
return { summary: `Generated ${count} with seed ${seed}.` };
|
||||
}
|
||||
|
||||
async function generateVideo(
|
||||
context: MediaJobContext,
|
||||
providers: ProviderSource,
|
||||
): Promise<MediaOperationResult> {
|
||||
const input = GenerateVideoInputSchema.parse(context.input);
|
||||
const provider = await requireProvider(providers);
|
||||
const source = input.sourceArtifactId
|
||||
? await fetchImage(context, input.sourceArtifactId)
|
||||
: undefined;
|
||||
const aspect = input.aspect
|
||||
? presetAspect(input.aspect)
|
||||
: source
|
||||
? source.width / source.height
|
||||
: presetAspect("landscape");
|
||||
const size = generationSize(aspect, VIDEO_PIXELS, 32);
|
||||
const seed = input.seed ?? randomInt(2 ** 48 - 1);
|
||||
const clips = await provider.generateVideo({
|
||||
prompt: input.prompt,
|
||||
negativePrompt: input.negativePrompt,
|
||||
...size,
|
||||
durationSeconds: input.durationSeconds,
|
||||
seed,
|
||||
source: source && {
|
||||
bytes: await coverImage(source.path, size.width, size.height),
|
||||
mimeType: "image/png",
|
||||
},
|
||||
signal: context.signal,
|
||||
onProgress: reportProgress(context),
|
||||
});
|
||||
await addGeneratedOutputs(context, clips, fileStem(input.prompt, "video"), source?.artifactId);
|
||||
return { summary: `Generated a ${input.durationSeconds} s video with seed ${seed}.` };
|
||||
}
|
||||
|
||||
function defaultFormat(mimeType: string): ImageFormat {
|
||||
if (mimeType === "image/jpeg") return "jpeg";
|
||||
if (mimeType === "image/webp") return "webp";
|
||||
return "png";
|
||||
}
|
||||
|
||||
async function editImage(
|
||||
context: MediaJobContext,
|
||||
providers: ProviderSource,
|
||||
): Promise<MediaOperationResult> {
|
||||
const input = EditImageInputSchema.parse(context.input);
|
||||
const source = await fetchImage(context, input.artifactId);
|
||||
let png: Uint8Array = await editImageGeometry(source.path, input);
|
||||
if (input.prompt) {
|
||||
const provider = await requireProvider(providers);
|
||||
await assertGpuReady(provider, IMAGE_GENERATION_MIN_VRAM_MB, context.signal);
|
||||
const { width, height } = await imageSize(png);
|
||||
// Repaint near the model's native size, then return to the edited size.
|
||||
const size = generationSize(width / height, IMAGE_PIXELS, 64);
|
||||
const [repainted] = await provider.generateImage({
|
||||
prompt: input.prompt,
|
||||
...size,
|
||||
source: { bytes: await coverImage(png, size.width, size.height), mimeType: "image/png" },
|
||||
strength: input.strength,
|
||||
signal: context.signal,
|
||||
onProgress: reportProgress(context),
|
||||
});
|
||||
if (!repainted) throw new Error("ComfyUI produced no image");
|
||||
png = await coverImage(repainted.bytes, width, height);
|
||||
}
|
||||
const format = input.format ?? defaultFormat(source.mimeType);
|
||||
const encoded = await encodeImage(png, format);
|
||||
const mimeType = IMAGE_FORMAT_MIME_TYPES[format];
|
||||
const stem = source.name.replace(/\.[^.]+$/, "") || "image";
|
||||
const name = `${stem}-edit${attachmentExtensionForMimeType(mimeType)}`;
|
||||
const path = join(context.tempDir, name);
|
||||
await writeFile(path, encoded.bytes);
|
||||
await context.addOutput({
|
||||
path,
|
||||
name,
|
||||
mimeType,
|
||||
width: encoded.width,
|
||||
height: encoded.height,
|
||||
parentArtifactId: source.artifactId,
|
||||
});
|
||||
return { summary: `Edited ${source.name} into ${name} (${encoded.width}×${encoded.height}).` };
|
||||
}
|
||||
|
||||
/**
|
||||
* Media operations for image and video generation on ComfyUI and for image edits with sharp.
|
||||
* Generation is available when ComfyUI is configured; edits whenever sharp loads (repainting
|
||||
* an edit also needs ComfyUI).
|
||||
*/
|
||||
export function createImageOperations(deps: ImageOperationsDeps): MediaOperationRegistration[] {
|
||||
const providers: ProviderSource = async () => {
|
||||
const { comfyui } = await deps.loadSettings();
|
||||
return comfyui ? createComfyUiProvider({ ...comfyui, fetch: deps.fetch }) : null;
|
||||
};
|
||||
const comfyUiConfigured = async () => Boolean((await deps.loadSettings()).comfyui);
|
||||
const status = async (signal?: AbortSignal) => (await providers())?.gpuStatus(signal);
|
||||
const sourceIds = (input: { sourceArtifactId?: string }) =>
|
||||
input.sourceArtifactId ? [input.sourceArtifactId] : [];
|
||||
const image: MediaOperationRegistration<GenerateImageInput> = {
|
||||
operation: "generate_image",
|
||||
inputSchema: GenerateImageInputSchema,
|
||||
title: (input) => input.prompt,
|
||||
available: comfyUiConfigured,
|
||||
handler: { operation: "generate_image", run: (context) => generateImage(context, providers) },
|
||||
gpu: { minVramMb: IMAGE_GENERATION_MIN_VRAM_MB, status },
|
||||
artifactIds: sourceIds,
|
||||
};
|
||||
const video: MediaOperationRegistration<GenerateVideoInput> = {
|
||||
operation: "generate_video",
|
||||
inputSchema: GenerateVideoInputSchema,
|
||||
title: (input) => input.prompt,
|
||||
available: comfyUiConfigured,
|
||||
handler: { operation: "generate_video", run: (context) => generateVideo(context, providers) },
|
||||
gpu: { minVramMb: VIDEO_GENERATION_MIN_VRAM_MB, status },
|
||||
artifactIds: sourceIds,
|
||||
};
|
||||
const edit: MediaOperationRegistration<EditImageInput> = {
|
||||
operation: "edit_image",
|
||||
inputSchema: EditImageInputSchema,
|
||||
title: editImageTitle,
|
||||
available: sharpAvailable,
|
||||
handler: { operation: "edit_image", run: (context) => editImage(context, providers) },
|
||||
artifactIds: (input) => [input.artifactId],
|
||||
};
|
||||
return [image, video, edit] as MediaOperationRegistration[];
|
||||
}
|
||||
393
packages/adapters/src/testing/comfyui-emulator.ts
Normal file
393
packages/adapters/src/testing/comfyui-emulator.ts
Normal file
@ -0,0 +1,393 @@
|
||||
import { createHash } from "node:crypto";
|
||||
|
||||
/** A 1 s, 32×32 H.264 MP4 (ffmpeg testsrc), small enough to inline. */
|
||||
const TINY_MP4_BASE64 = [
|
||||
"AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAM+bW9vdgAAAGxtdmhkAAAAAAAAAAAAAAAAAAAD6AAAA+gAAQAAAQAA",
|
||||
"AAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAA",
|
||||
"Ao10cmFrAAAAXHRraGQAAAADAAAAAAAAAAAAAAABAAAAAAAAA+gAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAA",
|
||||
"AAAAAAAAAAAAAABAAAAAACAAAAAgAAAAAAAkZWR0cwAAABxlbHN0AAAAAAAAAAEAAAPoAAAQAAABAAAAAAIFbWRpYQAAACBtZGhk",
|
||||
"AAAAAAAAAAAAAAAAAABAAAAAQABVxAAAAAAALWhkbHIAAAAAAAAAAHZpZGUAAAAAAAAAAAAAAABWaWRlb0hhbmRsZXIAAAABsG1p",
|
||||
"bmYAAAAUdm1oZAAAAAEAAAAAAAAAAAAAACRkaW5mAAAAHGRyZWYAAAAAAAAAAQAAAAx1cmwgAAAAAQAAAXBzdGJsAAAAwHN0c2QA",
|
||||
"AAAAAAAAAQAAALBhdmMxAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAACAAIABIAAAASAAAAAAAAAABDExhdmMgbGlieDI2NAAAAAAA",
|
||||
"AAAAAAAAAAAAAAAAAAAAGP//AAAANmF2Y0MBZAAK/+EAGGdkAAqscgRJbARAAAADAEAAAAQDxIlhGAEAB2joQ4OSyLD9+PgAAAAA",
|
||||
"EHBhc3AAAAABAAAAAQAAABRidHJ0AAAAAAAADAgAAAwIAAAAGHN0dHMAAAAAAAAAAQAAAAgAAAgAAAAAFHN0c3MAAAAAAAAAAQAA",
|
||||
"AAEAAAAYY3R0cwAAAAAAAAABAAAACAAAEAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAAgAAAABAAAANHN0c3oAAAAAAAAAAAAAAAgA",
|
||||
"AADTAAAADAAAABgAAAAdAAAAGgAAAB0AAAAfAAAAFwAAABRzdGNvAAAAAAAAAAEAAANuAAAAPXVkdGEAAAA1bWV0YQAAAAAAAAAh",
|
||||
"aGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAIaWxzdAAAAAhmcmVlAAABiW1kYXQAAADPZYiBABP/115QL0x5O3/uan23",
|
||||
"okcn80QwiyZelIq0iY0T1naTwGHVgXgsN7YBwlEikXvo+AyEjKk/afikvNFJbO7NjVpKr2f7e+xmO5rxkaelhrWKVT83gKOM5wKs",
|
||||
"QoFpCKDo56fFf6IYs3U1LAHAxZY3IIAUCFca5RyJ/ibKbmkxBCXXiBuUCCuL9FhMjcqhW2misGdpwuLgknHPr/37MowWA2zcOn9B",
|
||||
"OHR5mfgyt4/0x3o95hpnI+2axZFkXHhYwK55h134UUvr+r3B9b/xAAAACEGaCC2JP6KAAAAAFEGaEEvwIGTKYTe7qicJ6jU1YvCB",
|
||||
"AAAAGUGaGGk8CB5MpgTfu1UOB2sZPRDVFTCp+IAAAAAWQZogiXwIH5MpgER/wf8hVV5nVt8U2QAAABlBmiipvAgf5MpgBFfFjke1",
|
||||
"iXbU9gFuIojBAAAAG0GaMMn8CB/5MpgARn/KNfso3F7XaeLq3d3VjQAAABNBmjjojwIH/5MpgAEIf9P6j9Ph",
|
||||
].join("");
|
||||
|
||||
export const COMFYUI_EMULATOR_MP4 = new Uint8Array(Buffer.from(TINY_MP4_BASE64, "base64"));
|
||||
|
||||
const MIB = 1024 * 1024;
|
||||
|
||||
type EmulatedNode = { class_type: string; inputs: Record<string, unknown> };
|
||||
type EmulatedGraph = Record<string, EmulatedNode>;
|
||||
type FileRef = { filename: string; subfolder: string; type: string };
|
||||
|
||||
interface EmulatedPrompt {
|
||||
id: string;
|
||||
graph: EmulatedGraph;
|
||||
polls: number;
|
||||
failure?: string;
|
||||
interrupted?: boolean;
|
||||
/** Dropped from the queue before it ran. */
|
||||
removed?: boolean;
|
||||
/** Reported as finished by the history. */
|
||||
done?: boolean;
|
||||
outputs?: Record<string, Record<string, unknown>>;
|
||||
}
|
||||
|
||||
export interface EmulatedComfyUiRequest {
|
||||
method: string;
|
||||
path: string;
|
||||
authorization: string | null;
|
||||
body?: unknown;
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return Boolean(value) && typeof value === "object" && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function json(body: unknown, status = 200): Response {
|
||||
return Response.json(body, { status });
|
||||
}
|
||||
|
||||
/**
|
||||
* In-process substitute for a ComfyUI server's HTTP API, so the real provider runs in offline
|
||||
* tests: it validates submitted workflows like ComfyUI does, "runs" them after a few history
|
||||
* polls, and returns small generated PNGs or an MP4 fixture.
|
||||
*/
|
||||
export class ComfyUiEmulator {
|
||||
readonly requests: EmulatedComfyUiRequest[] = [];
|
||||
readonly prompts: EmulatedPrompt[] = [];
|
||||
readonly uploads = new Map<string, Uint8Array>();
|
||||
readonly deleted: string[] = [];
|
||||
readonly interrupted: string[] = [];
|
||||
checkpoints = ["sd_xl_base_1.0.safetensors", "ltx-video-2b-v0.9.5.safetensors"];
|
||||
/** `options` answers like current ComfyUI (`["COMBO", {options}]`), `legacy` like older ones. */
|
||||
checkpointListFormat: "options" | "legacy" = "options";
|
||||
/** null reports a CPU-only server. */
|
||||
gpu: { totalMb: number; freeMb: number; torchReservedMb: number; torchFreeMb: number } | null = {
|
||||
totalMb: 24_576,
|
||||
freeMb: 20_000,
|
||||
torchReservedMb: 0,
|
||||
torchFreeMb: 0,
|
||||
};
|
||||
/** History polls that still report a submitted prompt as running. */
|
||||
pendingPolls = 1;
|
||||
/** Answer the next request with this HTTP status. */
|
||||
failNextStatus: number | undefined;
|
||||
/** Fail the next submitted prompt while it runs, with this exception message. */
|
||||
failNextExecution: string | undefined;
|
||||
private readonly files = new Map<string, Uint8Array>();
|
||||
private nextPrompt = 1;
|
||||
|
||||
/** `basePath` serves the API under a prefix, like ComfyUI behind a reverse proxy. */
|
||||
constructor(private readonly options: { apiKey?: string; basePath?: string } = {}) {}
|
||||
|
||||
readonly fetch: typeof fetch = async (input, init) => {
|
||||
init?.signal?.throwIfAborted();
|
||||
const url = new URL(input instanceof Request ? input.url : String(input));
|
||||
const method = init?.method ?? "GET";
|
||||
const headers = new Headers(init?.headers);
|
||||
const authorization = headers.get("authorization");
|
||||
const prefix = `/${(this.options.basePath ?? "").replace(/^\/+|\/+$/g, "")}/`.replace(
|
||||
"//",
|
||||
"/",
|
||||
);
|
||||
if (!url.pathname.startsWith(prefix)) return json({ error: "Not found" }, 404);
|
||||
const path = url.pathname.slice(prefix.length);
|
||||
const body =
|
||||
typeof init?.body === "string" && headers.get("content-type")?.includes("json")
|
||||
? JSON.parse(init.body)
|
||||
: undefined;
|
||||
this.requests.push({ method, path, authorization, body });
|
||||
if (this.failNextStatus) {
|
||||
const status = this.failNextStatus;
|
||||
this.failNextStatus = undefined;
|
||||
return new Response("Upstream failure", { status });
|
||||
}
|
||||
if (this.options.apiKey && authorization !== `Bearer ${this.options.apiKey}`) {
|
||||
return json({ error: "Unauthorized" }, 401);
|
||||
}
|
||||
const route = `${method} ${path.split("/")[0]}`;
|
||||
if (route === "GET system_stats") return this.systemStats();
|
||||
if (route === "GET object_info" && path === "object_info/CheckpointLoaderSimple") {
|
||||
return this.checkpointInfo();
|
||||
}
|
||||
if (route === "POST prompt") return this.submit(body);
|
||||
if (route === "GET history") return this.history(decodeURIComponent(path.slice(8)));
|
||||
if (route === "GET view") return this.view(url.searchParams);
|
||||
if (route === "POST upload") return this.upload(input, init);
|
||||
if (route === "GET queue") return this.queue();
|
||||
if (route === "POST queue") {
|
||||
for (const id of Array.isArray(body?.delete) ? body.delete : []) {
|
||||
this.deleted.push(id);
|
||||
// Like ComfyUI, deleting only drops pending prompts; a running one needs an interrupt.
|
||||
const prompt = this.prompts.find((entry) => entry.id === id);
|
||||
if (prompt && prompt !== this.unfinished()[0]) prompt.removed = true;
|
||||
}
|
||||
return json({});
|
||||
}
|
||||
if (route === "POST interrupt") {
|
||||
const id = String(body?.prompt_id ?? "");
|
||||
this.interrupted.push(id);
|
||||
const running = this.unfinished()[0];
|
||||
if (running && (!id || running.id === id)) running.interrupted = true;
|
||||
return json({});
|
||||
}
|
||||
return json({ error: "Not found" }, 404);
|
||||
};
|
||||
|
||||
private systemStats(): Response {
|
||||
const gpu = this.gpu;
|
||||
return json({
|
||||
system: { os: "posix", comfyui_version: "0.3.60", embedded_python: false },
|
||||
devices: [
|
||||
gpu
|
||||
? {
|
||||
name: "cuda:0 Emulated GPU",
|
||||
type: "cuda",
|
||||
index: 0,
|
||||
vram_total: gpu.totalMb * MIB,
|
||||
vram_free: gpu.freeMb * MIB,
|
||||
torch_vram_total: gpu.torchReservedMb * MIB,
|
||||
torch_vram_free: gpu.torchFreeMb * MIB,
|
||||
}
|
||||
: { name: "cpu", type: "cpu", index: null, vram_total: 64 * 1024 * MIB, vram_free: 0 },
|
||||
],
|
||||
});
|
||||
}
|
||||
|
||||
private checkpointInfo(): Response {
|
||||
const spec =
|
||||
this.checkpointListFormat === "legacy"
|
||||
? [this.checkpoints, { tooltip: "Checkpoint" }]
|
||||
: ["COMBO", { options: this.checkpoints, tooltip: "Checkpoint" }];
|
||||
return json({
|
||||
CheckpointLoaderSimple: {
|
||||
input: { required: { ckpt_name: spec } },
|
||||
output: ["MODEL", "CLIP", "VAE"],
|
||||
name: "CheckpointLoaderSimple",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private submit(body: unknown): Response {
|
||||
const graph = isRecord(body) ? body.prompt : undefined;
|
||||
if (!isRecord(graph) || typeof (body as { client_id?: unknown }).client_id !== "string") {
|
||||
return json(
|
||||
{ error: { type: "invalid_prompt", message: "Cannot execute because prompt is invalid" } },
|
||||
400,
|
||||
);
|
||||
}
|
||||
const nodeErrors: Record<string, unknown> = {};
|
||||
const reject = (id: string, node: unknown, message: string, details: string) => {
|
||||
nodeErrors[id] = {
|
||||
errors: [{ type: "invalid", message, details, extra_info: {} }],
|
||||
dependent_outputs: [],
|
||||
class_type: isRecord(node) ? node.class_type : undefined,
|
||||
};
|
||||
};
|
||||
for (const [id, node] of Object.entries(graph)) {
|
||||
if (!isRecord(node) || typeof node.class_type !== "string" || !isRecord(node.inputs)) {
|
||||
reject(id, node, "Invalid node", "class_type and inputs are required");
|
||||
continue;
|
||||
}
|
||||
for (const [name, value] of Object.entries(node.inputs)) {
|
||||
if (typeof value === "string" && value.includes("{{")) {
|
||||
reject(id, node, "Unfilled placeholder", `${name}: ${value}`);
|
||||
}
|
||||
if (Array.isArray(value) && !isRecord(graph[String(value[0])])) {
|
||||
reject(id, node, "Missing linked node", `${name}: ${String(value[0])}`);
|
||||
}
|
||||
}
|
||||
const checkpoint = node.inputs.ckpt_name;
|
||||
if (typeof checkpoint === "string" && !this.checkpoints.includes(checkpoint)) {
|
||||
reject(id, node, "Value not in list", `ckpt_name: '${checkpoint}' not in list`);
|
||||
}
|
||||
if (node.class_type === "LoadImage" && !this.uploads.has(String(node.inputs.image))) {
|
||||
reject(id, node, "Invalid image file", `image - ${String(node.inputs.image)}`);
|
||||
}
|
||||
}
|
||||
if (Object.keys(nodeErrors).length > 0) {
|
||||
return json(
|
||||
{
|
||||
error: {
|
||||
type: "prompt_outputs_failed_validation",
|
||||
message: "Prompt outputs failed validation",
|
||||
},
|
||||
node_errors: nodeErrors,
|
||||
},
|
||||
400,
|
||||
);
|
||||
}
|
||||
const id = `prompt-${this.nextPrompt++}`;
|
||||
this.prompts.push({
|
||||
id,
|
||||
graph: graph as EmulatedGraph,
|
||||
polls: 0,
|
||||
failure: this.failNextExecution,
|
||||
});
|
||||
this.failNextExecution = undefined;
|
||||
return json({ prompt_id: id, number: this.nextPrompt, node_errors: {} });
|
||||
}
|
||||
|
||||
private async history(id: string): Promise<Response> {
|
||||
const prompt = this.prompts.find((entry) => entry.id === id);
|
||||
if (!prompt || prompt.removed) return json({});
|
||||
if (!prompt.done && !prompt.interrupted && prompt.polls < this.pendingPolls) {
|
||||
prompt.polls += 1;
|
||||
return json({});
|
||||
}
|
||||
prompt.done = true;
|
||||
if (prompt.interrupted) {
|
||||
return json({
|
||||
[id]: {
|
||||
prompt: [0, id, prompt.graph, {}, []],
|
||||
outputs: {},
|
||||
status: {
|
||||
status_str: "error",
|
||||
completed: false,
|
||||
messages: [["execution_interrupted", { prompt_id: id, node_id: "5" }]],
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
if (prompt.failure) {
|
||||
return json({
|
||||
[id]: {
|
||||
prompt: [0, id, prompt.graph, {}, []],
|
||||
outputs: {},
|
||||
status: {
|
||||
status_str: "error",
|
||||
completed: false,
|
||||
messages: [
|
||||
["execution_start", { prompt_id: id }],
|
||||
[
|
||||
"execution_error",
|
||||
{
|
||||
prompt_id: id,
|
||||
node_id: "5",
|
||||
node_type: "KSampler",
|
||||
exception_message: prompt.failure,
|
||||
exception_type: "RuntimeError",
|
||||
traceback: ["..."],
|
||||
},
|
||||
],
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
prompt.outputs ??= await this.execute(prompt);
|
||||
return json({
|
||||
[id]: {
|
||||
prompt: [0, id, prompt.graph, {}, Object.keys(prompt.outputs)],
|
||||
outputs: prompt.outputs,
|
||||
status: { status_str: "success", completed: true, messages: [] },
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private async execute(prompt: EmulatedPrompt): Promise<Record<string, Record<string, unknown>>> {
|
||||
const outputs: Record<string, Record<string, unknown>> = {};
|
||||
let counter = 0;
|
||||
const save = (bytes: Uint8Array, extension: string, type: string): FileRef => {
|
||||
counter += 1;
|
||||
const ref = {
|
||||
filename: `${prompt.id}_${String(counter).padStart(5, "0")}_.${extension}`,
|
||||
subfolder: type === "output" ? "evoo" : "",
|
||||
type,
|
||||
};
|
||||
this.files.set(fileKey(ref), bytes);
|
||||
return ref;
|
||||
};
|
||||
for (const [id, node] of Object.entries(prompt.graph)) {
|
||||
if (node.class_type === "SaveImage" || node.class_type === "PreviewImage") {
|
||||
const type = node.class_type === "SaveImage" ? "output" : "temp";
|
||||
outputs[id] = { images: [save(await this.renderImage(prompt.graph), "png", type)] };
|
||||
} else if (node.class_type === "SaveVideo") {
|
||||
outputs[id] = { images: [save(COMFYUI_EMULATOR_MP4, "mp4", "output")], animated: [true] };
|
||||
} else if (node.class_type === "VHS_VideoCombine") {
|
||||
outputs[id] = { gifs: [save(COMFYUI_EMULATOR_MP4, "mp4", "output")] };
|
||||
}
|
||||
}
|
||||
return outputs;
|
||||
}
|
||||
|
||||
/** A flat color derived from the seed, at the workflow's size (or its input image's). */
|
||||
private async renderImage(graph: EmulatedGraph): Promise<Uint8Array> {
|
||||
const { default: sharp } = await import("sharp");
|
||||
const nodes = Object.values(graph);
|
||||
const input = (name: string) =>
|
||||
nodes.map((node) => node.inputs[name]).find((value) => typeof value === "number");
|
||||
let width = Number(input("width") ?? 0);
|
||||
let height = Number(input("height") ?? 0);
|
||||
const loaded = nodes.find((node) => node.class_type === "LoadImage");
|
||||
if (loaded && !(width && height)) {
|
||||
const metadata = await sharp(this.uploads.get(String(loaded.inputs.image))).metadata();
|
||||
width = metadata.width ?? 64;
|
||||
height = metadata.height ?? 64;
|
||||
}
|
||||
const seed = String(input("seed") ?? input("noise_seed") ?? 0);
|
||||
const [r = 0, g = 0, b = 0] = createHash("sha256").update(seed).digest();
|
||||
const png = await sharp({
|
||||
create: { width: width || 64, height: height || 64, channels: 3, background: { r, g, b } },
|
||||
})
|
||||
.png()
|
||||
.toBuffer();
|
||||
return new Uint8Array(png);
|
||||
}
|
||||
|
||||
private view(params: URLSearchParams): Response {
|
||||
const bytes = this.files.get(
|
||||
fileKey({
|
||||
filename: params.get("filename") ?? "",
|
||||
subfolder: params.get("subfolder") ?? "",
|
||||
type: params.get("type") ?? "",
|
||||
}),
|
||||
);
|
||||
if (!bytes) return new Response("Not found", { status: 404 });
|
||||
return new Response(Buffer.from(bytes), {
|
||||
headers: { "Content-Type": "application/octet-stream" },
|
||||
});
|
||||
}
|
||||
|
||||
private async upload(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
|
||||
const form = await new Request(String(input), {
|
||||
method: "POST",
|
||||
headers: init?.headers,
|
||||
body: init?.body,
|
||||
}).formData();
|
||||
const image = form.get("image");
|
||||
if (!(image instanceof File) || form.get("type") !== "input") {
|
||||
return json({ error: "Invalid upload" }, 400);
|
||||
}
|
||||
this.uploads.set(image.name, new Uint8Array(await image.arrayBuffer()));
|
||||
return json({ name: image.name, subfolder: "", type: "input" });
|
||||
}
|
||||
|
||||
private unfinished(): EmulatedPrompt[] {
|
||||
return this.prompts.filter((prompt) => !prompt.done && !prompt.removed);
|
||||
}
|
||||
|
||||
private queue(): Response {
|
||||
const running = this.unfinished();
|
||||
return json({
|
||||
queue_running: running.slice(0, 1).map((prompt) => [0, prompt.id, prompt.graph, {}, []]),
|
||||
queue_pending: running
|
||||
.slice(1)
|
||||
.map((prompt, index) => [index + 1, prompt.id, prompt.graph, {}, []]),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function fileKey(ref: FileRef): string {
|
||||
return `${ref.type}/${ref.subfolder}/${ref.filename}`;
|
||||
}
|
||||
@ -10,6 +10,8 @@ export * from "./integration-settings.js";
|
||||
export * from "./local-settings.js";
|
||||
export * from "./mcp.js";
|
||||
export * from "./media.js";
|
||||
export * from "./media-generation.js";
|
||||
export * from "./media-jobs.js";
|
||||
export * from "./media-settings.js";
|
||||
export * from "./media-transcription.js";
|
||||
export * from "./model-tiers.js";
|
||||
|
||||
42
packages/contracts/src/media-generation.test.ts
Normal file
42
packages/contracts/src/media-generation.test.ts
Normal file
@ -0,0 +1,42 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
EditImageInputSchema,
|
||||
GenerateImageInputSchema,
|
||||
GenerateVideoInputSchema,
|
||||
} from "./media-generation.js";
|
||||
|
||||
describe("media generation inputs", () => {
|
||||
it("fills defaults and drops unknown fields such as continue_after", () => {
|
||||
expect(GenerateImageInputSchema.parse({ prompt: " a fox ", continue_after: true })).toEqual({
|
||||
prompt: "a fox",
|
||||
count: 1,
|
||||
});
|
||||
expect(GenerateVideoInputSchema.parse({ prompt: "waves" })).toEqual({
|
||||
prompt: "waves",
|
||||
durationSeconds: 5,
|
||||
});
|
||||
});
|
||||
|
||||
it("bounds image count, size and video duration", () => {
|
||||
expect(GenerateImageInputSchema.safeParse({ prompt: "x", count: 5 }).success).toBe(false);
|
||||
expect(GenerateImageInputSchema.safeParse({ prompt: "x", width: 512 }).success).toBe(false);
|
||||
expect(
|
||||
GenerateImageInputSchema.safeParse({ prompt: "x", width: 512, height: 4096 }).success,
|
||||
).toBe(false);
|
||||
expect(GenerateVideoInputSchema.safeParse({ prompt: "x", durationSeconds: 11 }).success).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
it("needs at least one edit and valid edit values", () => {
|
||||
expect(EditImageInputSchema.safeParse({ artifactId: "a" }).success).toBe(false);
|
||||
expect(EditImageInputSchema.safeParse({ artifactId: "a", rotate: 45 }).success).toBe(false);
|
||||
expect(EditImageInputSchema.safeParse({ artifactId: "a", resize: {} }).success).toBe(false);
|
||||
expect(
|
||||
EditImageInputSchema.safeParse({ artifactId: "a", extend: { top: 4, color: "red" } }).success,
|
||||
).toBe(false);
|
||||
expect(
|
||||
EditImageInputSchema.parse({ artifactId: "a", resize: { aspect: "vertical" }, rotate: 90 }),
|
||||
).toEqual({ artifactId: "a", resize: { aspect: "vertical", fit: "crop" }, rotate: 90 });
|
||||
});
|
||||
});
|
||||
107
packages/contracts/src/media-generation.ts
Normal file
107
packages/contracts/src/media-generation.ts
Normal file
@ -0,0 +1,107 @@
|
||||
import * as z from "zod";
|
||||
import { Id } from "./ids.js";
|
||||
import { MediaFitSchema, MediaPresetIdSchema } from "./media.js";
|
||||
|
||||
export const MEDIA_IMAGE_MAX_COUNT = 4;
|
||||
export const MEDIA_VIDEO_MAX_SECONDS = 10;
|
||||
|
||||
const Prompt = z.string().trim().min(1).max(4000);
|
||||
const NegativePrompt = z.string().trim().max(4000).optional().describe("What to avoid.");
|
||||
const Seed = z
|
||||
.number()
|
||||
.int()
|
||||
.min(0)
|
||||
.max(Number.MAX_SAFE_INTEGER)
|
||||
.optional()
|
||||
.describe("Reuse a seed to reproduce a result.");
|
||||
const Strength = z.number().min(0).max(1).optional();
|
||||
const ASPECTS = "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9";
|
||||
const aspect = (fallback: string) => MediaPresetIdSchema.describe(`${ASPECTS}. ${fallback}`);
|
||||
const Side = z.number().int().min(64).max(2048);
|
||||
const Pixels = z.number().int().min(0).max(16_384);
|
||||
|
||||
export const GenerateImageInputSchema = z
|
||||
.object({
|
||||
prompt: Prompt.describe("What the image shows."),
|
||||
negativePrompt: NegativePrompt,
|
||||
aspect: aspect("Default square, or the source image's aspect ratio.").optional(),
|
||||
width: Side.optional().describe("Exact width in pixels, with height; overrides aspect."),
|
||||
height: Side.optional(),
|
||||
count: z.number().int().min(1).max(MEDIA_IMAGE_MAX_COUNT).default(1),
|
||||
seed: Seed,
|
||||
sourceArtifactId: Id.optional().describe("Image artifact to start from (image-to-image)."),
|
||||
strength: Strength.describe("With a source: how far to move away from it (default 0.6)."),
|
||||
})
|
||||
.refine((input) => (input.width === undefined) === (input.height === undefined), {
|
||||
message: "Set width and height together",
|
||||
path: ["width"],
|
||||
});
|
||||
export type GenerateImageInput = z.infer<typeof GenerateImageInputSchema>;
|
||||
|
||||
export const GenerateVideoInputSchema = z.object({
|
||||
prompt: Prompt.describe("What happens in the clip: subject, motion, camera."),
|
||||
negativePrompt: NegativePrompt,
|
||||
aspect: aspect("Default landscape, or the start frame's aspect ratio.").optional(),
|
||||
durationSeconds: z.number().min(1).max(MEDIA_VIDEO_MAX_SECONDS).default(5),
|
||||
seed: Seed,
|
||||
sourceArtifactId: Id.optional().describe("Image artifact to use as the first frame."),
|
||||
});
|
||||
export type GenerateVideoInput = z.infer<typeof GenerateVideoInputSchema>;
|
||||
|
||||
export const ImageFormatSchema = z.enum(["png", "jpeg", "webp"]);
|
||||
export type ImageFormat = z.infer<typeof ImageFormatSchema>;
|
||||
|
||||
/** Edits apply in this order: crop, rotate, flip, resize, extend, repaint, format. */
|
||||
export const EditImageInputSchema = z
|
||||
.object({
|
||||
artifactId: Id.describe("Image artifact to edit."),
|
||||
crop: z
|
||||
.object({ left: Pixels, top: Pixels, width: Pixels.min(1), height: Pixels.min(1) })
|
||||
.optional()
|
||||
.describe("Area to keep, in source pixels."),
|
||||
rotate: z.literal([90, 180, 270]).optional().describe("Clockwise, in degrees."),
|
||||
flip: z.enum(["horizontal", "vertical"]).optional(),
|
||||
resize: z
|
||||
.object({
|
||||
aspect: aspect("Resizes to the format's exact size.").optional(),
|
||||
width: z.number().int().min(1).max(8192).optional(),
|
||||
height: z.number().int().min(1).max(8192).optional(),
|
||||
fit: MediaFitSchema.default("crop").describe(
|
||||
"How a different aspect ratio fills the size: crop, pad, or blur (blurred background).",
|
||||
),
|
||||
})
|
||||
.refine((resize) => resize.aspect || resize.width || resize.height, {
|
||||
message: "Set aspect, width or height",
|
||||
})
|
||||
.optional()
|
||||
.describe("Width or height alone keeps the aspect ratio."),
|
||||
extend: z
|
||||
.object({
|
||||
top: Pixels.default(0),
|
||||
right: Pixels.default(0),
|
||||
bottom: Pixels.default(0),
|
||||
left: Pixels.default(0),
|
||||
color: z
|
||||
.string()
|
||||
.regex(/^#[0-9a-fA-F]{6}$/)
|
||||
.optional()
|
||||
.describe("#RRGGBB; transparent when omitted and the format allows it."),
|
||||
})
|
||||
.optional()
|
||||
.describe("Borders to add, in pixels."),
|
||||
prompt: Prompt.optional().describe("Repaint the image with this prompt."),
|
||||
strength: Strength.describe("How far the repaint moves away from the image (default 0.6)."),
|
||||
format: ImageFormatSchema.optional().describe("Default keeps the source format."),
|
||||
})
|
||||
.refine(
|
||||
(input) =>
|
||||
input.crop ||
|
||||
input.rotate ||
|
||||
input.flip ||
|
||||
input.resize ||
|
||||
input.extend ||
|
||||
input.prompt ||
|
||||
input.format,
|
||||
{ message: "Nothing to edit" },
|
||||
);
|
||||
export type EditImageInput = z.infer<typeof EditImageInputSchema>;
|
||||
36
packages/contracts/src/media-jobs.ts
Normal file
36
packages/contracts/src/media-jobs.ts
Normal file
@ -0,0 +1,36 @@
|
||||
import * as z from "zod";
|
||||
import { Id } from "./ids.js";
|
||||
import { MediaJobStatusSchema, MediaOperationSchema, MediaOutputSchema } from "./media.js";
|
||||
|
||||
/** Operations a user starts from a media card (export, transcribe, edit); generation goes through a bot. */
|
||||
export const MediaCardOperationSchema = z.enum(["edit_image", "export_video", "transcribe"]);
|
||||
export type MediaCardOperation = z.infer<typeof MediaCardOperationSchema>;
|
||||
|
||||
export const MediaJobStartInputSchema = z
|
||||
.object({
|
||||
/** The conversation, like other thread procedures: a bot chat or a group. */
|
||||
botId: Id.optional(),
|
||||
groupId: Id.optional(),
|
||||
/** The message the action came from; its bot posts the new card. */
|
||||
messageId: Id,
|
||||
operation: MediaCardOperationSchema,
|
||||
/** The operation's input; validated by the operation. */
|
||||
input: z.record(z.string(), z.unknown()),
|
||||
})
|
||||
.refine((input) => Boolean(input.botId) !== Boolean(input.groupId), {
|
||||
message: "Pass botId or groupId",
|
||||
});
|
||||
export type MediaJobStartInput = z.infer<typeof MediaJobStartInputSchema>;
|
||||
|
||||
export const MediaJobSnapshotSchema = z.object({
|
||||
id: Id,
|
||||
operation: MediaOperationSchema,
|
||||
status: MediaJobStatusSchema,
|
||||
title: z.string(),
|
||||
progress: z.number().min(0).max(1).nullable(),
|
||||
error: z.string().nullable(),
|
||||
outputs: z.array(MediaOutputSchema),
|
||||
});
|
||||
export type MediaJobSnapshotDto = z.infer<typeof MediaJobSnapshotSchema>;
|
||||
|
||||
export const MediaJobStartOutputSchema = z.object({ jobId: Id, messageId: Id });
|
||||
@ -82,6 +82,12 @@ import {
|
||||
IntegrationProviderConfigSchema,
|
||||
IntegrationSetupStateSchema,
|
||||
} from "./integration-settings.js";
|
||||
import {
|
||||
MediaJobSnapshotSchema,
|
||||
MediaJobStartInputSchema,
|
||||
MediaJobStartOutputSchema,
|
||||
} from "./media-jobs.js";
|
||||
import { MediaSettingsUpdateSchema, MediaSettingsViewSchema } from "./media-settings.js";
|
||||
import { ModelTierMappingSchema, SetModelTierInput } from "./model-tiers.js";
|
||||
import {
|
||||
CreateProjectGoalInput,
|
||||
@ -690,6 +696,18 @@ export const appContract = {
|
||||
get: oc.output(IntegrationSetupStateSchema),
|
||||
save: oc.input(IntegrationProviderConfigSchema).output(z.object({ ok: z.literal(true) })),
|
||||
},
|
||||
media: {
|
||||
/** Deployment-wide media backends; deployment owner only. */
|
||||
settings: {
|
||||
get: oc.output(MediaSettingsViewSchema),
|
||||
update: oc.input(MediaSettingsUpdateSchema).output(MediaSettingsViewSchema),
|
||||
},
|
||||
jobs: {
|
||||
start: oc.input(MediaJobStartInputSchema).output(MediaJobStartOutputSchema),
|
||||
get: oc.input(z.object({ jobId: Id })).output(MediaJobSnapshotSchema),
|
||||
cancel: oc.input(z.object({ jobId: Id })).output(MediaJobSnapshotSchema),
|
||||
},
|
||||
},
|
||||
connections: {
|
||||
catalog: oc
|
||||
.input(z.object({ query: z.string().optional(), connectorId: z.string().optional() }))
|
||||
|
||||
@ -34,6 +34,11 @@
|
||||
"development": "./src/plot/plot-spec.ts",
|
||||
"default": "./src/plot/plot-spec.ts"
|
||||
},
|
||||
"./media/workflows": {
|
||||
"types": "./src/media/workflow-template.ts",
|
||||
"development": "./src/media/workflow-template.ts",
|
||||
"default": "./src/media/workflow-template.ts"
|
||||
},
|
||||
"./node/desktop-runtime": {
|
||||
"types": "./src/node/desktop-runtime.ts",
|
||||
"development": "./src/node/desktop-runtime.ts",
|
||||
|
||||
@ -1,12 +1,10 @@
|
||||
import type { AttachmentKind, AttachmentMimeType, MessageBlock } from "@rakazo/contracts";
|
||||
import {
|
||||
ATTACHMENT_MAX_BASE64_LENGTH,
|
||||
ATTACHMENT_MAX_BYTES,
|
||||
type AttachmentKind,
|
||||
type AttachmentMimeType,
|
||||
attachmentKindOf,
|
||||
isAllowedAttachmentMimeType,
|
||||
isAttachmentImageMimeType,
|
||||
type MessageBlock,
|
||||
} from "@rakazo/contracts";
|
||||
|
||||
export class AttachmentValidationError extends Error {
|
||||
|
||||
113
packages/core/src/media/workflow-template.test.ts
Normal file
113
packages/core/src/media/workflow-template.test.ts
Normal file
@ -0,0 +1,113 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
DEFAULT_WORKFLOWS,
|
||||
fillWorkflow,
|
||||
parseWorkflowTemplate,
|
||||
WorkflowTemplateError,
|
||||
} from "./workflow-template.js";
|
||||
|
||||
const minimal = {
|
||||
"1": { class_type: "CLIPTextEncode", inputs: { text: "{{prompt}}", clip: ["2", 1] } },
|
||||
"2": { class_type: "EmptyLatentImage", inputs: { width: "{{width}}", height: 512 } },
|
||||
"3": { class_type: "SaveImage", inputs: { filename_prefix: "run/{{seed}}-{{prompt}}" } },
|
||||
};
|
||||
|
||||
describe("parseWorkflowTemplate", () => {
|
||||
it("accepts JSON text or parsed graphs and lists the placeholders", () => {
|
||||
for (const source of [minimal, JSON.stringify(minimal)]) {
|
||||
const template = parseWorkflowTemplate(source);
|
||||
expect([...template.placeholders].sort()).toEqual(["prompt", "seed", "width"]);
|
||||
}
|
||||
});
|
||||
|
||||
it.each([
|
||||
["not json", "Workflow is not valid JSON"],
|
||||
["[]", "Workflow must be a JSON object"],
|
||||
["{}", "Workflow must have 1 to 1000 nodes"],
|
||||
[JSON.stringify({ nodes: [], links: [] }), "Workflow must be exported in ComfyUI's API format"],
|
||||
[JSON.stringify({ "1": { inputs: {} } }), "Workflow node 1 needs a class_type and inputs"],
|
||||
[
|
||||
JSON.stringify({ "1": { class_type: "SaveImage", inputs: { text: "plain" } } }),
|
||||
"Workflow needs a {{prompt}} placeholder",
|
||||
],
|
||||
[
|
||||
JSON.stringify({
|
||||
"1": { class_type: "X", inputs: { a: "{{prompt}}", b: "{{cfg}}" } },
|
||||
}),
|
||||
"Unknown workflow placeholders: {{cfg}}",
|
||||
],
|
||||
])("rejects %s", (source, message) => {
|
||||
expect(() => parseWorkflowTemplate(source)).toThrow(new WorkflowTemplateError(message));
|
||||
});
|
||||
|
||||
it("rejects placeholder names that only exist on Object.prototype", () => {
|
||||
const source = {
|
||||
"1": {
|
||||
class_type: "X",
|
||||
inputs: { a: "{{prompt}}", b: "{{constructor}}", c: "{{__proto__}}" },
|
||||
},
|
||||
};
|
||||
expect(() => parseWorkflowTemplate(source)).toThrow(
|
||||
"Unknown workflow placeholders: {{constructor}}, {{__proto__}}",
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects oversized workflow text", () => {
|
||||
expect(() => parseWorkflowTemplate(" ".repeat(1_000_001))).toThrow("Workflow is too large");
|
||||
});
|
||||
});
|
||||
|
||||
describe("fillWorkflow", () => {
|
||||
it("replaces whole values with typed values and embedded ones with text, without mutating", () => {
|
||||
const template = parseWorkflowTemplate(minimal);
|
||||
const graph = fillWorkflow(template, { prompt: "a $& cat", width: 768, seed: 7 });
|
||||
expect(graph["1"]?.inputs).toEqual({ text: "a $& cat", clip: ["2", 1] });
|
||||
expect(graph["2"]?.inputs).toEqual({ width: 768, height: 512 });
|
||||
expect(graph["3"]?.inputs.filename_prefix).toBe("run/7-a $& cat");
|
||||
expect(minimal["2"].inputs.width).toBe("{{width}}");
|
||||
});
|
||||
|
||||
it("refuses missing or mistyped values", () => {
|
||||
const template = parseWorkflowTemplate(minimal);
|
||||
expect(() => fillWorkflow(template, { prompt: "cat", seed: 1 })).toThrow(
|
||||
"Workflow value for {{width}} is missing",
|
||||
);
|
||||
expect(() => fillWorkflow(template, { prompt: "cat", width: 1.5, seed: 1 })).toThrow(
|
||||
"{{width}}",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("DEFAULT_WORKFLOWS", () => {
|
||||
it("only image-to-* defaults need an input image", () => {
|
||||
expect(DEFAULT_WORKFLOWS.textToImage.placeholders.has("image")).toBe(false);
|
||||
expect(DEFAULT_WORKFLOWS.textToVideo.placeholders.has("image")).toBe(false);
|
||||
expect(DEFAULT_WORKFLOWS.imageToImage.placeholders.has("image")).toBe(true);
|
||||
expect(DEFAULT_WORKFLOWS.imageToVideo.placeholders.has("image")).toBe(true);
|
||||
});
|
||||
|
||||
it("fill into graphs whose links point at existing nodes", () => {
|
||||
const values = {
|
||||
prompt: "a lighthouse",
|
||||
negative_prompt: "",
|
||||
model: "model.safetensors",
|
||||
image: "input.png",
|
||||
width: 512,
|
||||
height: 512,
|
||||
seed: 1,
|
||||
steps: 20,
|
||||
strength: 0.6,
|
||||
frames: 25,
|
||||
fps: 24,
|
||||
};
|
||||
for (const template of Object.values(DEFAULT_WORKFLOWS)) {
|
||||
const graph = fillWorkflow(template, values);
|
||||
expect(JSON.stringify(graph)).not.toContain("{{");
|
||||
for (const node of Object.values(graph)) {
|
||||
for (const input of Object.values(node.inputs)) {
|
||||
if (Array.isArray(input)) expect(graph[String(input[0])]).toBeDefined();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
159
packages/core/src/media/workflow-template.ts
Normal file
159
packages/core/src/media/workflow-template.ts
Normal file
@ -0,0 +1,159 @@
|
||||
import imageToImage from "./workflows/image-to-image.json" with { type: "json" };
|
||||
import imageToVideo from "./workflows/image-to-video.json" with { type: "json" };
|
||||
import textToImage from "./workflows/text-to-image.json" with { type: "json" };
|
||||
import textToVideo from "./workflows/text-to-video.json" with { type: "json" };
|
||||
|
||||
/**
|
||||
* Values a ComfyUI workflow template can ask for. A string value that is exactly `{{name}}`
|
||||
* becomes the typed value (so `"{{width}}"` turns into the number 1024); `{{name}}` inside a
|
||||
* longer string is replaced by its text. Templates are data: nothing is ever evaluated.
|
||||
*/
|
||||
export const WORKFLOW_PLACEHOLDERS = {
|
||||
prompt: "string",
|
||||
negative_prompt: "string",
|
||||
/** Checkpoint file name: the configured image model or the first installed checkpoint. */
|
||||
model: "string",
|
||||
/** Uploaded input image (image-to-image, image-to-video). */
|
||||
image: "string",
|
||||
width: "integer",
|
||||
height: "integer",
|
||||
seed: "integer",
|
||||
steps: "integer",
|
||||
/** 0–1 denoise for image-to-image. */
|
||||
strength: "number",
|
||||
frames: "integer",
|
||||
fps: "number",
|
||||
} as const satisfies Record<string, "string" | "integer" | "number">;
|
||||
|
||||
export type WorkflowPlaceholder = keyof typeof WORKFLOW_PLACEHOLDERS;
|
||||
export type WorkflowValues = {
|
||||
[Name in WorkflowPlaceholder]?: (typeof WORKFLOW_PLACEHOLDERS)[Name] extends "string"
|
||||
? string
|
||||
: number;
|
||||
};
|
||||
|
||||
/** One node of a workflow in ComfyUI's API format ("Export (API)" in ComfyUI). */
|
||||
export interface ComfyWorkflowNode {
|
||||
class_type: string;
|
||||
inputs: Record<string, unknown>;
|
||||
}
|
||||
export type ComfyWorkflow = Record<string, ComfyWorkflowNode>;
|
||||
|
||||
export interface WorkflowTemplate {
|
||||
readonly graph: ComfyWorkflow;
|
||||
/** The placeholders the graph uses; `image` marks a template that needs an input image. */
|
||||
readonly placeholders: ReadonlySet<WorkflowPlaceholder>;
|
||||
}
|
||||
|
||||
export class WorkflowTemplateError extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = "WorkflowTemplateError";
|
||||
}
|
||||
}
|
||||
|
||||
const PLACEHOLDER_PATTERN = /\{\{([a-z_]+)\}\}/g;
|
||||
const WHOLE_PLACEHOLDER = /^\{\{([a-z_]+)\}\}$/;
|
||||
const MAX_WORKFLOW_CHARS = 1_000_000;
|
||||
const MAX_WORKFLOW_NODES = 1_000;
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return Boolean(value) && typeof value === "object" && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function isPlaceholder(name: string): name is WorkflowPlaceholder {
|
||||
return Object.hasOwn(WORKFLOW_PLACEHOLDERS, name);
|
||||
}
|
||||
|
||||
function visitStrings(value: unknown, visit: (text: string) => void): void {
|
||||
if (typeof value === "string") visit(value);
|
||||
else if (Array.isArray(value)) for (const entry of value) visitStrings(entry, visit);
|
||||
else if (isRecord(value)) for (const entry of Object.values(value)) visitStrings(entry, visit);
|
||||
}
|
||||
|
||||
/** Validates an API-format workflow (JSON text or parsed) and lists its placeholders. */
|
||||
export function parseWorkflowTemplate(source: unknown): WorkflowTemplate {
|
||||
let value = source;
|
||||
if (typeof source === "string") {
|
||||
if (source.length > MAX_WORKFLOW_CHARS) {
|
||||
throw new WorkflowTemplateError("Workflow is too large");
|
||||
}
|
||||
try {
|
||||
value = JSON.parse(source);
|
||||
} catch {
|
||||
throw new WorkflowTemplateError("Workflow is not valid JSON");
|
||||
}
|
||||
}
|
||||
if (!isRecord(value)) throw new WorkflowTemplateError("Workflow must be a JSON object");
|
||||
if (Array.isArray(value.nodes)) {
|
||||
throw new WorkflowTemplateError("Workflow must be exported in ComfyUI's API format");
|
||||
}
|
||||
const nodes = Object.entries(value);
|
||||
if (nodes.length === 0 || nodes.length > MAX_WORKFLOW_NODES) {
|
||||
throw new WorkflowTemplateError(`Workflow must have 1 to ${MAX_WORKFLOW_NODES} nodes`);
|
||||
}
|
||||
for (const [id, node] of nodes) {
|
||||
if (!isRecord(node) || typeof node.class_type !== "string" || !isRecord(node.inputs)) {
|
||||
throw new WorkflowTemplateError(`Workflow node ${id} needs a class_type and inputs`);
|
||||
}
|
||||
}
|
||||
const placeholders = new Set<WorkflowPlaceholder>();
|
||||
const unknown = new Set<string>();
|
||||
visitStrings(value, (text) => {
|
||||
for (const [, name = ""] of text.matchAll(PLACEHOLDER_PATTERN)) {
|
||||
if (isPlaceholder(name)) placeholders.add(name);
|
||||
else unknown.add(`{{${name}}}`);
|
||||
}
|
||||
});
|
||||
if (unknown.size > 0) {
|
||||
throw new WorkflowTemplateError(`Unknown workflow placeholders: ${[...unknown].join(", ")}`);
|
||||
}
|
||||
if (!placeholders.has("prompt")) {
|
||||
throw new WorkflowTemplateError("Workflow needs a {{prompt}} placeholder");
|
||||
}
|
||||
return { graph: value as unknown as ComfyWorkflow, placeholders };
|
||||
}
|
||||
|
||||
function replacePlaceholders(value: unknown, values: WorkflowValues): unknown {
|
||||
if (typeof value === "string") {
|
||||
const whole = WHOLE_PLACEHOLDER.exec(value)?.[1];
|
||||
if (whole) return values[whole as WorkflowPlaceholder];
|
||||
return value.replace(PLACEHOLDER_PATTERN, (_, name: string) =>
|
||||
String(values[name as WorkflowPlaceholder]),
|
||||
);
|
||||
}
|
||||
if (Array.isArray(value)) return value.map((entry) => replacePlaceholders(entry, values));
|
||||
if (isRecord(value)) {
|
||||
return Object.fromEntries(
|
||||
Object.entries(value).map(([key, entry]) => [key, replacePlaceholders(entry, values)]),
|
||||
);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
/** Returns a copy of the template's graph with every placeholder replaced by its typed value. */
|
||||
export function fillWorkflow(template: WorkflowTemplate, values: WorkflowValues): ComfyWorkflow {
|
||||
for (const name of template.placeholders) {
|
||||
const value = values[name];
|
||||
const type = WORKFLOW_PLACEHOLDERS[name];
|
||||
const valid =
|
||||
type === "string"
|
||||
? typeof value === "string"
|
||||
: typeof value === "number" &&
|
||||
Number.isFinite(value) &&
|
||||
(type === "number" || Number.isInteger(value));
|
||||
if (!valid) throw new WorkflowTemplateError(`Workflow value for {{${name}}} is missing`);
|
||||
}
|
||||
return replacePlaceholders(template.graph, values) as ComfyWorkflow;
|
||||
}
|
||||
|
||||
/**
|
||||
* Built-in workflows. Images use core nodes only and run on any SD-family checkpoint; video
|
||||
* uses ComfyUI's core LTX-Video nodes (model files in workflows/README.md).
|
||||
*/
|
||||
export const DEFAULT_WORKFLOWS = {
|
||||
textToImage: parseWorkflowTemplate(textToImage),
|
||||
imageToImage: parseWorkflowTemplate(imageToImage),
|
||||
textToVideo: parseWorkflowTemplate(textToVideo),
|
||||
imageToVideo: parseWorkflowTemplate(imageToVideo),
|
||||
} as const;
|
||||
35
packages/core/src/media/workflows/README.md
Normal file
35
packages/core/src/media/workflows/README.md
Normal file
@ -0,0 +1,35 @@
|
||||
# ComfyUI workflows
|
||||
|
||||
Default workflows for `generate_image`, `generate_video` and repainting in `edit_image`, in
|
||||
ComfyUI's API format. A custom workflow in the Media settings replaces the matching default:
|
||||
one that uses `{{image}}` replaces the image-to-image (or image-to-video) default, one without
|
||||
it replaces the text-to-image (or text-to-video) default.
|
||||
|
||||
## Placeholders
|
||||
|
||||
A string value that is exactly `{{name}}` becomes the typed value; `{{name}}` inside a longer
|
||||
string is replaced by its text. Every workflow needs `{{prompt}}`; other names are optional.
|
||||
|
||||
| Placeholder | Value |
|
||||
|---|---|
|
||||
| `{{prompt}}`, `{{negative_prompt}}` | text |
|
||||
| `{{model}}` | the configured image model, or the first installed checkpoint |
|
||||
| `{{image}}` | the uploaded input image (for `LoadImage`) |
|
||||
| `{{width}}`, `{{height}}` | output size in pixels |
|
||||
| `{{seed}}`, `{{steps}}` | integers |
|
||||
| `{{strength}}` | 0–1, image-to-image denoise |
|
||||
| `{{frames}}`, `{{fps}}` | video length in frames (8n+1) and frame rate |
|
||||
|
||||
To write your own, build it in ComfyUI, replace the values above with placeholders and export it
|
||||
with **Export (API)**.
|
||||
|
||||
## Model files
|
||||
|
||||
- Images: any SD 1.5 or SDXL checkpoint in `models/checkpoints/`.
|
||||
- Video (LTX-Video, core nodes, about 12 GB of VRAM):
|
||||
- `models/checkpoints/ltx-video-2b-v0.9.5.safetensors` from
|
||||
[Lightricks/LTX-Video](https://huggingface.co/Lightricks/LTX-Video)
|
||||
- `models/text_encoders/t5xxl_fp16.safetensors` from
|
||||
[comfyanonymous/flux_text_encoders](https://huggingface.co/comfyanonymous/flux_text_encoders)
|
||||
|
||||
The video workflows need a ComfyUI release with the `CreateVideo` and `SaveVideo` nodes.
|
||||
45
packages/core/src/media/workflows/image-to-image.json
Normal file
45
packages/core/src/media/workflows/image-to-image.json
Normal file
@ -0,0 +1,45 @@
|
||||
{
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": { "ckpt_name": "{{model}}" }
|
||||
},
|
||||
"2": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{prompt}}", "clip": ["1", 1] }
|
||||
},
|
||||
"3": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{negative_prompt}}", "clip": ["1", 1] }
|
||||
},
|
||||
"4": {
|
||||
"class_type": "LoadImage",
|
||||
"inputs": { "image": "{{image}}" }
|
||||
},
|
||||
"5": {
|
||||
"class_type": "VAEEncode",
|
||||
"inputs": { "pixels": ["4", 0], "vae": ["1", 2] }
|
||||
},
|
||||
"6": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"positive": ["2", 0],
|
||||
"negative": ["3", 0],
|
||||
"latent_image": ["5", 0],
|
||||
"seed": "{{seed}}",
|
||||
"steps": "{{steps}}",
|
||||
"cfg": 7,
|
||||
"sampler_name": "dpmpp_2m",
|
||||
"scheduler": "karras",
|
||||
"denoise": "{{strength}}"
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": { "samples": ["6", 0], "vae": ["1", 2] }
|
||||
},
|
||||
"8": {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": { "images": ["7", 0], "filename_prefix": "evoo/image" }
|
||||
}
|
||||
}
|
||||
86
packages/core/src/media/workflows/image-to-video.json
Normal file
86
packages/core/src/media/workflows/image-to-video.json
Normal file
@ -0,0 +1,86 @@
|
||||
{
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": { "ckpt_name": "ltx-video-2b-v0.9.5.safetensors" }
|
||||
},
|
||||
"2": {
|
||||
"class_type": "CLIPLoader",
|
||||
"inputs": { "clip_name": "t5xxl_fp16.safetensors", "type": "ltxv" }
|
||||
},
|
||||
"3": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{prompt}}", "clip": ["2", 0] }
|
||||
},
|
||||
"4": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{negative_prompt}}", "clip": ["2", 0] }
|
||||
},
|
||||
"5": {
|
||||
"class_type": "LoadImage",
|
||||
"inputs": { "image": "{{image}}" }
|
||||
},
|
||||
"6": {
|
||||
"class_type": "LTXVImgToVideo",
|
||||
"inputs": {
|
||||
"positive": ["3", 0],
|
||||
"negative": ["4", 0],
|
||||
"vae": ["1", 2],
|
||||
"image": ["5", 0],
|
||||
"width": "{{width}}",
|
||||
"height": "{{height}}",
|
||||
"length": "{{frames}}",
|
||||
"batch_size": 1,
|
||||
"strength": 1
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"class_type": "LTXVConditioning",
|
||||
"inputs": { "positive": ["6", 0], "negative": ["6", 1], "frame_rate": "{{fps}}" }
|
||||
},
|
||||
"8": {
|
||||
"class_type": "LTXVScheduler",
|
||||
"inputs": {
|
||||
"steps": "{{steps}}",
|
||||
"max_shift": 2.05,
|
||||
"base_shift": 0.95,
|
||||
"stretch": true,
|
||||
"terminal": 0.1,
|
||||
"latent": ["6", 2]
|
||||
}
|
||||
},
|
||||
"9": {
|
||||
"class_type": "KSamplerSelect",
|
||||
"inputs": { "sampler_name": "euler" }
|
||||
},
|
||||
"10": {
|
||||
"class_type": "SamplerCustom",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"add_noise": true,
|
||||
"noise_seed": "{{seed}}",
|
||||
"cfg": 3,
|
||||
"positive": ["7", 0],
|
||||
"negative": ["7", 1],
|
||||
"sampler": ["9", 0],
|
||||
"sigmas": ["8", 0],
|
||||
"latent_image": ["6", 2]
|
||||
}
|
||||
},
|
||||
"11": {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": { "samples": ["10", 0], "vae": ["1", 2] }
|
||||
},
|
||||
"12": {
|
||||
"class_type": "CreateVideo",
|
||||
"inputs": { "images": ["11", 0], "fps": "{{fps}}" }
|
||||
},
|
||||
"13": {
|
||||
"class_type": "SaveVideo",
|
||||
"inputs": {
|
||||
"video": ["12", 0],
|
||||
"filename_prefix": "evoo/video",
|
||||
"format": "mp4",
|
||||
"codec": "h264"
|
||||
}
|
||||
}
|
||||
}
|
||||
41
packages/core/src/media/workflows/text-to-image.json
Normal file
41
packages/core/src/media/workflows/text-to-image.json
Normal file
@ -0,0 +1,41 @@
|
||||
{
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": { "ckpt_name": "{{model}}" }
|
||||
},
|
||||
"2": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{prompt}}", "clip": ["1", 1] }
|
||||
},
|
||||
"3": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{negative_prompt}}", "clip": ["1", 1] }
|
||||
},
|
||||
"4": {
|
||||
"class_type": "EmptyLatentImage",
|
||||
"inputs": { "width": "{{width}}", "height": "{{height}}", "batch_size": 1 }
|
||||
},
|
||||
"5": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"positive": ["2", 0],
|
||||
"negative": ["3", 0],
|
||||
"latent_image": ["4", 0],
|
||||
"seed": "{{seed}}",
|
||||
"steps": "{{steps}}",
|
||||
"cfg": 7,
|
||||
"sampler_name": "dpmpp_2m",
|
||||
"scheduler": "karras",
|
||||
"denoise": 1
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": { "samples": ["5", 0], "vae": ["1", 2] }
|
||||
},
|
||||
"7": {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": { "images": ["6", 0], "filename_prefix": "evoo/image" }
|
||||
}
|
||||
}
|
||||
77
packages/core/src/media/workflows/text-to-video.json
Normal file
77
packages/core/src/media/workflows/text-to-video.json
Normal file
@ -0,0 +1,77 @@
|
||||
{
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": { "ckpt_name": "ltx-video-2b-v0.9.5.safetensors" }
|
||||
},
|
||||
"2": {
|
||||
"class_type": "CLIPLoader",
|
||||
"inputs": { "clip_name": "t5xxl_fp16.safetensors", "type": "ltxv" }
|
||||
},
|
||||
"3": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{prompt}}", "clip": ["2", 0] }
|
||||
},
|
||||
"4": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": { "text": "{{negative_prompt}}", "clip": ["2", 0] }
|
||||
},
|
||||
"5": {
|
||||
"class_type": "LTXVConditioning",
|
||||
"inputs": { "positive": ["3", 0], "negative": ["4", 0], "frame_rate": "{{fps}}" }
|
||||
},
|
||||
"6": {
|
||||
"class_type": "EmptyLTXVLatentVideo",
|
||||
"inputs": {
|
||||
"width": "{{width}}",
|
||||
"height": "{{height}}",
|
||||
"length": "{{frames}}",
|
||||
"batch_size": 1
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"class_type": "LTXVScheduler",
|
||||
"inputs": {
|
||||
"steps": "{{steps}}",
|
||||
"max_shift": 2.05,
|
||||
"base_shift": 0.95,
|
||||
"stretch": true,
|
||||
"terminal": 0.1,
|
||||
"latent": ["6", 0]
|
||||
}
|
||||
},
|
||||
"8": {
|
||||
"class_type": "KSamplerSelect",
|
||||
"inputs": { "sampler_name": "euler" }
|
||||
},
|
||||
"9": {
|
||||
"class_type": "SamplerCustom",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"add_noise": true,
|
||||
"noise_seed": "{{seed}}",
|
||||
"cfg": 3,
|
||||
"positive": ["5", 0],
|
||||
"negative": ["5", 1],
|
||||
"sampler": ["8", 0],
|
||||
"sigmas": ["7", 0],
|
||||
"latent_image": ["6", 0]
|
||||
}
|
||||
},
|
||||
"10": {
|
||||
"class_type": "VAEDecode",
|
||||
"inputs": { "samples": ["9", 0], "vae": ["1", 2] }
|
||||
},
|
||||
"11": {
|
||||
"class_type": "CreateVideo",
|
||||
"inputs": { "images": ["10", 0], "fps": "{{fps}}" }
|
||||
},
|
||||
"12": {
|
||||
"class_type": "SaveVideo",
|
||||
"inputs": {
|
||||
"video": ["11", 0],
|
||||
"filename_prefix": "evoo/video",
|
||||
"format": "mp4",
|
||||
"codec": "h264"
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -106,6 +106,7 @@ async function main() {
|
||||
"packages/testkit/src/connections.test.ts",
|
||||
"packages/testkit/src/bot-secrets.test.ts",
|
||||
"packages/testkit/src/delegation.postgres.test.ts",
|
||||
"packages/testkit/src/media.postgres.test.ts",
|
||||
"packages/testkit/src/user-profile.postgres.test.ts",
|
||||
"packages/db/src/space-membership.postgres.test.ts",
|
||||
"packages/db/src/main-agent.postgres.test.ts",
|
||||
@ -117,7 +118,9 @@ async function main() {
|
||||
"packages/adapters/src/realtime.postgres.test.ts",
|
||||
"packages/adapters/src/job-reconciler.postgres.test.ts",
|
||||
"packages/adapters/src/cloud-agent.postgres.test.ts",
|
||||
"packages/adapters/src/media/image-operations.postgres.test.ts",
|
||||
"packages/adapters/src/media/transcription.postgres.test.ts",
|
||||
"packages/adapters/src/media-runtime.postgres.test.ts",
|
||||
"packages/adapters/src/builtin-memory-provider.postgres.test.ts",
|
||||
"packages/adapters/src/profile-consolidation.postgres.test.ts",
|
||||
];
|
||||
|
||||
378
packages/testkit/src/media.postgres.test.ts
Normal file
378
packages/testkit/src/media.postgres.test.ts
Normal file
@ -0,0 +1,378 @@
|
||||
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 { createThreadMessage } from "@rakazo/db";
|
||||
import { afterAll, beforeAll, describe, expect, it } from "vitest";
|
||||
import type { createApp } from "../../../apps/api/src/app.ts";
|
||||
import { sessionCookieHeader } from "./index.js";
|
||||
|
||||
type App = { request: (input: string, init?: RequestInit) => Response | Promise<Response> };
|
||||
type AppHandles = Awaited<ReturnType<typeof createApp>>;
|
||||
|
||||
process.env.WAKEUP_DRIVER = "memory";
|
||||
process.env.SANDBOX_PROVIDER = "fake";
|
||||
process.env.AGENT_RUNTIME = "scripted";
|
||||
|
||||
const hasDb = process.env.VERIFY_DATABASE === "1" && Boolean(process.env.DATABASE_URL);
|
||||
const describeWithDatabase = hasDb ? describe.sequential : describe.skip;
|
||||
|
||||
const PNG_BASE64 = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3]).toString(
|
||||
"base64",
|
||||
);
|
||||
|
||||
// testkit has no zod of its own; the contracts' schemas carry it.
|
||||
const SourceInputSchema = MediaOutputSchema.pick({ artifactId: true });
|
||||
|
||||
/** Fake card operations: an export that copies its source, and a transcription that waits. */
|
||||
function fakeOperations(): MediaOperationRegistration[] {
|
||||
const exportVideo: MediaOperationRegistration<{ artifactId: string }> = {
|
||||
operation: "export_video",
|
||||
inputSchema: SourceInputSchema,
|
||||
title: () => "Export",
|
||||
available: async () => true,
|
||||
artifactIds: (input) => [input.artifactId],
|
||||
handler: {
|
||||
operation: "export_video",
|
||||
run: async (context) => {
|
||||
const { artifactId } = context.input as { artifactId: string };
|
||||
const source = await context.fetchArtifact(artifactId);
|
||||
const out = path.join(context.tempDir, "export.png");
|
||||
await writeFile(out, Buffer.from(PNG_BASE64, "base64"));
|
||||
await context.addOutput({
|
||||
path: out,
|
||||
name: "export.png",
|
||||
mimeType: "image/png",
|
||||
parentArtifactId: source.artifactId,
|
||||
preset: "square",
|
||||
});
|
||||
return { summary: "Exported." };
|
||||
},
|
||||
},
|
||||
};
|
||||
const transcribe: MediaOperationRegistration<{ artifactId: string }> = {
|
||||
operation: "transcribe",
|
||||
inputSchema: SourceInputSchema,
|
||||
title: () => "Transcript",
|
||||
available: async () => true,
|
||||
artifactIds: (input) => [input.artifactId],
|
||||
handler: {
|
||||
operation: "transcribe",
|
||||
run: async (context) => {
|
||||
await new Promise<void>((resolve) => {
|
||||
context.signal.addEventListener("abort", () => resolve(), { once: true });
|
||||
});
|
||||
context.signal.throwIfAborted();
|
||||
return { summary: "" };
|
||||
},
|
||||
},
|
||||
};
|
||||
return [exportVideo, transcribe] as MediaOperationRegistration[];
|
||||
}
|
||||
|
||||
describeWithDatabase("media RPCs (PostgreSQL)", () => {
|
||||
let handles: AppHandles;
|
||||
let app: App;
|
||||
const stamp = `${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
const dataDir = mkdtempSync(path.join(tmpdir(), "rakazo-media-"));
|
||||
|
||||
beforeAll(async () => {
|
||||
const { createApp } = await import("../../../apps/api/src/app.ts");
|
||||
handles = await createApp({
|
||||
databaseUrl: process.env.DATABASE_URL!,
|
||||
dataDir,
|
||||
sandboxProvider: "fake",
|
||||
agentRuntime: "scripted",
|
||||
wakeupDriver: "memory",
|
||||
signupsEnabled: "true",
|
||||
composio: new ComposioEmulator(),
|
||||
mediaOperations: fakeOperations(),
|
||||
});
|
||||
app = handles.app;
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await handles?.prisma.integrationProviderConfig.deleteMany({ where: { id: "media" } });
|
||||
await handles?.stop();
|
||||
rmSync(dataDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
async function user(label: string) {
|
||||
const cookie = await signup(app, `${label}-${stamp}@rakazo.test`, label);
|
||||
const me = await rpc<{ userId: string; spaceId: string }>(app, cookie, "me");
|
||||
const bot = await rpc<{ id: string }>(app, cookie, "bots/create", {
|
||||
name: `${label} bot`,
|
||||
title: "",
|
||||
description: "",
|
||||
instructions: "",
|
||||
notifyOnFinish: false,
|
||||
});
|
||||
const thread = await handles.prisma.thread.findUniqueOrThrow({ where: { botId: bot.id } });
|
||||
const card = await createThreadMessage(handles.prisma, {
|
||||
threadId: thread.id,
|
||||
role: "bot",
|
||||
botId: bot.id,
|
||||
blocks: [{ kind: "text", text: "Here is your clip." }],
|
||||
});
|
||||
const artifact = await rpc<{ id: string }>(app, cookie, "artifacts/create", {
|
||||
botId: bot.id,
|
||||
name: "clip.png",
|
||||
mimeType: "image/png",
|
||||
contentBase64: PNG_BASE64,
|
||||
});
|
||||
return { cookie, me, bot, thread, card, artifact };
|
||||
}
|
||||
|
||||
async function waitForStatus(jobId: string, status: string) {
|
||||
const start = Date.now();
|
||||
while (Date.now() - start < 20_000) {
|
||||
const job = await handles.prisma.mediaJob.findUnique({ where: { id: jobId } });
|
||||
if (job?.status === status) return job;
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
}
|
||||
throw new Error(`media job ${jobId} never reached ${status}`);
|
||||
}
|
||||
|
||||
it("lets only the deployment owner read and change media settings, and never returns keys", async () => {
|
||||
const owner = await user("media-owner");
|
||||
const other = await user("media-other");
|
||||
await handles.prisma.deploymentSettings.upsert({
|
||||
where: { id: "default" },
|
||||
create: { id: "default", ownerUserId: owner.me.userId },
|
||||
update: { ownerUserId: owner.me.userId },
|
||||
});
|
||||
|
||||
const update = {
|
||||
comfyui: {
|
||||
baseUrl: "http://127.0.0.1:8188/",
|
||||
apiKey: "comfy-test-key",
|
||||
imageModel: "sd_xl_base_1.0.safetensors",
|
||||
imageWorkflow: '{"3":{"class_type":"KSampler","inputs":{}}}',
|
||||
videoWorkflow: "",
|
||||
},
|
||||
transcription: { baseUrl: "http://127.0.0.1:8080", model: "whisper-1", apiKey: "" },
|
||||
};
|
||||
expect((await raw(app, other.cookie, "media/settings/get", {})).status).toBe(403);
|
||||
expect((await raw(app, other.cookie, "media/settings/update", update)).status).toBe(403);
|
||||
expect(await handles.prisma.integrationProviderConfig.count({ where: { id: "media" } })).toBe(
|
||||
0,
|
||||
);
|
||||
|
||||
const saved = await rpc<MediaSettingsView>(app, owner.cookie, "media/settings/update", update);
|
||||
expect(saved).toEqual({
|
||||
comfyui: {
|
||||
baseUrl: "http://127.0.0.1:8188",
|
||||
hasApiKey: true,
|
||||
imageModel: "sd_xl_base_1.0.safetensors",
|
||||
imageWorkflow: update.comfyui.imageWorkflow,
|
||||
},
|
||||
transcription: { baseUrl: "http://127.0.0.1:8080/v1", model: "whisper-1", hasApiKey: false },
|
||||
});
|
||||
expect(JSON.stringify(saved)).not.toContain("comfy-test-key");
|
||||
expect(await rpc(app, owner.cookie, "media/settings/get")).toEqual(saved);
|
||||
const row = await handles.prisma.integrationProviderConfig.findUniqueOrThrow({
|
||||
where: { id: "media" },
|
||||
});
|
||||
expect(row.ciphertext).not.toContain("comfy-test-key");
|
||||
expect(row.ciphertext).not.toContain("127.0.0.1");
|
||||
|
||||
// An omitted key stays with the same server, never follows the settings elsewhere.
|
||||
const { apiKey: _key, ...withoutKey } = update.comfyui;
|
||||
const kept = await rpc<MediaSettingsView>(app, owner.cookie, "media/settings/update", {
|
||||
...update,
|
||||
comfyui: withoutKey,
|
||||
});
|
||||
expect(kept.comfyui?.hasApiKey).toBe(true);
|
||||
const moved = await rpc<MediaSettingsView>(app, owner.cookie, "media/settings/update", {
|
||||
...update,
|
||||
comfyui: { ...withoutKey, baseUrl: "http://127.0.0.1:8189" },
|
||||
});
|
||||
expect(moved.comfyui?.hasApiKey).toBe(false);
|
||||
|
||||
for (const baseUrl of ["http://169.254.169.254", "http://user:pass@127.0.0.1:8188"]) {
|
||||
const refused = await raw(app, owner.cookie, "media/settings/update", {
|
||||
...update,
|
||||
comfyui: { ...update.comfyui, baseUrl },
|
||||
});
|
||||
expect(refused.status, baseUrl).toBe(400);
|
||||
}
|
||||
const badWorkflow = await raw(app, owner.cookie, "media/settings/update", {
|
||||
...update,
|
||||
comfyui: { ...update.comfyui, imageWorkflow: "[1, 2" },
|
||||
});
|
||||
expect(badWorkflow.status).toBe(400);
|
||||
|
||||
const previous = process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC;
|
||||
try {
|
||||
delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC;
|
||||
expect(
|
||||
(
|
||||
await raw(app, owner.cookie, "media/settings/update", {
|
||||
comfyui: null,
|
||||
transcription: { baseUrl: "http://transcribe.example.com" },
|
||||
})
|
||||
).status,
|
||||
).toBe(400);
|
||||
process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = "1";
|
||||
// A key is never sent in cleartext to a public host.
|
||||
expect(
|
||||
(
|
||||
await raw(app, owner.cookie, "media/settings/update", {
|
||||
comfyui: null,
|
||||
transcription: { baseUrl: "http://transcribe.example.com", apiKey: "secret" },
|
||||
})
|
||||
).status,
|
||||
).toBe(400);
|
||||
const https = await rpc<MediaSettingsView>(app, owner.cookie, "media/settings/update", {
|
||||
comfyui: null,
|
||||
transcription: { baseUrl: "https://transcribe.example.com", apiKey: "secret" },
|
||||
});
|
||||
expect(https.transcription).toMatchObject({
|
||||
baseUrl: "https://transcribe.example.com/v1",
|
||||
hasApiKey: true,
|
||||
});
|
||||
} finally {
|
||||
if (previous === undefined) delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC;
|
||||
else process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = previous;
|
||||
}
|
||||
|
||||
expect(
|
||||
await rpc(app, owner.cookie, "media/settings/update", { comfyui: null, transcription: null }),
|
||||
).toEqual({ comfyui: null, transcription: null });
|
||||
expect(await handles.prisma.integrationProviderConfig.count({ where: { id: "media" } })).toBe(
|
||||
0,
|
||||
);
|
||||
});
|
||||
|
||||
it("starts, reads and cancels card jobs only for the actor's own threads and files", async () => {
|
||||
const alice = await user("media-alice");
|
||||
const mallory = await user("media-mallory");
|
||||
|
||||
const started = await rpc<{ jobId: string; messageId: string }>(
|
||||
app,
|
||||
alice.cookie,
|
||||
"media/jobs/start",
|
||||
{
|
||||
botId: alice.bot.id,
|
||||
messageId: alice.card.id,
|
||||
operation: "export_video",
|
||||
input: { artifactId: alice.artifact.id },
|
||||
},
|
||||
);
|
||||
const job = await waitForStatus(started.jobId, "done");
|
||||
expect(job).toMatchObject({
|
||||
spaceId: alice.me.spaceId,
|
||||
userId: alice.me.userId,
|
||||
botId: alice.bot.id,
|
||||
threadId: alice.thread.id,
|
||||
runId: null,
|
||||
continueAfter: false,
|
||||
});
|
||||
const snapshot = await rpc<{ status: string; outputs: MediaJobBlock["outputs"] }>(
|
||||
app,
|
||||
alice.cookie,
|
||||
"media/jobs/get",
|
||||
{ jobId: started.jobId },
|
||||
);
|
||||
expect(snapshot.status).toBe("done");
|
||||
expect(snapshot.outputs).toEqual([
|
||||
expect.objectContaining({
|
||||
name: "export.png",
|
||||
parentArtifactId: alice.artifact.id,
|
||||
preset: "square",
|
||||
}),
|
||||
]);
|
||||
const card = await handles.prisma.message.findUniqueOrThrow({
|
||||
where: { id: started.messageId },
|
||||
});
|
||||
expect(card).toMatchObject({ threadId: alice.thread.id, botId: alice.bot.id, role: "bot" });
|
||||
expect((card.blocks as MessageBlock[])[0]).toMatchObject({ kind: "media", status: "done" });
|
||||
|
||||
// Another user can neither read nor cancel the job, nor start one in Alice's conversation
|
||||
// or on her files.
|
||||
expect(
|
||||
(await raw(app, mallory.cookie, "media/jobs/get", { jobId: started.jobId })).status,
|
||||
).toBe(404);
|
||||
expect(
|
||||
(await raw(app, mallory.cookie, "media/jobs/cancel", { jobId: started.jobId })).status,
|
||||
).toBe(404);
|
||||
for (const body of [
|
||||
{
|
||||
botId: alice.bot.id,
|
||||
messageId: alice.card.id,
|
||||
operation: "export_video",
|
||||
input: { artifactId: mallory.artifact.id },
|
||||
},
|
||||
{
|
||||
botId: mallory.bot.id,
|
||||
messageId: alice.card.id,
|
||||
operation: "export_video",
|
||||
input: { artifactId: mallory.artifact.id },
|
||||
},
|
||||
{
|
||||
botId: mallory.bot.id,
|
||||
messageId: mallory.card.id,
|
||||
operation: "export_video",
|
||||
input: { artifactId: alice.artifact.id },
|
||||
},
|
||||
{
|
||||
botId: mallory.bot.id,
|
||||
messageId: mallory.card.id,
|
||||
operation: "generate_image",
|
||||
input: { prompt: "anything" },
|
||||
},
|
||||
]) {
|
||||
const response = await raw(app, mallory.cookie, "media/jobs/start", body);
|
||||
expect(response.status, JSON.stringify(body)).toBeGreaterThanOrEqual(400);
|
||||
}
|
||||
expect(await handles.prisma.mediaJob.count({ where: { userId: mallory.me.userId } })).toBe(0);
|
||||
|
||||
const slow = await rpc<{ jobId: string }>(app, alice.cookie, "media/jobs/start", {
|
||||
botId: alice.bot.id,
|
||||
messageId: alice.card.id,
|
||||
operation: "transcribe",
|
||||
input: { artifactId: alice.artifact.id },
|
||||
});
|
||||
await waitForStatus(slow.jobId, "running");
|
||||
const cancelled = await rpc<{ status: string }>(app, alice.cookie, "media/jobs/cancel", {
|
||||
jobId: slow.jobId,
|
||||
});
|
||||
expect(["running", "cancelled"]).toContain(cancelled.status);
|
||||
await waitForStatus(slow.jobId, "cancelled");
|
||||
});
|
||||
});
|
||||
|
||||
async function signup(app: App, email: string, name: string) {
|
||||
const res = await app.request("/api/auth/sign-up/email", {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json", origin: "http://127.0.0.1:5173" },
|
||||
body: JSON.stringify({ email, password: "password12", name }),
|
||||
});
|
||||
if (res.status >= 400) throw new Error(`signup failed ${res.status}: ${await res.text()}`);
|
||||
return sessionCookieHeader(res);
|
||||
}
|
||||
|
||||
function raw(app: App, cookie: string, proc: string, body: unknown = {}) {
|
||||
return app.request(`/rpc/${proc}`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json", cookie, origin: "http://127.0.0.1:5173" },
|
||||
body: JSON.stringify({ json: body }),
|
||||
});
|
||||
}
|
||||
|
||||
async function rpc<T>(app: App, cookie: string, proc: string, body: unknown = {}): Promise<T> {
|
||||
const res = await raw(app, cookie, proc, body);
|
||||
const text = await res.text();
|
||||
const parsed = JSON.parse(text) as { json?: T; error?: { message?: string } };
|
||||
if (res.status >= 400 || parsed.error) {
|
||||
throw new Error(`${proc} ${res.status}: ${parsed.error?.message ?? text}`);
|
||||
}
|
||||
return parsed.json as T;
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user