openclaw/extensions/openrouter/image-generation-provider.test.ts

// Openrouter tests cover image generation provider plugin behavior.
import {
  getProviderHttpMocks,
  installProviderHttpMockCleanup,
} from "openclaw/plugin-sdk/provider-http-test-mocks";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { buildOpenRouterImageGenerationProvider } from "./image-generation-provider.js";

const { postJsonRequestMock, resolveApiKeyForProviderMock, resolveProviderHttpRequestConfigMock } =
  getProviderHttpMocks();

installProviderHttpMockCleanup();

function requireOpenRouterPostBody(): {
  messages?: Array<{ content?: unknown }>;
} {
  const request = requireOpenRouterPostRequest();
  return request.body as { messages?: Array<{ content?: unknown }> };
}

function requireOpenRouterPostRequest(): Record<string, unknown> {
  const [call] = postJsonRequestMock.mock.calls;
  if (!call) {
    throw new Error("expected OpenRouter image generation request");
  }
  const [request] = call;
  if (!request || typeof request !== "object" || Array.isArray(request)) {
    throw new Error("expected OpenRouter image generation request");
  }
  return request as Record<string, unknown>;
}

function requireOpenRouterConfigRequest(): Record<string, unknown> {
  const [call] = resolveProviderHttpRequestConfigMock.mock.calls;
  if (!call) {
    throw new Error("expected OpenRouter image config request");
  }
  const [request] = call;
  if (!request || typeof request !== "object" || Array.isArray(request)) {
    throw new Error("expected OpenRouter image config request");
  }
  return request;
}

function requireHeaders(value: unknown): Headers {
  if (!(value instanceof Headers)) {
    throw new Error("expected OpenRouter image request headers");
  }
  return value;
}

function requireGeneratedImage(
  result: Awaited<
    ReturnType<ReturnType<typeof buildOpenRouterImageGenerationProvider>["generateImage"]>
  >,
  index: number,
) {
  const image = result.images[index];
  if (!image) {
    throw new Error(`expected OpenRouter generated image at index ${index}`);
  }
  return image;
}

