mirror of
https://github.com/MBombeck/HealthLog.git
synced 2026-11-09 13:01:00 +00:00
test(ai): pin the gateway provider's isolation and host policy
The properties this provider has to keep: it resolves against the base URL that was saved for it, the pinned OpenAI arm cannot see that URL and still posts to api.openai.com, and the OpenAI key never travels to the gateway. Each is asserted at the wire, not on internals. The SSRF floor is checked through the same alt-notation classes the Local provider is checked against, and the pinned tags are shown to keep requiring a public host even when the operator opened every private one. The AI settings connect-src is pinned to what it already allowed: the gateway is called server-side, so no CSP change is needed, and a future change that needs one will turn this red.
This commit is contained in:
Marc-André Bombeck committed
1 parent
ee620b2707
commit
3ef6a7119b
5 files changed
+852
No files matched your search
@@ -0,0 +1,61 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { NextRequest } from "next/server";
|
||||
|
||||
/**
|
||||
* The OpenAI-compatible gateway provider (#470) needs no CSP change, and this
|
||||
* pins that.
|
||||
*
|
||||
* The reasoning is worth stating because the opposite conclusion is the
|
||||
* tempting one: the gateway's host is user-configured, so if the browser ever
|
||||
* had to reach it, `connect-src` would need a wildcard — which would undo the
|
||||
* surgical gating the AI settings surface has today. It does not: the
|
||||
* provider is resolved and called server-side (`/api/ai/test`,
|
||||
* `/api/insights/*`), so the browser only ever talks to this origin.
|
||||
*
|
||||
* If a future change makes the settings page call a gateway directly, this
|
||||
* test goes red, and that is the moment to have the argument.
|
||||
*/
|
||||
|
||||
vi.mock("@/lib/process-type", () => ({
|
||||
shouldRunWeb: () => true,
|
||||
}));
|
||||
|
||||
import { proxy } from "../proxy";
|
||||
|
||||
function setNodeEnv(value: "development" | "production") {
|
||||
vi.stubEnv("NODE_ENV", value);
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
setNodeEnv("production");
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.unstubAllEnvs();
|
||||
});
|
||||
|
||||
function connectSrcFor(pathname: string): string {
|
||||
const res = proxy(
|
||||
new NextRequest(`http://localhost${pathname}`, {
|
||||
headers: { cookie: "healthlog_session=sess-1" },
|
||||
}),
|
||||
);
|
||||
const csp = res.headers.get("content-security-policy") ?? "";
|
||||
const directive = csp
|
||||
.split(";")
|
||||
.map((part) => part.trim())
|
||||
.find((part) => part.startsWith("connect-src"));
|
||||
return directive ?? "";
|
||||
}
|
||||
|
||||
describe("AI settings connect-src is untouched by the gateway provider", () => {
|
||||
it("still allows exactly the two AI hosts it allowed before", () => {
|
||||
expect(connectSrcFor("/settings/ai")).toBe(
|
||||
"connect-src 'self' https://api.openai.com https://chatgpt.com",
|
||||
);
|
||||
});
|
||||
|
||||
it("adds no AI host anywhere else", () => {
|
||||
expect(connectSrcFor("/dashboard")).toBe("connect-src 'self'");
|
||||
});
|
||||
});
|
||||
@@ -261,6 +261,27 @@ describe("POST /api/ai/test — dropdown-aware override", () => {
|
||||
expect(body.data.model).toBe("claude-3-5-sonnet-latest");
|
||||
});
|
||||
|
||||
it("forwards the gateway's own override fields (#470)", async () => {
|
||||
makeProviderThatSucceeds("anthropic/claude-sonnet-4-6");
|
||||
const response = await POST(
|
||||
jsonRequest({
|
||||
provider: "OPENAI_COMPATIBLE",
|
||||
compatBaseUrl: "https://litellm.example.com/v1",
|
||||
compatKey: "gateway-secret",
|
||||
compatModel: "anthropic/claude-sonnet-4-6",
|
||||
}) as never,
|
||||
);
|
||||
expect(response.status).toBe(200);
|
||||
// The gateway's fields travel under their own names — the pinned OpenAI
|
||||
// arm reads `baseUrl` / `openaiKey` and must not be reachable from here.
|
||||
expect(resolveProviderForTest).toHaveBeenCalledWith("u-1", {
|
||||
provider: "OPENAI_COMPATIBLE",
|
||||
compatBaseUrl: "https://litellm.example.com/v1",
|
||||
compatKey: "gateway-secret",
|
||||
compatModel: "anthropic/claude-sonnet-4-6",
|
||||
});
|
||||
});
|
||||
|
||||
it("rejects unknown provider values with 422", async () => {
|
||||
const response = await POST(
|
||||
jsonRequest({ provider: "MALICIOUS" }) as never,
|
||||
|
||||
@@ -46,6 +46,8 @@ vi.mock("@/lib/logging/context", () => ({
|
||||
import { GET, PATCH } from "../route";
|
||||
import { prisma } from "@/lib/db";
|
||||
import { resolveProviderAvailability } from "@/lib/ai/provider";
|
||||
import { isPublicUrl } from "@/lib/validations/notifications";
|
||||
import { isLocalAiHostAllowed } from "@/lib/ai/local-host-allowlist";
|
||||
|
||||
function patchRequest(body: unknown): Request {
|
||||
return new Request("http://localhost/api/user/ai-provider", {
|
||||
@@ -192,3 +194,110 @@ describe("PATCH /api/user/ai-provider — v1.22 fields", () => {
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// v1.33.1 (#470) — the OpenAI-compatible gateway writes three columns of its
|
||||
// own. What matters at this layer: the SSRF floor applies to the gateway's
|
||||
// base URL exactly as it does to the Local one, and the column is nobody
|
||||
// else's — switching the active provider away must not hand a stale URL to a
|
||||
// key-carrying provider, and must not silently discard the user's gateway
|
||||
// either.
|
||||
describe("PATCH /api/user/ai-provider — the gateway's own columns (#470)", () => {
|
||||
const patch = PATCH as (req: Request) => Promise<Response>;
|
||||
|
||||
it("persists base URL, model and encrypted key", async () => {
|
||||
vi.mocked(isPublicUrl).mockReturnValue(true);
|
||||
const res = await patch(
|
||||
patchRequest({
|
||||
provider: "OPENAI_COMPATIBLE",
|
||||
compatBaseUrl: "https://litellm.example.com/v1",
|
||||
compatModel: "anthropic/claude-sonnet-4-6",
|
||||
compatKey: "gateway-secret",
|
||||
}),
|
||||
);
|
||||
expect(res.status).toBe(200);
|
||||
expect(vi.mocked(prisma.user.update)).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
data: expect.objectContaining({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiCompatBaseUrl: "https://litellm.example.com/v1",
|
||||
aiCompatModel: "anthropic/claude-sonnet-4-6",
|
||||
aiCompatKeyEncrypted: "enc:gateway-secret",
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("refuses a private base URL the operator has not allowlisted", async () => {
|
||||
vi.mocked(isPublicUrl).mockReturnValue(false);
|
||||
vi.mocked(isLocalAiHostAllowed).mockReturnValue(false);
|
||||
const res = await patch(
|
||||
patchRequest({ compatBaseUrl: "http://169.254.169.254/v1" }),
|
||||
);
|
||||
expect(res.status).toBe(422);
|
||||
expect(vi.mocked(prisma.user.update)).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("accepts a private base URL the operator allowlisted", async () => {
|
||||
vi.mocked(isPublicUrl).mockReturnValue(false);
|
||||
vi.mocked(isLocalAiHostAllowed).mockReturnValue(true);
|
||||
const res = await patch(
|
||||
patchRequest({ compatBaseUrl: "http://litellm.lan:4000/v1" }),
|
||||
);
|
||||
expect(res.status).toBe(200);
|
||||
expect(vi.mocked(prisma.user.update)).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
data: expect.objectContaining({
|
||||
aiCompatBaseUrl: "http://litellm.lan:4000/v1",
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("leaves the gateway URL alone when the provider switches to OpenAI", async () => {
|
||||
vi.mocked(isPublicUrl).mockReturnValue(true);
|
||||
await patch(patchRequest({ provider: "OPENAI" }));
|
||||
const data = vi.mocked(prisma.user.update).mock.calls[0][0].data as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
// `aiBaseUrl` is the shared LOCAL column and IS cleared. The gateway's is
|
||||
// untouched — it belongs to one provider, and that provider is not this
|
||||
// one, so there is nothing for a stale value to redirect.
|
||||
expect(data.aiBaseUrl).toBeNull();
|
||||
expect(data).not.toHaveProperty("aiCompatBaseUrl");
|
||||
});
|
||||
});
|
||||
|
||||
describe("GET /api/user/ai-provider — the gateway's fields (#470)", () => {
|
||||
it("reports the base URL and model, and the key as presence only", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiModel: null,
|
||||
aiBaseUrl: null,
|
||||
aiAnthropicKeyEncrypted: null,
|
||||
aiLocalKeyEncrypted: null,
|
||||
aiOpenaiKeyEncrypted: null,
|
||||
aiCompatBaseUrl: "https://litellm.example.com/v1",
|
||||
aiCompatKeyEncrypted: "enc-gateway",
|
||||
aiCompatModel: "anthropic/claude-sonnet-4-6",
|
||||
aiResponseTimeoutSeconds: null,
|
||||
} as never);
|
||||
vi.mocked(resolveProviderAvailability).mockResolvedValue({
|
||||
aiAvailable: true,
|
||||
managedBy: "user",
|
||||
});
|
||||
|
||||
const res = await (GET as () => Promise<Response>)();
|
||||
const body = (await res.json()) as {
|
||||
data: {
|
||||
compatBaseUrl: string | null;
|
||||
compatModel: string | null;
|
||||
hasCompatKey: boolean;
|
||||
};
|
||||
};
|
||||
expect(body.data.compatBaseUrl).toBe("https://litellm.example.com/v1");
|
||||
expect(body.data.compatModel).toBe("anthropic/claude-sonnet-4-6");
|
||||
expect(body.data.hasCompatKey).toBe(true);
|
||||
expect(JSON.stringify(body)).not.toContain("enc-gateway");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,288 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { OpenAIClient } from "../openai-client";
|
||||
import { resetJsonModeDialectCache } from "../json-dialect";
|
||||
import { singleUserTurn } from "../types";
|
||||
|
||||
/**
|
||||
* The OpenAI-compatible gateway provider (#470), at the wire.
|
||||
*
|
||||
* Two properties carry the security weight of this provider and are pinned
|
||||
* here rather than left to review:
|
||||
*
|
||||
* 1. The pinned `admin-key` / `codex` tags keep `requirePublicHost: true`
|
||||
* no matter what the operator allowlisted. Only the gateway tag — whose
|
||||
* base URL a person typed — consults the allowlist.
|
||||
* 2. The SSRF floor under the gateway is `isPublicUrl`, which is the same
|
||||
* floor the Local provider sits on, including every IPv4 / IPv6
|
||||
* alt-notation class.
|
||||
*
|
||||
* `safeFetch` runs for real here (only the transport underneath it is
|
||||
* stubbed), so a private-host rejection in these tests is the production
|
||||
* rejection, not a mock of one.
|
||||
*/
|
||||
|
||||
// safeFetch's requirePublicHost path runs through undici's own `fetch`
|
||||
// (version-locked with its dispatcher). Delegate it to the global `fetch`
|
||||
// stub these tests install so the interception still applies.
|
||||
vi.mock("undici", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("undici")>();
|
||||
return {
|
||||
...actual,
|
||||
fetch: (input: unknown, init?: unknown) =>
|
||||
(globalThis.fetch as unknown as (i: unknown, n?: unknown) => unknown)(
|
||||
input,
|
||||
init,
|
||||
),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock("@/lib/logging/context", () => ({ annotate: vi.fn() }));
|
||||
|
||||
function okReply(content = '{"ok":true}') {
|
||||
return vi.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
choices: [{ message: { content }, finish_reason: "stop" }],
|
||||
usage: { total_tokens: 7 },
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
function errorReply(status: number, body: string) {
|
||||
return {
|
||||
ok: false,
|
||||
status,
|
||||
text: () => Promise.resolve(body),
|
||||
};
|
||||
}
|
||||
|
||||
const jsonTurn = () =>
|
||||
singleUserTurn({
|
||||
system: "system",
|
||||
user: "user",
|
||||
responseFormat: "json",
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
resetJsonModeDialectCache();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.unstubAllEnvs();
|
||||
});
|
||||
|
||||
describe("OpenAI-compatible gateway — wire", () => {
|
||||
it("posts to the configured base URL and reports its own provider tag", async () => {
|
||||
const fetchMock = okReply();
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const client = new OpenAIClient({
|
||||
apiKey: "gateway-secret",
|
||||
model: "anthropic/claude-sonnet-4-6",
|
||||
baseUrl: "https://litellm.example.com/v1",
|
||||
providerType: "openai-compatible",
|
||||
});
|
||||
const result = await client.generateCompletion(jsonTurn());
|
||||
|
||||
expect(fetchMock.mock.calls[0][0]).toBe(
|
||||
"https://litellm.example.com/v1/chat/completions",
|
||||
);
|
||||
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe(
|
||||
"Bearer gateway-secret",
|
||||
);
|
||||
expect(result.providerType).toBe("openai-compatible");
|
||||
expect(result.model).toBe("anthropic/claude-sonnet-4-6");
|
||||
});
|
||||
|
||||
it("sends no Authorization header at all when the gateway needs no bearer", async () => {
|
||||
const fetchMock = okReply();
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
await new OpenAIClient({
|
||||
apiKey: "",
|
||||
model: "llama3.1:70b",
|
||||
baseUrl: "https://gateway.example.com/v1",
|
||||
providerType: "openai-compatible",
|
||||
}).generateCompletion(jsonTurn());
|
||||
|
||||
expect(fetchMock.mock.calls[0][1].headers).not.toHaveProperty(
|
||||
"Authorization",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("OpenAI-compatible gateway — host policy", () => {
|
||||
it("refuses a private host that the operator has not allowlisted", async () => {
|
||||
vi.stubGlobal("fetch", okReply());
|
||||
vi.stubEnv("ALLOW_LOCAL_AI_PRIVATE_HOSTS", "");
|
||||
|
||||
await expect(
|
||||
new OpenAIClient({
|
||||
apiKey: "",
|
||||
model: "m",
|
||||
baseUrl: "http://192.168.1.5:4000/v1",
|
||||
providerType: "openai-compatible",
|
||||
}).generateCompletion(jsonTurn()),
|
||||
).rejects.toMatchObject({ kind: "private_host" });
|
||||
});
|
||||
|
||||
// The alt-notation classes `isPublicUrl` exists to catch. A gateway is the
|
||||
// one AI surface whose host a user types, so every one of these has to be
|
||||
// refused through the gateway arm too — not only through the Local one.
|
||||
it.each([
|
||||
["decimal IPv4", "http://2130706433/v1"],
|
||||
["hex IPv4", "http://0x7f000001/v1"],
|
||||
["octal IPv4", "http://0177.0.0.1/v1"],
|
||||
["cloud metadata", "http://169.254.169.254/v1"],
|
||||
["IPv4-mapped IPv6", "http://[::ffff:169.254.169.254]/v1"],
|
||||
["IPv6 ULA", "http://[fd00::1]/v1"],
|
||||
["loopback", "http://127.0.0.1:8080/v1"],
|
||||
])("refuses %s", async (_label, baseUrl) => {
|
||||
vi.stubGlobal("fetch", okReply());
|
||||
vi.stubEnv("ALLOW_LOCAL_AI_PRIVATE_HOSTS", "");
|
||||
|
||||
await expect(
|
||||
new OpenAIClient({
|
||||
apiKey: "",
|
||||
model: "m",
|
||||
baseUrl,
|
||||
providerType: "openai-compatible",
|
||||
}).generateCompletion(jsonTurn()),
|
||||
).rejects.toMatchObject({ kind: "private_host" });
|
||||
});
|
||||
|
||||
it("reaches a private host the operator allowlisted by name", async () => {
|
||||
const fetchMock = okReply();
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
vi.stubEnv("ALLOW_LOCAL_AI_PRIVATE_HOSTS", "litellm.lan");
|
||||
|
||||
await new OpenAIClient({
|
||||
apiKey: "",
|
||||
model: "m",
|
||||
baseUrl: "http://litellm.lan:4000/v1",
|
||||
providerType: "openai-compatible",
|
||||
}).generateCompletion(jsonTurn());
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("keeps the pinned tags on requirePublicHost even when the operator allowlisted the host", async () => {
|
||||
vi.stubGlobal("fetch", okReply());
|
||||
// The operator opened every private host for their local AI endpoint.
|
||||
// That opt-in belongs to the gateway/local surface; a client carrying an
|
||||
// OpenAI-issued key must not inherit it.
|
||||
vi.stubEnv("ALLOW_LOCAL_AI_PRIVATE_HOSTS", "true");
|
||||
|
||||
await expect(
|
||||
new OpenAIClient({
|
||||
apiKey: "sk-user-openai-key",
|
||||
model: "gpt-4o",
|
||||
baseUrl: "http://192.168.1.5:4000/v1",
|
||||
}).generateCompletion(jsonTurn()),
|
||||
).rejects.toMatchObject({ kind: "private_host" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("OpenAI-compatible gateway — JSON-mode dialect", () => {
|
||||
it("retries once without response_format when the gateway rejects the field", async () => {
|
||||
const fetchMock = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(
|
||||
errorReply(400, '{"error":"Unknown parameter: response_format"}'),
|
||||
)
|
||||
.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
choices: [{ message: { content: '{"ok":true}' } }],
|
||||
usage: { total_tokens: 3 },
|
||||
}),
|
||||
});
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const result = await new OpenAIClient({
|
||||
apiKey: "k",
|
||||
model: "m",
|
||||
baseUrl: "https://strict-gateway.example.com/v1",
|
||||
providerType: "openai-compatible",
|
||||
}).generateCompletion(jsonTurn());
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
expect(JSON.parse(fetchMock.mock.calls[0][1].body)).toHaveProperty(
|
||||
"response_format",
|
||||
);
|
||||
expect(JSON.parse(fetchMock.mock.calls[1][1].body)).not.toHaveProperty(
|
||||
"response_format",
|
||||
);
|
||||
expect(result.content).toBe('{"ok":true}');
|
||||
});
|
||||
|
||||
it("remembers the no-flag dialect for that endpoint", async () => {
|
||||
const fetchMock = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce(
|
||||
errorReply(400, '{"error":"response_format is not supported"}'),
|
||||
)
|
||||
.mockResolvedValue({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
choices: [{ message: { content: "{}" } }],
|
||||
}),
|
||||
});
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const client = new OpenAIClient({
|
||||
apiKey: "k",
|
||||
model: "m",
|
||||
baseUrl: "https://strict-gateway.example.com/v1",
|
||||
providerType: "openai-compatible",
|
||||
});
|
||||
await client.generateCompletion(jsonTurn());
|
||||
await client.generateCompletion(jsonTurn());
|
||||
|
||||
// Three calls total: the rejected probe, its retry, and the second
|
||||
// generation — which must not probe again.
|
||||
expect(fetchMock).toHaveBeenCalledTimes(3);
|
||||
expect(JSON.parse(fetchMock.mock.calls[2][1].body)).not.toHaveProperty(
|
||||
"response_format",
|
||||
);
|
||||
});
|
||||
|
||||
it("does not retry an unrelated 4xx", async () => {
|
||||
const fetchMock = vi
|
||||
.fn()
|
||||
.mockResolvedValue(errorReply(404, '{"error":"model not found"}'));
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
await expect(
|
||||
new OpenAIClient({
|
||||
apiKey: "k",
|
||||
model: "nope",
|
||||
baseUrl: "https://gateway.example.com/v1",
|
||||
providerType: "openai-compatible",
|
||||
}).generateCompletion(jsonTurn()),
|
||||
).rejects.toMatchObject({ httpStatus: 404 });
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("never probes or retries on the pinned OpenAI endpoint", async () => {
|
||||
const fetchMock = vi
|
||||
.fn()
|
||||
.mockResolvedValue(
|
||||
errorReply(400, '{"error":"Unknown parameter: response_format"}'),
|
||||
);
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
await expect(
|
||||
new OpenAIClient({
|
||||
apiKey: "sk-key",
|
||||
model: "gpt-4o",
|
||||
baseUrl: "https://api.openai.com/v1",
|
||||
}).generateCompletion(jsonTurn()),
|
||||
).rejects.toThrow("OpenAI request failed (400)");
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,373 @@
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
/**
|
||||
* Resolution for the OpenAI-compatible gateway provider (#470).
|
||||
*
|
||||
* The claim this file has to keep true: the gateway resolves against the base
|
||||
* URL the user saved, and the pinned OpenAI arm cannot see that base URL or
|
||||
* be redirected by it. Both are asserted at the wire — each test resolves a
|
||||
* provider and runs a completion through a stubbed transport, so what is
|
||||
* checked is the request that would actually leave the process.
|
||||
*/
|
||||
|
||||
vi.mock("@/lib/db", () => ({
|
||||
prisma: {
|
||||
user: { findUnique: vi.fn(), update: vi.fn() },
|
||||
appSettings: { findUnique: vi.fn() },
|
||||
},
|
||||
}));
|
||||
vi.mock("@/lib/crypto", () => ({
|
||||
decrypt: vi.fn((v: string) => `decrypted:${v}`),
|
||||
encrypt: vi.fn((v: string) => `encrypted:${v}`),
|
||||
}));
|
||||
vi.mock("@/lib/ai/codex-oauth", () => ({
|
||||
refreshDeviceTokens: vi.fn(),
|
||||
encryptCodexCreds: vi.fn(),
|
||||
decryptCodexCreds: vi.fn(),
|
||||
}));
|
||||
vi.mock("@/lib/logging/context", () => ({ annotate: vi.fn() }));
|
||||
vi.mock("undici", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("undici")>();
|
||||
return {
|
||||
...actual,
|
||||
fetch: (input: unknown, init?: unknown) =>
|
||||
(globalThis.fetch as unknown as (i: unknown, n?: unknown) => unknown)(
|
||||
input,
|
||||
init,
|
||||
),
|
||||
};
|
||||
});
|
||||
|
||||
import {
|
||||
AITestConfigError,
|
||||
resolveProvider,
|
||||
resolveProviderAvailability,
|
||||
resolveProviderChain,
|
||||
resolveProviderForTest,
|
||||
} from "../provider";
|
||||
import {
|
||||
PROVIDER_CHAIN_DEFAULT,
|
||||
PROVIDER_CHAIN_TYPES,
|
||||
} from "../provider-chain";
|
||||
import { singleUserTurn } from "../types";
|
||||
import { prisma } from "@/lib/db";
|
||||
|
||||
/** A user row with every AI column present, overridable per test. */
|
||||
function userRow(overrides: Record<string, unknown> = {}) {
|
||||
return {
|
||||
aiProvider: null,
|
||||
aiModel: null,
|
||||
aiBaseUrl: null,
|
||||
aiAnthropicKeyEncrypted: null,
|
||||
aiLocalKeyEncrypted: null,
|
||||
aiOpenaiKeyEncrypted: null,
|
||||
aiCompatBaseUrl: null,
|
||||
aiCompatKeyEncrypted: null,
|
||||
aiCompatModel: null,
|
||||
aiProviderChain: null,
|
||||
useCentralCodex: false,
|
||||
codexConnectionStatus: null,
|
||||
codexAccessTokenEncrypted: null,
|
||||
codexRefreshTokenEncrypted: null,
|
||||
...overrides,
|
||||
} as never;
|
||||
}
|
||||
|
||||
function stubTransport() {
|
||||
const fetchMock = vi.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
choices: [{ message: { content: '{"ok":true}' } }],
|
||||
usage: { total_tokens: 1 },
|
||||
}),
|
||||
});
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
return fetchMock;
|
||||
}
|
||||
|
||||
const probe = () =>
|
||||
singleUserTurn({ system: "s", user: "u", responseFormat: "json" });
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
vi.mocked(prisma.appSettings.findUnique).mockResolvedValue(null as never);
|
||||
});
|
||||
|
||||
describe("the gateway resolves against its own columns", () => {
|
||||
it("posts to the saved base URL with the saved bearer and model", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiCompatBaseUrl: "https://litellm.example.com/v1",
|
||||
aiCompatKeyEncrypted: "enc-gateway",
|
||||
aiCompatModel: "anthropic/claude-sonnet-4-6",
|
||||
}),
|
||||
);
|
||||
|
||||
const provider = await resolveProvider("u-1");
|
||||
const result = await provider.generateCompletion(probe());
|
||||
|
||||
expect(fetchMock.mock.calls[0][0]).toBe(
|
||||
"https://litellm.example.com/v1/chat/completions",
|
||||
);
|
||||
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe(
|
||||
"Bearer decrypted:enc-gateway",
|
||||
);
|
||||
expect(JSON.parse(fetchMock.mock.calls[0][1].body).model).toBe(
|
||||
"anthropic/claude-sonnet-4-6",
|
||||
);
|
||||
expect(result.providerType).toBe("openai-compatible");
|
||||
});
|
||||
|
||||
it("works with no bearer at all — a gateway may need none", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiCompatBaseUrl: "https://gateway.example.com/v1",
|
||||
aiCompatModel: "qwen2.5",
|
||||
}),
|
||||
);
|
||||
|
||||
const provider = await resolveProvider("u-1");
|
||||
await provider.generateCompletion(probe());
|
||||
|
||||
expect(fetchMock.mock.calls[0][1].headers).not.toHaveProperty(
|
||||
"Authorization",
|
||||
);
|
||||
});
|
||||
|
||||
it("falls back to the shared model when no gateway model is set", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiModel: "shared-model",
|
||||
aiCompatBaseUrl: "https://gateway.example.com/v1",
|
||||
}),
|
||||
);
|
||||
|
||||
const provider = await resolveProvider("u-1");
|
||||
await provider.generateCompletion(probe());
|
||||
|
||||
expect(JSON.parse(fetchMock.mock.calls[0][1].body).model).toBe(
|
||||
"shared-model",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("the pinned OpenAI arm never reads the gateway's base URL", () => {
|
||||
it("still posts to api.openai.com when a gateway base URL is saved", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI",
|
||||
aiOpenaiKeyEncrypted: "enc-openai",
|
||||
// A configured gateway sits right there in the row.
|
||||
aiCompatBaseUrl: "https://attacker.example.com/v1",
|
||||
aiCompatKeyEncrypted: "enc-gateway",
|
||||
}),
|
||||
);
|
||||
|
||||
const provider = await resolveProvider("u-1");
|
||||
await provider.generateCompletion(probe());
|
||||
|
||||
expect(fetchMock.mock.calls[0][0]).toBe(
|
||||
"https://api.openai.com/v1/chat/completions",
|
||||
);
|
||||
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe(
|
||||
"Bearer decrypted:enc-openai",
|
||||
);
|
||||
});
|
||||
|
||||
it("resolves the chain's openai entry against api.openai.com regardless of the gateway", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiOpenaiKeyEncrypted: "enc-openai",
|
||||
aiCompatBaseUrl: "https://attacker.example.com/v1",
|
||||
aiProviderChain: [
|
||||
{ providerType: "openai", priority: 1, enabled: true },
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
const chain = await resolveProviderChain("u-1");
|
||||
expect(chain).toHaveLength(1);
|
||||
await chain[0].instance.generateCompletion(probe());
|
||||
|
||||
expect(fetchMock.mock.calls[0][0]).toBe(
|
||||
"https://api.openai.com/v1/chat/completions",
|
||||
);
|
||||
});
|
||||
|
||||
it("keeps the test endpoint's OpenAI arm pinned too", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI",
|
||||
aiOpenaiKeyEncrypted: "enc-openai",
|
||||
aiCompatBaseUrl: "https://attacker.example.com/v1",
|
||||
}),
|
||||
);
|
||||
|
||||
const provider = await resolveProviderForTest("u-1", {});
|
||||
await provider.generateCompletion(probe());
|
||||
|
||||
expect(fetchMock.mock.calls[0][0]).toBe(
|
||||
"https://api.openai.com/v1/chat/completions",
|
||||
);
|
||||
});
|
||||
|
||||
it("never sends the OpenAI key to the gateway", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiOpenaiKeyEncrypted: "enc-openai",
|
||||
aiCompatBaseUrl: "https://gateway.example.com/v1",
|
||||
aiCompatModel: "m",
|
||||
}),
|
||||
);
|
||||
|
||||
const provider = await resolveProvider("u-1");
|
||||
await provider.generateCompletion(probe());
|
||||
|
||||
const authorization = fetchMock.mock.calls[0][1].headers.Authorization;
|
||||
expect(authorization).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("an unconfigured gateway entry is skipped, not an error", () => {
|
||||
it("drops the chain entry when no base URL is saved", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiCompatModel: "m",
|
||||
aiProviderChain: [
|
||||
{ providerType: "openai-compatible", priority: 1, enabled: true },
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(resolveProviderChain("u-1")).resolves.toEqual([]);
|
||||
});
|
||||
|
||||
it("drops the chain entry when no model can be resolved", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiCompatBaseUrl: "https://gateway.example.com/v1",
|
||||
aiProviderChain: [
|
||||
{ providerType: "openai-compatible", priority: 1, enabled: true },
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(resolveProviderChain("u-1")).resolves.toEqual([]);
|
||||
});
|
||||
|
||||
it("resolves the entry once both are present", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiCompatBaseUrl: "https://gateway.example.com/v1",
|
||||
aiCompatModel: "m",
|
||||
aiProviderChain: [
|
||||
{ providerType: "openai-compatible", priority: 1, enabled: true },
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
const chain = await resolveProviderChain("u-1");
|
||||
expect(chain.map((e) => e.providerType)).toEqual(["openai-compatible"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("chain vocabulary", () => {
|
||||
it("offers the gateway as a chain type", () => {
|
||||
expect(PROVIDER_CHAIN_TYPES).toContain("openai-compatible");
|
||||
});
|
||||
|
||||
it("keeps it out of the default chain — it can only resolve once configured", () => {
|
||||
expect(
|
||||
PROVIDER_CHAIN_DEFAULT.map((entry) => entry.providerType),
|
||||
).not.toContain("openai-compatible");
|
||||
});
|
||||
});
|
||||
|
||||
describe("availability reporting", () => {
|
||||
it("reports a configured gateway as the user's own provider", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiCompatBaseUrl: "https://gateway.example.com/v1",
|
||||
aiCompatModel: "m",
|
||||
aiProviderChain: [
|
||||
{ providerType: "openai-compatible", priority: 1, enabled: true },
|
||||
],
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(resolveProviderAvailability("u-1")).resolves.toEqual({
|
||||
aiAvailable: true,
|
||||
managedBy: "user",
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("the connection test's gateway arm", () => {
|
||||
it("refuses a private base URL the operator has not allowlisted", async () => {
|
||||
vi.stubEnv("ALLOW_LOCAL_AI_PRIVATE_HOSTS", "");
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(
|
||||
userRow({
|
||||
aiProvider: "OPENAI_COMPATIBLE",
|
||||
aiCompatBaseUrl: "http://169.254.169.254/v1",
|
||||
aiCompatModel: "m",
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(
|
||||
resolveProviderForTest("u-1", { provider: "OPENAI_COMPATIBLE" }),
|
||||
).rejects.toBeInstanceOf(AITestConfigError);
|
||||
vi.unstubAllEnvs();
|
||||
});
|
||||
|
||||
it("refuses an unsaved gateway config with no base URL", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(userRow());
|
||||
|
||||
await expect(
|
||||
resolveProviderForTest("u-1", { provider: "OPENAI_COMPATIBLE" }),
|
||||
).rejects.toThrow(/base URL/i);
|
||||
});
|
||||
|
||||
it("refuses an unsaved gateway config with no model", async () => {
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(userRow());
|
||||
|
||||
await expect(
|
||||
resolveProviderForTest("u-1", {
|
||||
provider: "OPENAI_COMPATIBLE",
|
||||
compatBaseUrl: "https://gateway.example.com/v1",
|
||||
}),
|
||||
).rejects.toThrow(/model/i);
|
||||
});
|
||||
|
||||
it("tests an unsaved gateway config without persisting anything", async () => {
|
||||
const fetchMock = stubTransport();
|
||||
vi.mocked(prisma.user.findUnique).mockResolvedValue(userRow());
|
||||
|
||||
const provider = await resolveProviderForTest("u-1", {
|
||||
provider: "OPENAI_COMPATIBLE",
|
||||
compatBaseUrl: "https://unsaved.example.com/v1",
|
||||
compatKey: "typed-but-unsaved",
|
||||
compatModel: "typed-model",
|
||||
});
|
||||
await provider.generateCompletion(probe());
|
||||
|
||||
expect(fetchMock.mock.calls[0][0]).toBe(
|
||||
"https://unsaved.example.com/v1/chat/completions",
|
||||
);
|
||||
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe(
|
||||
"Bearer typed-but-unsaved",
|
||||
);
|
||||
expect(prisma.user.update).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user