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
5 changes: 5 additions & 0 deletions src/app/api/v1/audio/speech/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv
import { parseSpeechModel } from "@omniroute/open-sse/config/audioRegistry.ts";
import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -41,6 +42,10 @@ export async function POST(request) {
return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model");
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, body.model);
if (policy.rejection) return policy.rejection;
Comment on lines +45 to +47

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Policy enforcement runs before parsing/normalizing the speech model. parseSpeechModel supports bare model IDs; if allowedModels is stored as provider/model (as in /v1/models), restricted keys may be incorrectly denied when clients omit the provider prefix. Consider normalizing (parseSpeechModel -> provider/model) before calling enforceApiKeyPolicy.

Copilot uses AI. Check for mistakes.

const { provider } = parseSpeechModel(body.model);
if (!provider) {
return errorResponse(
Expand Down
5 changes: 5 additions & 0 deletions src/app/api/v1/audio/transcriptions/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv
import { parseTranscriptionModel } from "@omniroute/open-sse/config/audioRegistry.ts";
import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -43,6 +44,10 @@ export async function POST(request) {
return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model");
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, model as string);
if (policy.rejection) return policy.rejection;
Comment on lines +47 to +49

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Policy enforcement runs before parsing/normalizing the transcription model. parseTranscriptionModel supports bare model IDs; if allowedModels is stored as provider/model, restricted keys may be incorrectly denied when clients omit the provider prefix. Consider normalizing (parseTranscriptionModel -> provider/model) before calling enforceApiKeyPolicy.

Copilot uses AI. Check for mistakes.

const { provider } = parseTranscriptionModel(model);
if (!provider) {
return errorResponse(
Expand Down
5 changes: 5 additions & 0 deletions src/app/api/v1/embeddings/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
import * as log from "@/sse/utils/logger";
import { toJsonErrorPayload } from "@/shared/utils/upstreamError";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -78,6 +79,10 @@ export async function POST(request) {
return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing input");
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, body.model);
if (policy.rejection) return policy.rejection;

Comment on lines +82 to +85

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Policy enforcement runs before parsing/normalizing the embedding model. Because parseEmbeddingModel supports bare model IDs (no provider prefix), a restricted key with allowedModels like "openai/text-embedding-3-large" will be denied if the client sends "text-embedding-3-large". Consider parsing first and passing a normalized provider/model string into enforceApiKeyPolicy.

Copilot uses AI. Check for mistakes.
// Parse model to get provider
const { provider } = parseEmbeddingModel(body.model);
if (!provider) {
Expand Down
5 changes: 5 additions & 0 deletions src/app/api/v1/images/generations/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
import * as log from "@/sse/utils/logger";
import { toJsonErrorPayload } from "@/shared/utils/upstreamError";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -75,6 +76,10 @@ export async function POST(request) {
return errorResponse(HTTP_STATUS.BAD_REQUEST, "Invalid prompt: expected a non-empty string");
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, body.model);
if (policy.rejection) return policy.rejection;

Comment on lines +79 to +82

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Policy enforcement runs before parsing/normalizing the image model. Since parseImageModel supports bare model IDs, a restricted key may be incorrectly denied when the client omits the provider prefix (while allowedModels typically uses provider/model). Consider normalizing (parseImageModel -> provider/model) before calling enforceApiKeyPolicy.

Copilot uses AI. Check for mistakes.
// Parse model to get provider
const { provider } = parseImageModel(body.model);
if (!provider) {
Expand Down
6 changes: 6 additions & 0 deletions src/app/api/v1/moderations/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv
import { parseModerationModel } from "@omniroute/open-sse/config/moderationRegistry.ts";
import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -38,6 +39,11 @@ export async function POST(request) {
}

const model = body.model || "omni-moderation-latest";

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, model);
if (policy.rejection) return policy.rejection;
Comment on lines 41 to +45

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Policy enforcement runs before parsing/normalizing the moderation model. parseModerationModel supports both "openai/model" and bare model IDs; /v1/models returns provider/model IDs. Consider building a canonical provider/model string (e.g., from parseModerationModel) before calling enforceApiKeyPolicy so allowedModels checks behave consistently.

Copilot uses AI. Check for mistakes.

const { provider } = parseModerationModel(model);

// Default to openai if no provider prefix
Expand Down
5 changes: 5 additions & 0 deletions src/app/api/v1/providers/[provider]/embeddings/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { getRegistryEntry } from "@omniroute/open-sse/config/providerRegistry.ts
import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/services/auth";
import { handleEmbedding } from "@omniroute/open-sse/handlers/embeddings.ts";
import * as log from "@/sse/utils/logger";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -53,6 +54,10 @@ export async function POST(request, { params }) {
body.model = `${providerAlias}/${body.model}`;
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, body.model);
if (policy.rejection) return policy.rejection;

// Validate provider match
if (body.model) {
const prefix = body.model.split("/")[0];
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv
import { getImageProvider } from "@omniroute/open-sse/config/imageRegistry.ts";
import * as log from "@/sse/utils/logger";
import { toJsonErrorPayload } from "@/shared/utils/upstreamError";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -60,6 +61,10 @@ export async function POST(request, { params }) {
body.model = `${rawProvider}/${body.model}`;
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, body.model);
if (policy.rejection) return policy.rejection;

// Validate provider match
const modelProvider = body.model.split("/")[0];
if (modelProvider !== rawProvider) {
Expand Down
5 changes: 5 additions & 0 deletions src/app/api/v1/rerank/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv
import { parseRerankModel } from "@omniroute/open-sse/config/rerankRegistry.ts";
import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy";

/**
* Handle CORS preflight
Expand Down Expand Up @@ -44,6 +45,10 @@ export async function POST(request) {
return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model");
}

// Enforce API key policies (model restrictions + budget limits)
const policy = await enforceApiKeyPolicy(request, body.model);
if (policy.rejection) return policy.rejection;
Comment on lines +48 to +50

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Policy enforcement runs before parsing/normalizing the rerank model. Since parseRerankModel supports bare model IDs, allowedModels entries stored as provider/model may not match when clients omit the provider prefix. Consider parsing first and passing a normalized provider/model string into enforceApiKeyPolicy.

Copilot uses AI. Check for mistakes.

const { provider } = parseRerankModel(body.model);
if (!provider) {
return errorResponse(
Expand Down
104 changes: 104 additions & 0 deletions src/shared/utils/apiKeyPolicy.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
/**
* API Key Policy Enforcement — Shared middleware for all /v1/* endpoints.
*
* Enforces API key policies: model restrictions and budget limits.
* Should be called after API key authentication in every endpoint that
* accepts a model parameter.
*
* @module shared/utils/apiKeyPolicy
*/

import { extractApiKey } from "@/sse/services/auth";
import { getApiKeyMetadata, isModelAllowedForKey } from "@/lib/localDb";
import { checkBudget } from "@/domain/costRules";
import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";

export interface ApiKeyPolicyResult {
/** API key string (null if no key provided) */
apiKey: string | null;
/** Metadata from DB (null if no key or key not found) */
apiKeyInfo: any | null;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The apiKeyInfo property is typed as any. To leverage TypeScript's type safety and improve code clarity, consider defining a specific interface for the API key metadata (e.g., ApiKeyMetadata) and using it here. This would help prevent potential runtime errors and make the code easier to understand and maintain.

/** If set, the request should be rejected with this Response */
rejection: Response | null;
}

/**
* Enforce API key policies for a request.
*
* Checks:
* 1. Model restriction — if the key has `allowedModels`, verify the requested model is permitted
* 2. Budget limit — if the key has a budget configured, verify it hasn't been exceeded
*
* @param request - The incoming HTTP request
* @param modelStr - The model ID from the request body
* @returns ApiKeyPolicyResult with apiKey, metadata, and optional rejection response
*
* @example
* ```ts
* const policy = await enforceApiKeyPolicy(request, body.model);
* if (policy.rejection) return policy.rejection;
* // proceed with request, optionally use policy.apiKeyInfo
* ```
*/
export async function enforceApiKeyPolicy(
request: Request,
modelStr: string | null
): Promise<ApiKeyPolicyResult> {
const apiKey = extractApiKey(request);
Comment on lines +44 to +48

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This new shared policy layer introduces security-critical behavior (model allowlist patterns + budget rejection) but currently has no automated tests. Adding unit tests covering exact match, prefix (provider/*), wildcard patterns, and budget-exceeded responses would help prevent regressions as more /v1 endpoints adopt this helper.

Copilot uses AI. Check for mistakes.

// No API key = local mode, skip policy checks
if (!apiKey) {
return { apiKey: null, apiKeyInfo: null, rejection: null };
}

// Fetch key metadata (includes allowedModels)
let apiKeyInfo: any = null;
try {
apiKeyInfo = await getApiKeyMetadata(apiKey);
} catch {
// If metadata fetch fails, don't block — degrade gracefully
return { apiKey, apiKeyInfo: null, rejection: null };
Comment on lines +59 to +61

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If getApiKeyMetadata() throws (DB unavailable/transient error), this helper currently fails open and skips both model and budget enforcement. That can reintroduce the original bypass during outages. Consider at least logging the error and/or failing closed (e.g., 503) when an API key is presented, or making fail-open vs fail-closed configurable.

Suggested change
} catch {
// If metadata fetch fails, don't block — degrade gracefully
return { apiKey, apiKeyInfo: null, rejection: null };
} catch (err) {
// If metadata fetch fails, fail closed to avoid bypassing policy checks
console.error("Failed to fetch API key metadata", err);
return {
apiKey,
apiKeyInfo: null,
rejection: errorResponse(
HTTP_STATUS.SERVICE_UNAVAILABLE,
"Unable to validate API key metadata. Please try again later."
),
};

Copilot uses AI. Check for mistakes.
}
Comment on lines +57 to +62

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The try...catch block for getApiKeyMetadata currently suppresses errors. While degrading gracefully is the correct behavior for the user, logging the error is crucial for maintainability. It would help you debug issues with the database or the metadata retrieval logic. You'll need to import the logger: import * as log from "@/sse/utils/logger";

Suggested change
try {
apiKeyInfo = await getApiKeyMetadata(apiKey);
} catch {
// If metadata fetch fails, don't block — degrade gracefully
return { apiKey, apiKeyInfo: null, rejection: null };
}
try {
apiKeyInfo = await getApiKeyMetadata(apiKey);
} catch (error) {
// If metadata fetch fails, don't block — degrade gracefully but log it
log.warn("API_POLICY", "Failed to fetch API key metadata. Policies will not be applied.", { error });
return { apiKey, apiKeyInfo: null, rejection: null };
}


// Key not found in DB — skip policy (auth layer handles validation)
if (!apiKeyInfo) {
return { apiKey, apiKeyInfo: null, rejection: null };
}

// ── Check 1: Model restriction ──
if (modelStr && apiKeyInfo.allowedModels && apiKeyInfo.allowedModels.length > 0) {
const allowed = await isModelAllowedForKey(apiKey, modelStr);

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

enforceApiKeyPolicy already fetched apiKeyInfo (including allowedModels), but isModelAllowedForKey() will fetch metadata again internally. Even with caching, this is duplicated work for restricted keys. Consider refactoring the model-allowed check to reuse the already-fetched allowedModels/metadata.

Suggested change
const allowed = await isModelAllowedForKey(apiKey, modelStr);
const allowed =
Array.isArray(apiKeyInfo.allowedModels) &&
apiKeyInfo.allowedModels.includes(modelStr);

Copilot uses AI. Check for mistakes.
if (!allowed) {
return {
apiKey,
apiKeyInfo,
rejection: errorResponse(
HTTP_STATUS.FORBIDDEN,
`Model "${modelStr}" is not allowed for this API key`
),
};
}
}

// ── Check 2: Budget limit ──
if (apiKeyInfo.id) {
try {
const budgetOk = checkBudget(apiKeyInfo.id);
if (!budgetOk.allowed) {
return {
apiKey,
apiKeyInfo,
rejection: errorResponse(
HTTP_STATUS.RATE_LIMITED,
budgetOk.reason || "Budget limit exceeded"
),
};
}
} catch {
// Budget check is best-effort — don't block on errors
}
Comment on lines +98 to +100

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the metadata fetch, this catch block suppresses errors from the budget check. To help with debugging and monitoring, it's best to log these errors, even if the request is allowed to proceed. This requires importing the logger if you haven't already: import * as log from "@/sse/utils/logger";

} catch (error) {
      // Budget check is best-effort — don't block on errors, but log them.
      log.warn("API_POLICY", "Budget check failed. Request will be allowed.", { error });
    }

}

return { apiKey, apiKeyInfo, rejection: null };
}
31 changes: 10 additions & 21 deletions src/sse/handlers/chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import {
} from "@omniroute/open-sse/utils/proxyFetch.ts";
import * as log from "../utils/logger";
import { updateProviderCredentials, checkAndRefreshToken } from "../services/tokenRefresh";
import { getSettings, getCombos, getApiKeyMetadata } from "@/lib/localDb";
import { getSettings, getCombos } from "@/lib/localDb";
import { resolveProxyForConnection } from "@/lib/localDb";
import { logProxyEvent } from "../../lib/proxyLogger";
import { logTranslationEvent } from "../../lib/translatorEvents";
Expand All @@ -34,8 +34,9 @@ import { getCircuitBreaker, CircuitBreakerOpenError } from "../../shared/utils/c
import { isModelAvailable, setModelUnavailable } from "../../domain/modelAvailability";
import { RequestTelemetry, recordTelemetry } from "../../shared/utils/requestTelemetry";
import { generateRequestId } from "../../shared/utils/requestId";
import { checkBudget, recordCost } from "../../domain/costRules";
import { recordCost } from "../../domain/costRules";
import { logAuditEvent } from "../../lib/compliance/index";
import { enforceApiKeyPolicy } from "../../shared/utils/apiKeyPolicy";

/**
* Handle chat completion request
Expand Down Expand Up @@ -97,15 +98,8 @@ export async function handleChat(request: any, clientRawRequest: any = null) {
// Log API key (masked)
const authHeader = request.headers.get("Authorization");
const apiKey = extractApiKey(request);
let apiKeyInfo = null;
if (authHeader && apiKey) {
const masked = log.maskKey(apiKey);
log.debug("AUTH", `API Key: ${masked}`);
try {
apiKeyInfo = await getApiKeyMetadata(apiKey);
} catch {
apiKeyInfo = null;
}
log.debug("AUTH", `API Key: ${log.maskKey(apiKey)}`);
} else {
log.debug("AUTH", "No API key provided (local mode)");
}
Expand All @@ -129,19 +123,14 @@ export async function handleChat(request: any, clientRawRequest: any = null) {
return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model");
}

// Pipeline: Budget check (if API key has budget limits)
// Pipeline: API key policy enforcement (model restrictions + budget limits)
telemetry.startPhase("policy");
if (apiKeyInfo?.id) {
try {
const budgetOk = checkBudget(apiKeyInfo.id);
if (!budgetOk.allowed) {
log.warn("BUDGET", `API key ${apiKeyInfo.id} exceeded budget: ${budgetOk.reason}`);
return errorResponse(429, budgetOk.reason || "Budget limit exceeded");
}
} catch {
// Budget check is best-effort — don't block on errors
}
const policy = await enforceApiKeyPolicy(request, modelStr);
if (policy.rejection) {
log.warn("POLICY", `API key policy rejected: ${modelStr} (key=${policy.apiKeyInfo?.id || "unknown"})`);
return policy.rejection;
}
const apiKeyInfo = policy.apiKeyInfo;
Comment on lines +126 to +133

Copilot AI Feb 25, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Model strings can be provided without a provider prefix (e.g. "whisper-1", "text-embedding-3-large"), but allowedModels are typically stored as provider/model (as returned by /v1/models). Since policy enforcement runs before model normalization, restricted keys may be incorrectly rejected (or require clients to change model format). Consider normalizing to a canonical provider/model string (via getModelInfo) before calling enforceApiKeyPolicy, while still allowing combo names to pass through.

Copilot uses AI. Check for mistakes.
telemetry.endPhase();

// Check if model is a combo (has multiple models with fallback)
Expand Down