describe("openrouter image generation provider", () => {
  beforeEach(() => {
    resolveApiKeyForProviderMock.mockResolvedValue({ apiKey: "openrouter-key" });
  });

  it("builds provider metadata and capabilities", () => {
    const provider = buildOpenRouterImageGenerationProvider();
    expect(provider.id).toBe("openrouter");
    expect(provider.label).toBe("OpenRouter");
    expect(provider.defaultModel).toBe("google/gemini-3.1-flash-image-preview");
    expect(provider.models).toContain("google/gemini-3-pro-image-preview");
    expect(provider.capabilities.generate.maxCount).toBe(4);
    expect(provider.capabilities.generate.supportsAspectRatio).toBe(true);
    expect(provider.capabilities.edit.enabled).toBe(true);
    expect(provider.capabilities.edit.maxInputImages).toBe(5);
  });

  it("sends chat completion image requests with Gemini image config and count", async () => {
    const release = vi.fn(async () => {});
    postJsonRequestMock.mockResolvedValue({
      response: {
        json: async () => ({
          choices: [
            {
              message: {
                images: [
                  {
                    imageUrl: {
                      url: `data:image/png;base64,${Buffer.from("png-one").toString("base64")}`,
                    },
                  },
                ],
              },
            },
          ],
        }),
      },
      release,
    });

    const provider = buildOpenRouterImageGenerationProvider();
    const result = await provider.generateImage({
      provider: "openrouter",
      model: "google/gemini-3.1-flash-image-preview",
      prompt: "draw a sticker",
      aspectRatio: "16:9",
      resolution: "2K",
      count: 2,
      timeoutMs: 12_345,
      ssrfPolicy: { allowRfc2544BenchmarkRange: true },
      cfg: {
        models: {
          providers: {
            openrouter: {
              baseUrl: "https://custom.openrouter.test/api/v1",
              models: [],
            },
          },
        },
      },
    });

    expect(resolveApiKeyForProviderMock).toHaveBeenCalledOnce();
    expect(resolveApiKeyForProviderMock).toHaveBeenCalledWith({
      provider: "openrouter",
      cfg: {
        models: {
          providers: {
            openrouter: {
              baseUrl: "https://custom.openrouter.test/api/v1",
              models: [],
            },
          },
        },
      },
      agentDir: undefined,
      store: undefined,
    });
    expect(resolveProviderHttpRequestConfigMock).toHaveBeenCalledOnce();
    expect(requireOpenRouterConfigRequest()).toEqual({
      baseUrl: "https://custom.openrouter.test/api/v1",
      defaultBaseUrl: "https://openrouter.ai/api/v1",
      allowPrivateNetwork: false,
      defaultHeaders: {
        Authorization: "Bearer openrouter-key",
        "HTTP-Referer": "https://openclaw.ai",
        "X-OpenRouter-Title": "OpenClaw",
      },
      request: undefined,
      provider: "openrouter",
      capability: "image",
      transport: "http",
    });
    expect(postJsonRequestMock).toHaveBeenCalledOnce();
    const request = requireOpenRouterPostRequest();
    const headers = requireHeaders(request.headers);
    expect(Object.fromEntries(headers.entries())).toEqual({
      authorization: "Bearer openrouter-key",
      "http-referer": "https://openclaw.ai",
      "x-openrouter-title": "OpenClaw",
    });
    expect(request).toEqual({
      url: "https://custom.openrouter.test/api/v1/chat/completions",
      headers,
      body: {
        model: "google/gemini-3.1-flash-image-preview",
        messages: [
          {
            role: "user",
            content: "draw a sticker",
          },
        ],
        modalities: ["image", "text"],
        n: 2,
        image_config: {
          aspect_ratio: "16:9",
          image_size: "2K",
        },
      },
      timeoutMs: 12_345,
      fetchFn: fetch,
      allowPrivateNetwork: false,
      ssrfPolicy: { allowRfc2544BenchmarkRange: true },
      dispatcherPolicy: undefined,
    });
    const image = requireGeneratedImage(result, 0);
    expect(image.buffer.toString()).toBe("png-one");
    expect(image.mimeType).toBe("image/png");
    expect(release).toHaveBeenCalledOnce();
  });

  it("applies configured image request transport without weakening private-network policy", async () => {
    const requestPolicy = {
      allowPrivateNetwork: true,
      headers: { "X-OpenRouter-Trace": "image-trace" },
      auth: { mode: "authorization-bearer" as const, token: "override-image-token" },
      proxy: { mode: "explicit-proxy" as const, url: "http://proxy.example.test:8443" },
      tls: { ca: "synthetic-provider-ca", serverName: "provider.example.test" },
    };
    const dispatcherPolicy = {
      mode: "explicit-proxy" as const,
      proxyUrl: requestPolicy.proxy.url,
    };
    resolveProviderHttpRequestConfigMock.mockImplementationOnce((params) => {
      const headers = new Headers(params.defaultHeaders);
      for (const [name, value] of Object.entries(params.request?.headers ?? {})) {
        headers.set(name, value);
      }
      if (params.request?.auth?.mode === "authorization-bearer") {
        headers.set("Authorization", `Bearer ${params.request.auth.token}`);
      }
      return {
        baseUrl: params.baseUrl ?? params.defaultBaseUrl,
        allowPrivateNetwork:
          (params.allowPrivateNetwork ?? params.request?.allowPrivateNetwork) === true,
        headers,
        dispatcherPolicy,
      };
    });
    const release = vi.fn(async () => {});
    postJsonRequestMock.mockResolvedValue({
      response: Response.json({
        choices: [
          {
            message: {
              images: [{ image_url: { url: "data:image/png;base64,cG5n" } }],
            },
          },
        ],
      }),
      release,
    });

    const result = await buildOpenRouterImageGenerationProvider().generateImage({
      provider: "openrouter",
      model: "google/gemini-3.1-flash-image-preview",
      prompt: "draw through the configured transport",
      cfg: {
        models: {
          providers: {
            openrouter: {
              baseUrl: "https://custom.openrouter.test/api/v1",
              request: requestPolicy,
              models: [],
            },
          },
        },
      },
    });

    expect(requireOpenRouterConfigRequest()).toMatchObject({
      baseUrl: "https://custom.openrouter.test/api/v1",
      provider: "openrouter",
      capability: "image",
      allowPrivateNetwork: false,
      request: requestPolicy,
    });
    const request = requireOpenRouterPostRequest();
    const headers = requireHeaders(request.headers);
    expect(headers.get("authorization")).toBe("Bearer override-image-token");
    expect(headers.get("x-openrouter-trace")).toBe("image-trace");
    expect(request).toMatchObject({
      allowPrivateNetwork: false,
      dispatcherPolicy,
    });
    expect(requireGeneratedImage(result, 0).buffer.toString()).toBe("png");
    expect(release).toHaveBeenCalledOnce();
  });

  it("uses a 180s default timeout when no request timeout is provided", async () => {
    const release = vi.fn(async () => {});
    postJsonRequestMock.mockResolvedValue({
      response: {
        json: async () => ({
          choices: [
            {
              message: {
                images: [
                  {
                    imageUrl: {
                      url: `data:image/png;base64,${Buffer.from("png-one").toString("base64")}`,
                    },
                  },
                ],
              },
            },
          ],
        }),
      },
      release,
    });

    const provider = buildOpenRouterImageGenerationProvider();
    await provider.generateImage({
      provider: "openrouter",
      model: "google/gemini-3.1-flash-image-preview",
      prompt: "draw a sticker",
      cfg: {},
    });

    expect(postJsonRequestMock).toHaveBeenCalledWith(
      expect.objectContaining({
        timeoutMs: 180_000,
      }),
    );
  });

  it("sends reference images as data URLs for edit-style requests", async () => {
    postJsonRequestMock.mockResolvedValue({
      response: {
        json: async () => ({
          choices: [
            {
              message: {
                content: [
                  {
                    type: "image_url",
                    image_url: {
                      url: `data:image/webp;base64,${Buffer.from("webp-one").toString("base64")}`,
                    },
                  },
                ],
              },
            },
          ],
        }),
      },
      release: vi.fn(async () => {}),
    });

    const provider = buildOpenRouterImageGenerationProvider();
    const result = await provider.generateImage({
      provider: "openrouter",
      model: "google/gemini-3.1-flash-image-preview",
      prompt: "turn this into watercolor",
      inputImages: [{ buffer: Buffer.from("source-image"), mimeType: "image/png" }],
      cfg: {},
    });

    const body = requireOpenRouterPostBody();
    expect(body.messages?.[0]?.content).toEqual([
      { type: "text", text: "turn this into watercolor" },
      {
        type: "image_url",
        image_url: {
          url: `data:image/png;base64,${Buffer.from("source-image").toString("base64")}`,
        },
      },
    ]);
    const image = requireGeneratedImage(result, 0);
    expect(image.buffer.toString()).toBe("webp-one");
    expect(image.mimeType).toBe("image/webp");
  });

  it("wraps wrong-shape successful OpenRouter image responses", async () => {
    postJsonRequestMock.mockResolvedValue({
      response: {
        json: async () => ({ choices: { message: {} } }),
      },
      release: vi.fn(async () => {}),
    });

    const provider = buildOpenRouterImageGenerationProvider();
    await expect(
      provider.generateImage({
        provider: "openrouter",
        model: "google/gemini-3.1-flash-image-preview",
        prompt: "bad shape",
        cfg: {},
      }),
    ).rejects.toThrow("OpenRouter image generation response malformed");
  });

  it("extracts image fallbacks from string content and raw b64 parts", async () => {
    const png = Buffer.from("png-inline").toString("base64");
    const raw = Buffer.from("raw-inline").toString("base64");
    postJsonRequestMock.mockResolvedValue({
      response: {
        json: async () => ({
          choices: [
            {
              message: {
                content: `done data:image/png;base64,${png}`,
              },
            },
            {
              message: {
                content: [{ b64_json: raw }],
              },
            },
          ],
        }),
      },
      release: vi.fn(async () => {}),
    });

    const result = await buildOpenRouterImageGenerationProvider().generateImage({
      provider: "openrouter",
      model: "google/gemini-3.1-flash-image-preview",
      prompt: "draw image fallbacks",
      cfg: {},
    });

    expect(result.images.map((image) => image.buffer.toString())).toEqual([
      "png-inline",
      "raw-inline",
    ]);
  });

  it("rejects invalid raw image parts in strict extraction mode", async () => {
    postJsonRequestMock.mockResolvedValue({
      response: {
        json: async () => ({
          choices: [
            {
              message: {
                content: [{ b64_json: "not-base64!" }],
              },
            },
          ],
        }),
      },
      release: vi.fn(async () => {}),
    });

    await expect(
      buildOpenRouterImageGenerationProvider().generateImage({
        provider: "openrouter",
        model: "google/gemini-3.1-flash-image-preview",
        prompt: "draw invalid fallback",
        cfg: {},
      }),
    ).rejects.toThrow("OpenRouter image generation response malformed");
  });
});