OmniRoute/tests/unit/guardrails/visionBridgeHelpers.callVisionModel.test.ts

/**
 * Tests for callVisionModel helper function.
 */

import test from "node:test";
import assert from "node:assert/strict";
import dns from "node:dns";
import { callVisionModel, type VisionModelConfig } from "@/lib/guardrails/visionBridgeHelpers";
import { createProviderConnection } from "@/lib/db/providers";
import { resetDbInstance } from "@/lib/db/core";

// Store original fetch
const originalFetch = globalThis.fetch;

// Stub DNS for fetchRemoteImage's GHSA-cmhj-wh2f-9cgx DNS-rebinding guard
// (assertHostnameResolvesPublic in src/shared/network/remoteImageFetch.ts).
// These tests mock globalThis.fetch with example.com hosts that don't actually
// resolve in CI; the call path (callVisionModel -> fetchRemoteImageAsDataUri)
// does not expose a way to inject a `lookup` stub through to fetchRemoteImage,
// so we monkey-patch dns.promises.lookup with a pass-through public-IP
// resolver. Node --test runs each test file in its own process, so this
// rebinding does not leak across files.
const originalDnsLookup = dns.promises.lookup;
(dns.promises as { lookup: unknown }).lookup = (async (
  _hostname: string,
  options?: { all?: boolean }
) => {
  const record = { address: "203.0.113.1", family: 4 };
  return options && options.all ? [record] : record;
}) as typeof dns.promises.lookup;
process.on("exit", () => {
  (dns.promises as { lookup: unknown }).lookup = originalDnsLookup;
});

// (#8430) getBestVisionModel now validates that a `fixedModel` has a usable
// connection (via hasUsableCredentialsForModel, which queries the real DB)
// before returning it — an unreachable fixedModel falls through to
// auto-selection and, with nothing else configured either, resolves to `null`,
// which callVisionModel turns into a hard "No vision-capable provider
// connected" error before it ever reaches the HTTP call these tests mock.
// The router/credential-selection logic itself is already covered by
// visionBridgeRouter.test.ts and repro-8430.test.ts; these tests exercise
// callVisionModel's own request/response handling, so they just need one
// usable connection seeded per provider they use ("openai/gpt-4o-mini",
// "anthropic/claude-3-haiku") so getBestVisionModel resolves the requested
// fixedModel unchanged instead of null.
test.before(async () => {
  await createProviderConnection({
    provider: "openai",
    authType: "apikey",
    name: "vision-bridge-test-openai",
    apiKey: "sk-test-openai",
  });
  await createProviderConnection({
    provider: "anthropic",
    authType: "apikey",
    name: "vision-bridge-test-anthropic",
    apiKey: "sk-test-anthropic",
  });
});

test.after(() => {
  resetDbInstance();
});

test("callVisionModel returns description on success", async () => {
  // Mock global fetch
  const mockResponse = {
    ok: true,
    json: async () => ({
      choices: [{ message: { content: "A beautiful sunset over the ocean" } }],
    }),
  };
  globalThis.fetch = async () => mockResponse as unknown as Response;

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    const result = await callVisionModel(
      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
      config
    );

    assert.strictEqual(result, "A beautiful sunset over the ocean");
  } finally {
    // Restore original fetch
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel can route a catalog model through the OmniRoute self-loop", async () => {
  let capturedUrl = "";
  let capturedBody: Record<string, unknown> = {};
  let capturedHeaders: Record<string, string> = {};
  const fetchImpl: typeof fetch = async (input, init) => {
    capturedUrl = String(input);
    capturedBody = JSON.parse(String(init?.body));
    capturedHeaders = (init?.headers ?? {}) as Record<string, string>;
    return Response.json({ choices: [{ message: { content: "GREEN_SCENE_2" } }] });
  };

  const result = await callVisionModel("data:image/png;base64,iVBORw0KGgo", {
    model: "openai/gpt-4o-mini",
    prompt: "Describe this frame",
    timeoutMs: 30000,
    maxImages: 1,
    routeThroughOmniRoute: true,
    fetchImpl,
  });

  const url = new URL(capturedUrl);
  assert.equal(url.hostname, "localhost");
  assert.equal(url.pathname, "/v1/chat/completions");
  assert.equal(capturedBody.model, "openai/gpt-4o-mini");
  assert.equal(capturedHeaders["x-omniroute-admission-bypass"], "internal");
  assert.match(capturedHeaders["x-omniroute-disabled-guardrails"], /video-bridge/);
  assert.equal(result, "GREEN_SCENE_2");
});

