Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/lib/db/apiKeys.ts
Original file line number Diff line number Diff line change
Expand Up @@ -368,6 +368,15 @@ async function getModelPermissionCandidates(modelId: string): Promise<string[]>
return Array.from(candidates);
}

export async function isModelBlockedByPatterns(
blockedModels: string[] | null | undefined,
modelId: string
): Promise<boolean> {
if (!blockedModels?.length) return false;
const candidates = await getModelPermissionCandidates(modelId);
return blockedModels.some((pattern) => modelPatternMatches(pattern, candidates));
}

async function getPublishedModelLookupTarget(
modelId: string
): Promise<{ providerId: string; modelId: string } | null> {
Expand Down
3 changes: 3 additions & 0 deletions src/shared/utils/apiKeyPolicy.ts
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ export interface ApiKeyMetadata {
name?: string;
modelAccessMode?: "all" | "restricted";
allowedModels?: string[];
blockedModels?: string[];
allowedCombos?: string[];
allowedConnections?: string[];
allowedQuotas?: string[];
Expand Down Expand Up @@ -319,6 +320,7 @@ async function validateStandardRoutingTarget(
const hasModelRestrictions =
apiKeyInfo.modelAccessMode === "restricted" ||
Boolean(apiKeyInfo.allowedModels?.length) ||
Boolean(apiKeyInfo.blockedModels?.length) ||
apiKeyInfo.disableNonPublicModels === true;
if (!requestedComboName && hasModelRestrictions && modelStr.startsWith("auto/")) {
requestedComboName = modelStr;
Expand Down Expand Up @@ -531,6 +533,7 @@ async function validateModelAccess(context: PolicyContext): Promise<Response | n
const hasModelRestrictions =
apiKeyInfo.modelAccessMode === "restricted" ||
Boolean(apiKeyInfo.allowedModels?.length) ||
Boolean(apiKeyInfo.blockedModels?.length) ||
apiKeyInfo.disableNonPublicModels === true;
if (!requestedComboName && hasModelRestrictions) {
if (modelStr.startsWith("auto/") || modelStr.startsWith("qtSd/")) {
Expand Down
9 changes: 8 additions & 1 deletion src/sse/handlers/chat/comboTargetKeyPolicy.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,11 @@
* inner target so #9057 holds.
*/

import { isModelBlockedByPatterns } from "@/lib/db/apiKeys";

export type ComboTargetKeyPolicyInfo = {
allowedModels?: string[] | null;
blockedModels?: string[] | null;
disableNonPublicModels?: boolean | null;
modelAccessMode?: string | null;
};
Expand Down Expand Up @@ -37,9 +40,13 @@ export async function comboTargetPassesKeyModelPolicy(opts: {
if (!apiKey || !apiKeyInfo) return true;

const hasModelRestrictions =
Boolean(apiKeyInfo.allowedModels?.length) || apiKeyInfo.disableNonPublicModels === true;
Boolean(apiKeyInfo.allowedModels?.length) ||
Boolean(apiKeyInfo.blockedModels?.length) ||
apiKeyInfo.disableNonPublicModels === true;
if (!hasModelRestrictions) return true;

if (await isModelBlockedByPatterns(apiKeyInfo.blockedModels, targetModelStr)) return false;

if (allowListCoversRequestedCombo(apiKeyInfo.allowedModels, requestedModelStr)) {
return true;
}
Expand Down
31 changes: 31 additions & 0 deletions tests/unit/api-key-policy.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -517,6 +517,37 @@ test("enforceApiKeyPolicy rejects disallowed models and exhausted budgets", asyn
assert.match(await readErrorMessage(overBudget.rejection), /Daily budget exceeded/);
});

test("enforceApiKeyPolicy applies blockedModels in all-access mode", async () => {
const key = await createKeyWithPolicy({
modelAccessMode: "all",
allowedModels: [],
blockedModels: ["gpt-6*", "*/gpt-6*"],
});
const policy = await loadPolicy("all-mode-blocked-models");

const blocked = await policy.enforceApiKeyPolicy(
makePolicyRequest(key.key),
"mbrouter/gpt-6-codex"
);
assert.equal(blocked.rejection.status, 403);

const allowed = await policy.enforceApiKeyPolicy(
makePolicyRequest(key.key),
"mbrouter/gpt-5.6-sol"
);
assert.equal(allowed.rejection, null);

const metadata = await apiKeysDb.getApiKeyMetadata(key.key);
assert.ok(metadata);
const rerouted = await policy.validateApiKeyRoutingTarget(
makePolicyRequest(key.key),
key.key,
metadata,
"gpt-6"
);
assert.equal(rerouted?.status, 403);
});

test("enforceApiKeyPolicy returns Anthropic error envelope for /v1/messages model denials", async () => {
const restrictedKey = await createKeyWithPolicy({
allowedModels: ["cc/*"],
Expand Down
47 changes: 47 additions & 0 deletions tests/unit/combo-restricted-key-target-policy-12886.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -77,3 +77,50 @@ test("#12886: unrestricted key skips the gate", async () => {
assert.equal(ok, true);
assert.equal(called, 0);
});

test("blockedModels still filters combo targets in all-access mode", async () => {
let called = 0;
const ok = await comboTargetPassesKeyModelPolicy({
apiKey: KEY,
apiKeyInfo: {
modelAccessMode: "all",
allowedModels: [],
blockedModels: ["deepseek/*"],
},
requestedModelStr: COMBO,
targetModelStr: INNER,
isModelAllowedForKey: async () => {
called += 1;
return false;
},
});
assert.equal(ok, false);
assert.equal(called, 0);
});

test("blockedModels takes precedence without disabling allowed combo targets", async () => {
const apiKeyInfo = {
modelAccessMode: "restricted",
allowedModels: [COMBO],
blockedModels: ["anthropic/*"],
};
const checker = allowListChecker([COMBO]);

const allowed = await comboTargetPassesKeyModelPolicy({
apiKey: KEY,
apiKeyInfo,
requestedModelStr: COMBO,
targetModelStr: INNER,
isModelAllowedForKey: checker,
});
const blocked = await comboTargetPassesKeyModelPolicy({
apiKey: KEY,
apiKeyInfo,
requestedModelStr: COMBO,
targetModelStr: OTHER,
isModelAllowedForKey: checker,
});

assert.equal(allowed, true);
assert.equal(blocked, false);
});