Files
diegosouzapw--omniroute/tests/unit/combo-account-allowlist-3266.test.ts
2026-07-13 13:39:12 +08:00

259 lines
9.2 KiB
TypeScript

/**
* #3266 — Per-combo account allowlist.
*
* A combo model step can carry a first-class `allowedConnectionIds` so a
* round-robin / weighted strategy is scoped to a subset of a provider's
* connections (e.g. {foo1, foo2}) without hand-pinning one step per account.
*
* Acceptance (owner): a round-robin scoped to {foo1, foo2} over an active pool
* {foo1..foo4} never selects foo3/foo4 on real chat requests.
*
* Coverage:
* 1. steps.ts parses `allowedConnectionIds` (trim + drop empty).
* 2. handleComboChat propagates the step allowlist onto target.allowedConnectionIds.
* 3. getProviderCredentials never hands back a connection outside the allowlist.
*/
import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-combo-allowlist-"));
process.env.DATA_DIR = TEST_DATA_DIR;
const core = await import("../../src/lib/db/core.ts");
const providersDb = await import("../../src/lib/db/providers.ts");
const { buildAutoCandidates, handleComboChat } = await import("../../open-sse/services/combo.ts");
const { getProviderCredentials } = await import("../../src/sse/services/auth.ts");
const { normalizeComboStep } = await import("../../src/lib/combos/steps.ts");
const { buildPrecisionComboModelStep } = await import("../../src/lib/combos/builderDraft.ts");
function createLog() {
return { info() {}, warn() {}, debug() {}, error() {} };
}
function okResponse(content: string) {
return new Response(JSON.stringify({ choices: [{ message: { content } }] }), {
status: 200,
headers: { "content-type": "application/json" },
});
}
async function resetStorage() {
core.resetDbInstance();
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
}
async function seedConn(name: string, tags?: string[]) {
return providersDb.createProviderConnection({
provider: "openai",
authType: "apikey",
name,
apiKey: `sk-${name}`,
isActive: true,
...(tags ? { providerSpecificData: { tags } } : {}),
});
}
test.beforeEach(async () => {
await resetStorage();
});
test.after(async () => {
await resetStorage();
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
});
// ── 1. Schema parse ─────────────────────────────────────────────────────────
test("normalizeComboStep parses allowedConnectionIds (trim + drop empty)", () => {
const step = normalizeComboStep({
kind: "model",
model: "openai/gpt-4o-mini",
allowedConnectionIds: [" foo1 ", "foo2", "", " "],
});
assert.ok(step && step.kind === "model");
assert.deepEqual((step as { allowedConnectionIds?: string[] }).allowedConnectionIds, [
"foo1",
"foo2",
]);
});
test("normalizeComboStep omits allowedConnectionIds when absent or all-empty", () => {
const bare = normalizeComboStep({ kind: "model", model: "openai/gpt-4o-mini" });
assert.ok(bare && bare.kind === "model");
assert.equal((bare as { allowedConnectionIds?: string[] }).allowedConnectionIds, undefined);
const emptyish = normalizeComboStep({
kind: "model",
model: "openai/gpt-4o-mini",
allowedConnectionIds: ["", " "],
});
assert.equal((emptyish as { allowedConnectionIds?: string[] }).allowedConnectionIds, undefined);
});
test("buildPrecisionComboModelStep carries an allowlist when auto-selecting (#3266)", () => {
const step = buildPrecisionComboModelStep({
providerId: "openai",
modelId: "gpt-4o-mini",
connectionId: null,
allowedConnectionIds: [" a ", "b", "", "b"],
});
// trims, drops empty, dedupes
assert.deepEqual((step as { allowedConnectionIds?: string[] }).allowedConnectionIds, ["a", "b"]);
});
test("buildPrecisionComboModelStep drops the allowlist when a single account is pinned (#3266)", () => {
const step = buildPrecisionComboModelStep({
providerId: "openai",
modelId: "gpt-4o-mini",
connectionId: "pinned-1",
allowedConnectionIds: ["a", "b"],
});
// a forced connectionId wins — allowlist is meaningless and must not be carried
assert.equal(step.connectionId, "pinned-1");
assert.equal((step as { allowedConnectionIds?: string[] }).allowedConnectionIds, undefined);
});
// ── 2. Propagation onto the resolved combo target ───────────────────────────
test("handleComboChat propagates a step allowlist onto target.allowedConnectionIds (#3266)", async () => {
const foo1 = await seedConn("foo1");
const foo2 = await seedConn("foo2");
await seedConn("foo3");
await seedConn("foo4");
let captured: string[] | null = null;
const response = await handleComboChat({
body: { model: "rr", messages: [{ role: "user", content: "hi" }] },
combo: {
name: "rr",
strategy: "round-robin",
models: [
{
kind: "model",
model: "openai/gpt-4o-mini",
allowedConnectionIds: [foo1.id, foo2.id],
},
],
},
handleSingleModel: async (
_body: unknown,
modelStr: string,
target: { allowedConnectionIds?: unknown }
) => {
captured = Array.isArray(target?.allowedConnectionIds) ? target.allowedConnectionIds : null;
return okResponse(modelStr);
},
log: createLog(),
});
assert.equal(response.status, 200);
assert.ok(captured, "target.allowedConnectionIds must be populated from the step");
assert.deepEqual([...captured!].sort(), [foo1.id, foo2.id].sort());
});
test("buildAutoCandidates expands dynamic auto steps only within allowedConnectionIds", async () => {
const foo1 = await seedConn("foo1");
const foo2 = await seedConn("foo2");
const foo3 = await seedConn("foo3");
const foo4 = await seedConn("foo4");
const allowed = new Set([foo1.id, foo2.id]);
const forbidden = new Set([foo3.id, foo4.id]);
const candidates = await buildAutoCandidates(
[
{
kind: "model",
stepId: "openai/gpt-4o-mini",
executionKey: "openai/gpt-4o-mini",
modelStr: "openai/gpt-4o-mini",
provider: "openai",
providerId: "openai",
connectionId: null,
allowedConnectionIds: [foo1.id, foo2.id],
weight: 1,
label: null,
},
],
"auto-allowlist"
);
const connectionIds = candidates.map((candidate) => candidate.connectionId).filter(Boolean);
assert.deepEqual([...connectionIds].sort(), [foo1.id, foo2.id].sort());
assert.ok(connectionIds.every((connectionId) => allowed.has(connectionId!)));
assert.ok(connectionIds.every((connectionId) => !forbidden.has(connectionId!)));
});
// ── 3. Acceptance: the credential selector never escapes the allowlist ───────
test("getProviderCredentials never selects a connection outside the step allowlist (#3266)", async () => {
const foo1 = await seedConn("foo1");
const foo2 = await seedConn("foo2");
const foo3 = await seedConn("foo3");
const foo4 = await seedConn("foo4");
const allowed = [foo1.id, foo2.id];
const forbidden = new Set([foo3.id, foo4.id]);
const seen = new Set<string>();
for (let i = 0; i < 24; i++) {
const cred = (await getProviderCredentials("openai", null, allowed)) as {
connectionId?: string;
} | null;
assert.ok(cred && cred.connectionId, "a credential within the allowlist must be returned");
assert.ok(
allowed.includes(cred!.connectionId!),
`selected ${cred!.connectionId} which is outside the allowlist`
);
assert.ok(!forbidden.has(cred!.connectionId!), "must never select foo3/foo4");
seen.add(cred!.connectionId!);
}
// Both allowed accounts should be reachable (round-robin spreads across the subset).
assert.ok(seen.size >= 1 && [...seen].every((id) => allowed.includes(id)));
});
// ── 4. Step allowlist + tag routing → most-restrictive (intersection) ───────
test("a step allowlist intersects with tag routing — most-restrictive wins (#3266)", async () => {
await seedConn("foo1", ["us"]);
const foo2 = await seedConn("foo2", ["eu"]);
const foo3 = await seedConn("foo3", ["eu"]);
await seedConn("foo4", ["us"]);
let captured: string[] | null = null;
const response = await handleComboChat({
body: {
model: "rr",
messages: [{ role: "user", content: "hi" }],
metadata: { tags: ["eu"] },
},
combo: {
name: "rr",
strategy: "priority",
models: [
{
kind: "model",
model: "openai/gpt-4o-mini",
// allowlist excludes foo4; tags=[eu] match foo2+foo3 → intersection {foo2,foo3}
allowedConnectionIds: [foo2.id, foo3.id, "nonexistent-but-allowed"],
},
],
},
handleSingleModel: async (
_body: unknown,
modelStr: string,
target: { allowedConnectionIds?: unknown }
) => {
captured = Array.isArray(target?.allowedConnectionIds) ? target.allowedConnectionIds : null;
return okResponse(modelStr);
},
log: createLog(),
});
assert.equal(response.status, 200);
assert.ok(captured, "target.allowedConnectionIds must be the intersection");
assert.deepEqual([...captured!].sort(), [foo2.id, foo3.id].sort());
});