test("callVisionModel throws on HTTP error", async () => {
  const mockResponse = {
    ok: false,
    status: 500,
    text: async () => "Internal Server Error",
  };
  globalThis.fetch = async () => mockResponse as unknown as Response;

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    await assert.rejects(
      async () => await callVisionModel("data:image/png;base64,iVBORw0KGgo", config),
      /Vision API error 500/
    );
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel throws on API error response", async () => {
  const mockResponse = {
    ok: true,
    json: async () => ({
      error: { message: "Invalid API key" },
    }),
  };
  globalThis.fetch = async () => mockResponse as unknown as Response;

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    await assert.rejects(
      async () => await callVisionModel("data:image/png;base64,iVBORw0KGgo", config),
      /Invalid API key/
    );
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel throws on empty response", async () => {
  const mockResponse = {
    ok: true,
    json: async () => ({
      choices: [{}],
    }),
  };
  globalThis.fetch = async () => mockResponse as unknown as Response;

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    await assert.rejects(
      async () => await callVisionModel("data:image/png;base64,iVBORw0KGgo", config),
      /empty or invalid/
    );
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel trims whitespace from response", async () => {
  const mockResponse = {
    ok: true,
    json: async () => ({
      choices: [{ message: { content: "  A test description  " } }],
    }),
  };
  globalThis.fetch = async () => mockResponse as unknown as Response;

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    const result = await callVisionModel("data:image/png;base64,iVBORw0KGgo", config);
    assert.strictEqual(result, "A test description");
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel passes custom API key", async () => {
  let capturedHeaders: Record<string, string> = {};

  const mockResponse = {
    ok: true,
    json: async () => ({
      choices: [{ message: { content: "Description" } }],
    }),
  };

  globalThis.fetch = async (url: URL | RequestInfo, init?: RequestInit) => {
    if (init?.headers) {
      capturedHeaders = init.headers as Record<string, string>;
    }
    return mockResponse as unknown as Response;
  };

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    await callVisionModel("data:image/png;base64,iVBORw0KGgo", config, "sk-custom-key");

    assert.strictEqual(capturedHeaders["Authorization"], "Bearer sk-custom-key");
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel uses correct request body format", async () => {
  let capturedBody: Record<string, unknown> = {};

  const mockResponse = {
    ok: true,
    json: async () => ({
      choices: [{ message: { content: "Description" } }],
    }),
  };

  globalThis.fetch = async (url: URL | RequestInfo, init?: RequestInit) => {
    if (init?.body) {
      capturedBody = JSON.parse(init.body as string);
    }
    return mockResponse as unknown as Response;
  };

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "What is in this image?",
      timeoutMs: 30000,
      maxImages: 10,
    };

    const imageUri = "data:image/png;base64,test123";
    await callVisionModel(imageUri, config);

    // Verify request structure
    assert.strictEqual(capturedBody.model, "gpt-4o-mini");
    assert.strictEqual(
      capturedBody.stream,
      false,
      "the JSON parser requires the internal vision request to opt out of SSE"
    );
    assert.ok(Array.isArray(capturedBody.messages));
    assert.strictEqual((capturedBody.messages as unknown[]).length, 1);

    const message = (capturedBody.messages as Array<{ role: string; content: unknown[] }>)[0];
    assert.strictEqual(message.role, "user");
    assert.ok(Array.isArray(message.content));
    assert.strictEqual(message.content.length, 2);

    // First content is image_url
    const imagePart = message.content[0] as {
      type: string;
      image_url: { url: string; detail: string };
    };
    assert.strictEqual(imagePart.type, "image_url");
    assert.strictEqual(imagePart.image_url.url, imageUri);
    assert.strictEqual(imagePart.image_url.detail, "high");

    // Second content is text prompt
    const textPart = message.content[1] as { type: string; text: string };
    assert.strictEqual(textPart.type, "text");
    assert.strictEqual(textPart.text, "What is in this image?");
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel fetches remote images before Anthropic requests", async () => {
  const fetchCalls: Array<{ url: string; init?: RequestInit }> = [];

  globalThis.fetch = async (url: URL | RequestInfo, init?: RequestInit) => {
    const requestUrl = String(url);
    fetchCalls.push({ url: requestUrl, init });

    if (requestUrl === "https://cdn.example.com/cat.png") {
      return new Response(Buffer.from("cat-image-bytes"), {
        status: 200,
        headers: { "Content-Type": "image/png" },
      });
    }

    return new Response(
      JSON.stringify({
        content: [{ type: "text", text: "A cat sitting on a chair" }],
      }),
      {
        status: 200,
        headers: { "Content-Type": "application/json" },
      }
    );
  };

  try {
    const config: VisionModelConfig = {
      model: "anthropic/claude-3-haiku",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
    };

    const result = await callVisionModel("https://cdn.example.com/cat.png", config, "sk-ant");

    assert.strictEqual(result, "A cat sitting on a chair");
    assert.strictEqual(fetchCalls.length, 2);
    assert.strictEqual(fetchCalls[0].url, "https://cdn.example.com/cat.png");
    assert.strictEqual(fetchCalls[1].url, "https://api.anthropic.com/v1/messages");

    const anthropicBody = JSON.parse(fetchCalls[1].init?.body as string);
    const imagePart = anthropicBody.messages[0].content[0];
    const imageSource = imagePart.source;
    assert.strictEqual(imageSource.type, "base64");
    assert.strictEqual(imageSource.media_type, "image/png");
    assert.strictEqual(imageSource.data, Buffer.from("cat-image-bytes").toString("base64"));
    // Compatibility guard for the global (not OpenCode-scoped) `detail: "high"`
    // default added to the OpenAI-compatible describe path: Anthropic's wire
    // format has no `detail` concept, so the describe self-loop must not leak
    // an OpenAI-only field into the Anthropic request body.
    assert.strictEqual(imagePart.detail, undefined);
    assert.strictEqual(imageSource.detail, undefined);
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("callVisionModel propagates an external abort to fetch and stops before fallback", async () => {
  const controller = new AbortController();
  let fetchCalls = 0;
  let fetchSignal: AbortSignal | null = null;

  globalThis.fetch = async (_url: URL | RequestInfo, init?: RequestInit) => {
    fetchCalls += 1;
    fetchSignal = init?.signal instanceof AbortSignal ? init.signal : null;
    controller.abort();
    const error = new Error("private aborted request detail");
    error.name = "AbortError";
    throw error;
  };

  try {
    const config: VisionModelConfig = {
      model: "openai/gpt-4o-mini",
      prompt: "Describe this image",
      timeoutMs: 30_000,
      maxImages: 10,
      signal: controller.signal,
    };

    await assert.rejects(
      () =>
        callVisionModel(
          "data:image/png;base64,iVBORw0KGgo",
          config,
          "sk-test",
          { maxFallbackAttempts: 2 },
          {
            hasUsableCredentials: async (model) =>
              model === "openai/gpt-4o-mini" || model.startsWith("anthropic/"),
          }
        ),
      /timed out|aborted/i
    );
    assert.equal(fetchCalls, 1, "an aborted parent request must not try a fallback model");
    assert.equal(fetchSignal?.aborted, true, "the parent abort must reach the active fetch");
  } finally {
    globalThis.fetch = originalFetch;
  }
});

// GHSA-34rg-3pqj-35g9 — the Anthropic describe self-call inlines the user's image URL to
// base64 through the same `fetchRemoteImageAsDataUri()` sink as the claude-wire reroute.
// The DNS stub at the top of this file answers a public IP for every hostname, so only the
// `public-only` string check stands between the request body and a loopback/RFC-1918 fetch.
for (const privateUrl of ["http://127.0.0.1:1/x.png", "http://192.168.1.50/x.png"]) {
  test(`callVisionModel never fetches a private image URL (${privateUrl}) for the Anthropic describe path (GHSA-34rg-3pqj-35g9)`, async () => {
    const fetchedUrls: string[] = [];
    const fetchImpl: typeof fetch = async (url) => {
      const requestUrl = String(url);
      fetchedUrls.push(requestUrl);
      if (requestUrl === privateUrl) {
        // Canary: on the vulnerable code these bytes are inlined into the Anthropic body.
        return new Response(Buffer.from("intranet-bytes"), {
          status: 200,
          headers: { "Content-Type": "image/png" },
        });
      }
      return new Response(JSON.stringify({ content: [{ type: "text", text: "described" }] }), {
        status: 200,
        headers: { "Content-Type": "application/json" },
      });
    };

    const config: VisionModelConfig = {
      model: "anthropic/claude-3-haiku",
      prompt: "Describe this image",
      timeoutMs: 30000,
      maxImages: 10,
      fetchImpl,
    };

    await assert.rejects(
      () => callVisionModel(privateUrl, config, "sk-ant", { maxFallbackAttempts: 1 }),
      /blocked/i
    );
    assert.deepStrictEqual(
      fetchedUrls,
      [],
      "neither the private download nor the self-call may happen"
    );
  });
}