TypeScript 97.4%
SQL 1%
JavaScript 0.9%
CSS 0.6%
1import "server-only";2import { and, eq, desc } from "drizzle-orm";3import { z } from "zod";4import { getDb, customEndpoints, type CustomEndpoint } from "@/db";5import { decryptSecret, encryptSecret, keyHint } from "@/lib/crypto/keys";6import { newId } from "@/lib/ids";7import { writeAudit } from "@/lib/audit";8import { log } from "@/lib/log";9import { ApiError } from "@/lib/api";10import type { AIProviderAdapter, PolyModel } from "@/lib/ai/core/types";11import { createCustomEndpointAdapter, discoverCustomModels, mergeCustomModels, normalizeBaseUrl, normalizeModelsPath, parseCustomModelKey, toCustomPolyModel, NO_KEY_SENTINEL, type CustomEndpointConfig, type ManualModel } from "@/lib/ai/providers/custom";12import { assertEndpointUrlAllowed, checkEndpointUrlSync } from "./ssrf";1314/**15 * Custom OpenAI-compatible endpoints — CRUD, secret handling, validation/discovery and the16 * request-time resolver used by the chat and Arena services.17 *18 * Secrets: the API key and the custom headers (as one JSON document) are AES-256-GCM envelopes19 * produced by `encryptSecret` with AAD `"<userId>|custom:<endpointId>"`, so a row copied to20 * another user or another endpoint cannot be decrypted. Nothing returned to the browser ever21 * contains key material — `PublicEndpoint` carries a hint and header *names* only.22 */2324export const MANUAL_MODEL_SCHEMA = z.object({25 id: z.string().trim().min(1).max(200),26 displayName: z.string().trim().max(120).optional(),27 contextTokens: z.number().int().positive().max(100_000_000).optional(),28 vision: z.boolean().optional(),29 tools: z.boolean().optional(),30 reasoning: z.boolean().optional(),31});3233const HEADER_NAME_RE = /^[A-Za-z0-9-]{1,80}$/;34const headersSchema = z35 .record(z.string(), z.string().max(2_000))36 .refine((h) => Object.keys(h).length <= 12, "At most 12 headers")37 .refine((h) => Object.keys(h).every((k) => HEADER_NAME_RE.test(k)), "Header names may only contain letters, digits and dashes")38 .refine((h) => !Object.keys(h).some((k) => ["host", "content-length", "transfer-encoding", "connection"].includes(k.toLowerCase())), "That header is managed by the HTTP client");3940export const ENDPOINT_INPUT_SCHEMA = z.object({41 name: z.string().trim().min(1, "Give the endpoint a name").max(80),42 baseUrl: z.string().trim().min(1, "Base URL is required").max(500),43 /** `undefined` = keep, `null`/`""` = clear, string = replace. */44 apiKey: z.string().max(1_024).nullable().optional(),45 headers: headersSchema.nullable().optional(),46 modelsPath: z.string().trim().max(200).nullable().optional(),47 manualModels: z.array(MANUAL_MODEL_SCHEMA).max(200).optional(),48});49export type EndpointInput = z.infer<typeof ENDPOINT_INPUT_SCHEMA>;50export const ENDPOINT_PATCH_SCHEMA = ENDPOINT_INPUT_SCHEMA.partial();51export type EndpointPatch = z.infer<typeof ENDPOINT_PATCH_SCHEMA>;5253export interface PublicEndpoint {54 id: string;55 name: string;56 baseUrl: string;57 hasKey: boolean;58 keyHint: string | null;59 /** Header names only — values stay encrypted. */60 headerNames: string[];61 modelsPath: string;62 manualModels: ManualModel[];63 status: "unverified" | "valid" | "invalid" | "error";64 lastValidatedAt: string | null;65 lastValidationError: string | null;66 lastLatencyMs: number | null;67 modelsAvailable: number | null;68 discoveredModels: { id: string; ownedBy?: string }[];69 discoveredAt: string | null;70 /** The host is loopback/private (only reachable when PolyLLM runs next to it). */71 isPrivate: boolean;72 createdAt: string;73 updatedAt: string;74}7576export interface EndpointProbeResult {77 ok: boolean;78 latencyMs: number;79 modelsAvailable: number;80 error?: string;81 errorCode?: string;82 /** Discovery + manual models, merged. */83 models: PolyModel[];84}8586const ctx = (userId: string, id: string) => ({ userId, provider: `custom:${id}` });8788function headerNamesOf(row: CustomEndpoint, userId: string): string[] {89 const h = decryptHeaders(row, userId);90 return h ? Object.keys(h) : [];91}9293export function toPublicEndpoint(row: CustomEndpoint, userId: string): PublicEndpoint {94 return {95 id: row.id,96 name: row.name,97 baseUrl: row.baseUrl,98 hasKey: Boolean(row.encryptedKey),99 keyHint: row.keyHint,100 headerNames: headerNamesOf(row, userId),101 modelsPath: row.modelsPath,102 manualModels: row.manualModels ?? [],103 status: (row.status as PublicEndpoint["status"]) ?? "unverified",104 lastValidatedAt: row.lastValidatedAt?.toISOString() ?? null,105 lastValidationError: row.lastValidationError,106 lastLatencyMs: row.lastLatencyMs,107 modelsAvailable: row.modelsAvailable,108 discoveredModels: (row.discoveredModels ?? []).map((m) => ({ id: m.id, ownedBy: m.ownedBy })),109 discoveredAt: row.discoveredAt?.toISOString() ?? null,110 isPrivate: checkEndpointUrlSync(row.baseUrl, { allowPrivate: true }).isPrivate,111 createdAt: row.createdAt.toISOString(),112 updatedAt: row.updatedAt.toISOString(),113 };114}115116// ---------------------------------------------------------------------------117// Secrets118// ---------------------------------------------------------------------------119function decryptKey(row: CustomEndpoint, userId: string): string | null {120 if (!row.encryptedKey) return null;121 try {122 return decryptSecret(row.encryptedKey, ctx(userId, row.id));123 } catch (e) {124 log.error("endpoint key decrypt failed", { endpointId: row.id, error: (e as Error).message });125 return null;126 }127}128129function decryptHeaders(row: CustomEndpoint, userId: string): Record<string, string> | null {130 if (!row.encryptedHeaders) return null;131 try {132 const parsed = JSON.parse(decryptSecret(row.encryptedHeaders, ctx(userId, row.id))) as Record<string, string>;133 return parsed && typeof parsed === "object" ? parsed : null;134 } catch (e) {135 log.error("endpoint headers decrypt failed", { endpointId: row.id, error: (e as Error).message });136 return null;137 }138}139140/** Runtime configuration with decrypted secrets — only build right before a request. */141export function toConfig(row: CustomEndpoint, userId: string): CustomEndpointConfig {142 return {143 id: row.id,144 name: row.name,145 baseUrl: normalizeBaseUrl(row.baseUrl),146 apiKey: decryptKey(row, userId),147 headers: decryptHeaders(row, userId),148 modelsPath: row.modelsPath,149 manualModels: row.manualModels ?? [],150 };151}152153// ---------------------------------------------------------------------------154// CRUD155// ---------------------------------------------------------------------------156export async function listEndpoints(userId: string): Promise<PublicEndpoint[]> {157 const rows = await getDb().select().from(customEndpoints).where(eq(customEndpoints.userId, userId)).orderBy(desc(customEndpoints.createdAt));158 return rows.map((r) => toPublicEndpoint(r, userId));159}160161export async function getEndpointRow(userId: string, id: string): Promise<CustomEndpoint | null> {162 const [row] = await getDb()163 .select()164 .from(customEndpoints)165 .where(and(eq(customEndpoints.userId, userId), eq(customEndpoints.id, id)))166 .limit(1);167 return row ?? null;168}169170export async function getEndpoint(userId: string, id: string): Promise<PublicEndpoint | null> {171 const row = await getEndpointRow(userId, id);172 return row ? toPublicEndpoint(row, userId) : null;173}174175export const MAX_ENDPOINTS_PER_USER = 20;176177async function validateBaseUrl(raw: string): Promise<string> {178 const check = await assertEndpointUrlAllowed(raw);179 if (!check.ok) throw new ApiError(400, check.message ?? "Invalid endpoint URL", check.reason ?? "INVALID_URL");180 return normalizeBaseUrl(raw);181}182183function cleanKey(k: string | null | undefined): string | null {184 const t = (k ?? "").trim();185 if (!t) return null;186 if (/\s/.test(t)) throw new ApiError(400, "API keys cannot contain whitespace.", "INVALID_KEY");187 return t;188}189190export async function createEndpoint(userId: string, input: EndpointInput, ip?: string | null): Promise<PublicEndpoint> {191 const db = getDb();192 const existing = await db.select({ id: customEndpoints.id }).from(customEndpoints).where(eq(customEndpoints.userId, userId));193 if (existing.length >= MAX_ENDPOINTS_PER_USER) throw new ApiError(400, `You can configure up to ${MAX_ENDPOINTS_PER_USER} endpoints.`, "LIMIT_REACHED");194 const baseUrl = await validateBaseUrl(input.baseUrl);195 const id = newId("cep");196 const key = cleanKey(input.apiKey);197 const headers = input.headers && Object.keys(input.headers).length ? input.headers : null;198 const now = new Date();199 const [row] = await db200 .insert(customEndpoints)201 .values({202 id,203 userId,204 name: input.name,205 baseUrl,206 encryptedKey: key ? encryptSecret(key, ctx(userId, id)) : null,207 keyHint: key ? keyHint(key) : null,208 encryptedHeaders: headers ? encryptSecret(JSON.stringify(headers), ctx(userId, id)) : null,209 modelsPath: normalizeModelsPath(input.modelsPath === undefined ? "/models" : input.modelsPath),210 manualModels: dedupeManual(input.manualModels ?? []),211 status: "unverified",212 createdAt: now,213 updatedAt: now,214 })215 .returning();216 await writeAudit({ userId, action: "endpoint.created", ipAddress: ip, meta: { endpointId: id, host: safeHost(baseUrl) } });217 return toPublicEndpoint(row, userId);218}219220export async function updateEndpoint(userId: string, id: string, patch: EndpointPatch, ip?: string | null): Promise<PublicEndpoint> {221 const row = await getEndpointRow(userId, id);222 if (!row) throw new ApiError(404, "Endpoint not found", "NOT_FOUND");223 const set: Partial<typeof customEndpoints.$inferInsert> = { updatedAt: new Date() };224 let invalidate = false;225 if (patch.name !== undefined) set.name = patch.name;226 if (patch.baseUrl !== undefined && normalizeBaseUrl(patch.baseUrl) !== row.baseUrl) {227 set.baseUrl = await validateBaseUrl(patch.baseUrl);228 invalidate = true;229 }230 if (patch.apiKey !== undefined) {231 const key = cleanKey(patch.apiKey);232 set.encryptedKey = key ? encryptSecret(key, ctx(userId, id)) : null;233 set.keyHint = key ? keyHint(key) : null;234 invalidate = true;235 }236 if (patch.headers !== undefined) {237 const headers = patch.headers && Object.keys(patch.headers).length ? patch.headers : null;238 set.encryptedHeaders = headers ? encryptSecret(JSON.stringify(headers), ctx(userId, id)) : null;239 invalidate = true;240 }241 if (patch.modelsPath !== undefined) {242 set.modelsPath = normalizeModelsPath(patch.modelsPath);243 invalidate = true;244 }245 if (patch.manualModels !== undefined) set.manualModels = dedupeManual(patch.manualModels);246 if (invalidate) {247 set.status = "unverified";248 set.lastValidationError = null;249 }250 const [updated] = await getDb().update(customEndpoints).set(set).where(eq(customEndpoints.id, id)).returning();251 await writeAudit({ userId, action: "endpoint.updated", ipAddress: ip, meta: { endpointId: id, fields: Object.keys(patch) } });252 return toPublicEndpoint(updated, userId);253}254255export async function deleteEndpoint(userId: string, id: string, ip?: string | null): Promise<void> {256 const res = await getDb()257 .delete(customEndpoints)258 .where(and(eq(customEndpoints.userId, userId), eq(customEndpoints.id, id)))259 .returning({ id: customEndpoints.id });260 if (!res.length) throw new ApiError(404, "Endpoint not found", "NOT_FOUND");261 await writeAudit({ userId, action: "endpoint.deleted", ipAddress: ip, meta: { endpointId: id } });262}263264function dedupeManual(list: ManualModel[]): ManualModel[] {265 const seen = new Set<string>();266 const out: ManualModel[] = [];267 for (const m of list) {268 const id = m.id.trim();269 if (!id || seen.has(id)) continue;270 seen.add(id);271 out.push({ ...m, id, displayName: m.displayName?.trim() || undefined });272 }273 return out;274}275276function safeHost(url: string): string {277 try {278 return new URL(url).host;279 } catch {280 return "?";281 }282}283284// ---------------------------------------------------------------------------285// Validation / discovery286// ---------------------------------------------------------------------------287/**288 * Probe the endpoint: `GET {baseUrl}{modelsPath}` when discovery is enabled (latency = that call);289 * otherwise a 1-token chat completion against the first manual model. Persists status, latency,290 * the discovery snapshot and the error (if any).291 */292export async function probeEndpoint(userId: string, id: string): Promise<EndpointProbeResult> {293 const row = await getEndpointRow(userId, id);294 if (!row) throw new ApiError(404, "Endpoint not found", "NOT_FOUND");295 const cfg = toConfig(row, userId);296 const t0 = Date.now();297 const now = new Date();298 const db = getDb();299 const urlCheck = await assertEndpointUrlAllowed(row.baseUrl);300 if (!urlCheck.ok) {301 await db.update(customEndpoints).set({ status: "invalid", lastValidatedAt: now, lastValidationError: urlCheck.message ?? "Blocked URL", updatedAt: now }).where(eq(customEndpoints.id, id));302 return { ok: false, latencyMs: 0, modelsAvailable: 0, error: urlCheck.message, errorCode: urlCheck.reason, models: [] };303 }304 try {305 let models: PolyModel[];306 let discovered: { id: string; ownedBy?: string; created?: number }[] | null = null;307 if (normalizeModelsPath(cfg.modelsPath)) {308 const found = await discoverCustomModels(cfg);309 discovered = found.map((m) => ({ id: m.id, ownedBy: typeof m.owned_by === "string" ? m.owned_by : undefined, created: typeof m.created === "number" ? m.created : undefined }));310 models = mergeCustomModels(cfg, found, cfg.manualModels ?? []);311 } else {312 const first = cfg.manualModels?.[0];313 if (!first) throw new ApiError(400, "Discovery is disabled and no manual model is declared — add a models path or at least one model id.", "NO_MODELS");314 const adapter = createCustomEndpointAdapter(cfg);315 const res = await adapter.chat({ provider: "custom", model: first.id, apiKey: cfg.apiKey ?? NO_KEY_SENTINEL, messages: [{ role: "user", content: [{ type: "text", text: "ping" }] }], settings: { maxTokens: 1 }, timeoutMs: 20_000 });316 if (res.finishReason === "error") throw new Error("The endpoint answered but the completion failed.");317 models = mergeCustomModels(cfg, [], cfg.manualModels ?? []);318 }319 const latencyMs = Date.now() - t0;320 await db321 .update(customEndpoints)322 .set({323 status: "valid",324 lastValidatedAt: now,325 lastValidationError: null,326 lastLatencyMs: latencyMs,327 modelsAvailable: models.length,328 ...(discovered ? { discoveredModels: discovered, discoveredAt: now } : {}),329 updatedAt: now,330 })331 .where(eq(customEndpoints.id, id));332 return { ok: true, latencyMs, modelsAvailable: models.length, models };333 } catch (e) {334 const latencyMs = Date.now() - t0;335 const err = e as { message?: string; code?: string; status?: number; name?: string };336 const message = friendlyProbeError(err);337 await db.update(customEndpoints).set({ status: "invalid", lastValidatedAt: now, lastValidationError: message.slice(0, 500), lastLatencyMs: latencyMs, updatedAt: now }).where(eq(customEndpoints.id, id));338 log.warn("endpoint probe failed", { endpointId: id, host: safeHost(row.baseUrl), error: message, code: err.code });339 return { ok: false, latencyMs, modelsAvailable: 0, error: message, errorCode: err.code ?? (err.status ? `HTTP_${err.status}` : "PROBE_FAILED"), models: [] };340 }341}342343function friendlyProbeError(err: { message?: string; code?: string; status?: number; name?: string; cause?: unknown }): string {344 const cause = (err as { cause?: { code?: string } }).cause;345 const code = err.code ?? cause?.code;346 if (err.name === "AbortError" || code === "UND_ERR_CONNECT_TIMEOUT" || code === "REQUEST_TIMEOUT") return "The endpoint did not answer within the timeout.";347 if (code === "ECONNREFUSED") return "Connection refused — nothing is listening at that address (from the PolyLLM server).";348 if (code === "ENOTFOUND" || code === "EAI_AGAIN") return "Host name could not be resolved.";349 if (code === "ECONNRESET" || code === "UND_ERR_SOCKET") return "The connection was reset by the endpoint.";350 if (code === "CERT_HAS_EXPIRED" || code === "DEPTH_ZERO_SELF_SIGNED_CERT" || code === "SELF_SIGNED_CERT_IN_CHAIN" || code === "UNABLE_TO_VERIFY_LEAF_SIGNATURE") return "TLS certificate could not be verified.";351 if (err.status === 401 || err.status === 403 || code === "INVALID_API_KEY") return "The endpoint rejected the credentials (401/403).";352 if (err.status === 404) return "404 — check the base URL (it should end with /v1) and the models path.";353 return err.message?.replace(/^fetch failed:?\s*/i, "").trim() || "Could not reach the endpoint.";354}355356// ---------------------------------------------------------------------------357// Registry surface + request-time resolution358// ---------------------------------------------------------------------------359/** Models from every endpoint of the user (discovery snapshot + manual declarations). Never calls the endpoints. */360export async function listCustomModels(userId: string): Promise<PolyModel[]> {361 const rows = await getDb().select().from(customEndpoints).where(eq(customEndpoints.userId, userId)).orderBy(desc(customEndpoints.createdAt));362 const out: PolyModel[] = [];363 for (const row of rows) {364 const cfg = { id: row.id, name: row.name };365 const merged = mergeCustomModels(366 cfg,367 (row.discoveredModels ?? []).map((m) => ({ id: m.id, owned_by: m.ownedBy, created: m.created })),368 row.manualModels ?? [],369 );370 for (const m of merged) out.push({ ...m, metadata: { ...m.metadata, endpointStatus: row.status, endpointBaseUrl: row.baseUrl } });371 }372 return out;373}374375export async function countValidEndpoints(userId: string): Promise<number> {376 const rows = await getDb().select({ id: customEndpoints.id }).from(customEndpoints).where(and(eq(customEndpoints.userId, userId), eq(customEndpoints.status, "valid")));377 return rows.length;378}379380export interface ResolvedCustomEndpoint {381 adapter: AIProviderAdapter;382 /** Decrypted key, or `NO_KEY_SENTINEL` when the endpoint has none (adapters need a non-empty bearer). */383 apiKey: string;384 model: PolyModel;385 endpoint: PublicEndpoint;386}387388/**389 * Resolve a `custom/<endpointId>:<modelId>` key for `userId` into a ready-to-use adapter, key and model.390 * Returns `null` when the key is not a custom key or the endpoint does not belong to the user.391 *392 * Integration (owned by the chat/arena services):393 * const custom = isCustomModelKey(modelKey) ? await resolveCustomEndpoint(userId, modelKey) : null;394 * const model = custom?.model ?? (await getModel(modelKey));395 * const apiKey = custom?.apiKey ?? (await getDecryptedKey(userId, model.provider));396 * const adapter = custom?.adapter ?? getAdapter(model.provider);397 */398export async function resolveCustomEndpoint(userId: string, modelKey: string): Promise<ResolvedCustomEndpoint | null> {399 const parsed = parseCustomModelKey(modelKey);400 if (!parsed) return null;401 const row = await getEndpointRow(userId, parsed.endpointId);402 if (!row) return null;403 const cfg = toConfig(row, userId);404 const manual = (row.manualModels ?? []).find((m) => m.id.trim() === parsed.modelId);405 const disc = (row.discoveredModels ?? []).find((m) => m.id === parsed.modelId);406 const model = toCustomPolyModel(cfg, parsed.modelId, manual, { ownedBy: disc?.ownedBy, created: disc?.created, source: disc ? "discovery" : "manual" });407 return { adapter: createCustomEndpointAdapter(cfg), apiKey: cfg.apiKey ?? NO_KEY_SENTINEL, model, endpoint: toPublicEndpoint(row, userId) };408}409410export { isCustomModelKey, parseCustomModelKey } from "@/lib/ai/providers/custom";411