import type { MediaCapabilities, MediaRuntime, StartMediaJobInput } from "@rakazo/adapter-kit"; import { STUDIO_BOT_TOOLS } from "@rakazo/db"; import { describe, expect, it, vi } from "vitest"; import { selectBuiltinToolsForRun, unknownBuiltinToolNames } from "./executor-tools.js"; import { MediaJobRejectedError, mediaWakePrompt } 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, analyze_video: 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("exposes a media tool only for operations the space can run, within the bot allowlist", () => { const base = { graphicalToolsAllowed: true, groupId: null, trigger: "user", semanticMemoryEnabled: false, messagingChannelRun: false, bot: { kind: "specialist", restrictBuiltinTools: false, allowedBuiltinTools: [] as string[] }, }; const mediaNames = (options: Partial[0]>) => selectBuiltinToolsForRun({ ...base, ...options }) .map((tool) => tool.name) .filter((name) => isMediaToolName(name)); expect(mediaNames({})).toEqual([]); expect(mediaNames({ mediaCapabilities: none })).toEqual([]); expect(mediaNames({ mediaCapabilities: { ...none, generate_image: true } })).toEqual([ "generate_image", ]); expect( mediaNames({ mediaCapabilities: { ...none, generate_image: true, edit_image: true }, bot: { kind: "specialist", restrictBuiltinTools: true, allowedBuiltinTools: ["edit_image"], }, }), ).toEqual(["edit_image"]); }); it("adds continue_after to every media tool schema", () => { expect( mediaToolInputSchema({ 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("always wakes the bot for tools whose result is only for it", async () => { const start = vi.fn(async (_input: StartMediaJobInput) => ({ jobId: "job-1", messageId: "message-1", })); const result = await executeMediaTool(runtime(start), scope, "analyze_video", { artifact_id: "video-1", title: "Cut points", continue_after: false, }); expect(start.mock.calls[0]?.[0]).toMatchObject({ operation: "analyze_video", input: { artifact_id: "video-1", title: "Cut points" }, continueAfter: true, }); expect((result as { note: string }).note).toContain("woken"); expect( mediaToolInputSchema({ properties: { title: { type: "string" } } }, { continueAfter: false }), ).toEqual({ type: "object", properties: { title: { type: "string" } } }); }); 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", ); }); it("includes both artifactId and materialized workspace path in mediaWakePrompt and keeps STUDIO_BOT_TOOLS valid", () => { expect(unknownBuiltinToolNames(STUDIO_BOT_TOOLS)).toEqual([]); const prompt = mediaWakePrompt({ id: "job-hero", operation: "generate_image", title: "Botanical bottle on travertine", status: "done", error: null, summary: "Generated 1 image.", outputs: [ { artifactId: "art_hero_1", kind: "image", mimeType: "image/png", name: "botanical-bottle.png", size: 4096, }, ], }); expect(prompt).toContain("botanical-bottle.png (art_hero_1, attachments/art_hero_1.png)"); }); });