evOo-local/packages/adapters/src/media-tools.test.ts

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("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<Parameters<typeof selectBuiltinToolsForRun>[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("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",
    );
  });
});