diff --git a/server/lib/blueprints/aiProviders.ts b/server/lib/blueprints/aiProviders.ts new file mode 100644 index 000000000..5dbf18bf0 --- /dev/null +++ b/server/lib/blueprints/aiProviders.ts @@ -0,0 +1,286 @@ +import { and, eq, inArray } from "drizzle-orm"; +import { + aiModels, + aiProviders, + resourceAiModels, + siteResourceAiModels, + Transaction +} from "@server/db"; +import { + AccessMode, + ModelListType, + clearPublicResourceAiConfig, + clearSiteResourceAiConfig, + isInferenceFieldsError, + resolveProviderAttachments, + setPublicResourceAiProviders, + setSiteResourceAiProviders +} from "@server/lib/aiInferenceResource"; + +export type BlueprintAiModelInput = { + model: string; + listType: ModelListType; +}; + +export type BlueprintAiProviderInput = { + provider: string; + accessMode: AccessMode; + enabled: boolean; + models: BlueprintAiModelInput[]; +}; + +async function resolveProviderNiceIds( + orgId: string, + niceIds: string[], + trx: Transaction +): Promise> { + const unique = [...new Set(niceIds)]; + if (unique.length === 0) { + return new Map(); + } + + const rows = await trx + .select({ + providerId: aiProviders.providerId, + niceId: aiProviders.niceId + }) + .from(aiProviders) + .where( + and( + eq(aiProviders.orgId, orgId), + inArray(aiProviders.niceId, unique) + ) + ); + + const byNiceId = new Map(rows.map((r) => [r.niceId, r.providerId])); + const missing = unique.filter((id) => !byNiceId.has(id)); + if (missing.length > 0) { + throw new Error( + `AI provider(s) not found in this org: ${missing.join(", ")}` + ); + } + + return byNiceId; +} + +async function resolveModelKeys( + providers: BlueprintAiProviderInput[], + providerIdByNiceId: Map, + trx: Transaction +): Promise> { + const providerIds = [ + ...new Set( + providers + .filter((p) => p.models.length > 0) + .map((p) => providerIdByNiceId.get(p.provider)!) + ) + ]; + + if (providerIds.length === 0) { + return new Map(); + } + + const rows = await trx + .select({ + modelId: aiModels.modelId, + modelKey: aiModels.modelKey, + providerId: aiModels.providerId + }) + .from(aiModels) + .where(inArray(aiModels.providerId, providerIds)); + + const byProviderAndKey = new Map(); + for (const row of rows) { + byProviderAndKey.set(`${row.providerId}::${row.modelKey}`, row.modelId); + } + + const modelIdByEntryKey = new Map(); + const missing: string[] = []; + for (const provider of providers) { + const providerId = providerIdByNiceId.get(provider.provider)!; + for (const m of provider.models) { + const modelId = byProviderAndKey.get(`${providerId}::${m.model}`); + if (modelId === undefined) { + missing.push(`${provider.provider}/${m.model}`); + continue; + } + modelIdByEntryKey.set(`${provider.provider}::${m.model}`, modelId); + } + } + + if (missing.length > 0) { + throw new Error(`AI model(s) not found: ${missing.join(", ")}`); + } + + return modelIdByEntryKey; +} + +async function validateModelEntries(input: { + orgId: string; + entries: { modelId: number; listType: ModelListType }[]; + selectProviderIds: number[]; + trx: Transaction; +}): Promise { + if (input.entries.length === 0) { + return; + } + + if (input.selectProviderIds.length === 0) { + throw new Error( + "Set at least one attached AI provider to access-mode 'select' before declaring models" + ); + } + + const modelIds = input.entries.map((e) => e.modelId); + const catalogRows = await input.trx + .select({ + modelId: aiModels.modelId, + listType: aiModels.listType, + providerId: aiModels.providerId, + enabled: aiModels.enabled + }) + .from(aiModels) + .innerJoin(aiProviders, eq(aiModels.providerId, aiProviders.providerId)) + .where( + and( + inArray(aiModels.modelId, modelIds), + inArray(aiModels.providerId, input.selectProviderIds), + eq(aiProviders.orgId, input.orgId) + ) + ); + + const catalogById = new Map(catalogRows.map((row) => [row.modelId, row])); + for (const entry of input.entries) { + const catalog = catalogById.get(entry.modelId); + if (!catalog) { + throw new Error( + `Model ${entry.modelId} does not exist or does not belong to a select-mode attached provider` + ); + } + if (catalog.listType !== entry.listType) { + throw new Error( + `Model ${entry.modelId} must use list-type "${catalog.listType}" to match the provider catalog entry` + ); + } + if (!catalog.enabled) { + throw new Error(`Model ${entry.modelId} is disabled on its provider`); + } + } +} + +type SyncInferenceAiConfigInput = { + orgId: string; + trx: Transaction; + mode: string; + providers: BlueprintAiProviderInput[]; +} & ( + | { scope: "public"; resourceId: number } + | { scope: "site"; siteResourceId: number } +); + +/** + * Fully declarative: makes the resource's attached AI providers/models match + * exactly what the blueprint declares (omitted providers/models are removed). + * Non-inference resources have any leftover AI config cleared. + */ +export async function syncInferenceAiConfig( + input: SyncInferenceAiConfigInput +): Promise { + const { orgId, trx, mode } = input; + + if (mode !== "inference") { + if (input.scope === "public") { + await clearPublicResourceAiConfig(input.resourceId, trx); + } else { + await clearSiteResourceAiConfig(input.siteResourceId, trx); + } + return; + } + + const providerIdByNiceId = await resolveProviderNiceIds( + orgId, + input.providers.map((p) => p.provider), + trx + ); + + const resolvedAttachments = await resolveProviderAttachments({ + orgId, + attachments: input.providers.map((p) => ({ + providerId: providerIdByNiceId.get(p.provider)!, + accessMode: p.accessMode, + enabled: p.enabled + })), + requireAtLeastOne: false + }); + if (isInferenceFieldsError(resolvedAttachments)) { + throw new Error(resolvedAttachments.error); + } + + if (input.scope === "public") { + await setPublicResourceAiProviders( + input.resourceId, + resolvedAttachments, + trx + ); + } else { + await setSiteResourceAiProviders( + input.siteResourceId, + resolvedAttachments, + trx + ); + } + + const modelIdByEntryKey = await resolveModelKeys( + input.providers, + providerIdByNiceId, + trx + ); + + const modelEntries = input.providers.flatMap((p) => + p.models.map((m) => ({ + modelId: modelIdByEntryKey.get(`${p.provider}::${m.model}`)!, + listType: m.listType + })) + ); + + const selectProviderIds = resolvedAttachments + .filter((a) => a.accessMode === "select") + .map((a) => a.providerId); + + await validateModelEntries({ + orgId, + entries: modelEntries, + selectProviderIds, + trx + }); + + if (input.scope === "public") { + await trx + .delete(resourceAiModels) + .where(eq(resourceAiModels.resourceId, input.resourceId)); + if (modelEntries.length > 0) { + await trx.insert(resourceAiModels).values( + modelEntries.map((m) => ({ + resourceId: input.resourceId, + modelId: m.modelId, + listType: m.listType + })) + ); + } + } else { + await trx + .delete(siteResourceAiModels) + .where( + eq(siteResourceAiModels.siteResourceId, input.siteResourceId) + ); + if (modelEntries.length > 0) { + await trx.insert(siteResourceAiModels).values( + modelEntries.map((m) => ({ + siteResourceId: input.siteResourceId, + modelId: m.modelId, + listType: m.listType + })) + ); + } + } +} diff --git a/server/lib/blueprints/privateResources.ts b/server/lib/blueprints/privateResources.ts index 2991a5a83..070ac68b4 100644 --- a/server/lib/blueprints/privateResources.ts +++ b/server/lib/blueprints/privateResources.ts @@ -29,11 +29,13 @@ import { tierMatrix } from "../billing/tierMatrix"; import { build } from "@server/build"; import { LimitId } from "../billing"; import { usageService } from "../billing/usageService"; +import { syncInferenceAiConfig } from "./aiProviders"; async function getDomainForSiteResource( siteResourceId: number | undefined, fullDomain: string, orgId: string, + isInference: boolean, trx: Transaction ): Promise<{ subdomain: string | null; domainId: string }> { const [fullDomainExists] = await trx @@ -43,6 +45,9 @@ async function getDomainForSiteResource( and( eq(siteResources.fullDomain, fullDomain), eq(siteResources.orgId, orgId), + // exclude looking at the ones on exit nodes if this is an inference resource, + // and vice versa, so inference and non-inference resources can share a full-domain + ne(siteResources.requiresExitNodeConnection, !isInference), siteResourceId ? ne(siteResources.siteResourceId, siteResourceId) : isNotNull(siteResources.siteResourceId) @@ -214,7 +219,7 @@ export async function updatePrivateResources( resourceStatusFromSite = siteSingle.status ?? "approved"; } - if (allSites.length === 0) { + if (resourceData.mode !== "inference" && allSites.length === 0) { throw new Error( `No valid sites found for private private resource ${resourceNiceId} in org ${orgId}` ); @@ -231,11 +236,16 @@ export async function updatePrivateResources( let domainInfo: | { subdomain: string | null; domainId: string } | undefined; - if (resourceData["full-domain"] && resourceData.mode === "http") { + if ( + resourceData["full-domain"] && + (resourceData.mode === "http" || + resourceData.mode === "inference") + ) { domainInfo = await getDomainForSiteResource( existingResource.siteResourceId, resourceData["full-domain"], orgId, + resourceData.mode === "inference", trx ); } @@ -265,6 +275,8 @@ export async function updatePrivateResources( } } + const isInference = resourceData.mode === "inference"; + // Update existing resource const [updatedResource] = await trx .update(siteResources) @@ -279,13 +291,15 @@ export async function updatePrivateResources( alias: resourceData.alias || null, disableIcmp: resourceData["disable-icmp"] || - (resourceData.mode == "http" ? true : false), // default to true for http resources, otherwise false + (resourceData.mode == "http" || isInference + ? true + : false), // default to true for http/inference resources, otherwise false tcpPortRangeString: - resourceData.mode == "http" + resourceData.mode == "http" || isInference ? "443,80" : resourceData["tcp-ports"], udpPortRangeString: - resourceData.mode == "http" + resourceData.mode == "http" || isInference ? "" : resourceData["udp-ports"], fullDomain: resourceData["full-domain"] || null, @@ -295,7 +309,9 @@ export async function updatePrivateResources( authDaemonMode: resourceData["auth-daemon"]?.mode || "native", authDaemonPort: resourceData["auth-daemon"]?.port || 22123, - status: resourceStatusFromSite + status: resourceStatusFromSite, + networkId: isInference ? null : undefined, + requiresExitNodeConnection: isInference }) .where( eq( @@ -307,7 +323,19 @@ export async function updatePrivateResources( const siteResourceId = existingResource.siteResourceId; - if (updatedResource.networkId) { + if (isInference) { + // inference resources are not attached to any site network + if (existingResource.networkId) { + await trx + .delete(siteNetworks) + .where( + eq( + siteNetworks.networkId, + existingResource.networkId + ) + ); + } + } else if (updatedResource.networkId) { await trx .delete(siteNetworks) .where( @@ -322,6 +350,23 @@ export async function updatePrivateResources( } } + await syncInferenceAiConfig({ + orgId, + trx, + mode: resourceData.mode, + scope: "site", + siteResourceId, + providers: resourceData["ai-providers"].map((p) => ({ + provider: p.provider, + accessMode: p["access-mode"], + enabled: p.enabled, + models: p.models.map((m) => ({ + model: m.model, + listType: m["list-type"] + })) + })) + }); + await trx .delete(clientSiteResources) .where(eq(clientSiteResources.siteResourceId, siteResourceId)); @@ -501,14 +546,20 @@ export async function updatePrivateResources( releaseAliasLock = release; } + const isInference = resourceData.mode === "inference"; + let domainInfo: | { subdomain: string | null; domainId: string } | undefined; - if (resourceData["full-domain"] && resourceData.mode === "http") { + if ( + resourceData["full-domain"] && + (resourceData.mode === "http" || isInference) + ) { domainInfo = await getDomainForSiteResource( undefined, resourceData["full-domain"], orgId, + isInference, trx ); } @@ -534,13 +585,16 @@ export async function updatePrivateResources( } } - const [network] = await trx - .insert(networks) - .values({ - scope: "resource", - orgId: orgId - }) - .returning(); + let network: typeof networks.$inferSelect | undefined; + if (!isInference) { + [network] = await trx + .insert(networks) + .values({ + scope: "resource", + orgId: orgId + }) + .returning(); + } // Create new resource const [newResource] = await trx @@ -548,8 +602,8 @@ export async function updatePrivateResources( .values({ orgId: orgId, niceId: resourceNiceId, - networkId: network.networkId, - defaultNetworkId: network.networkId, + networkId: network ? network.networkId : null, + defaultNetworkId: network ? network.networkId : null, name: resourceData.name || resourceNiceId, mode: resourceData.mode, ssl: resourceData.ssl, @@ -561,13 +615,15 @@ export async function updatePrivateResources( aliasAddress: aliasAddress, disableIcmp: resourceData["disable-icmp"] || - (resourceData.mode == "http" ? true : false), // default to true for http resources, otherwise false + (resourceData.mode == "http" || isInference + ? true + : false), // default to true for http/inference resources, otherwise false tcpPortRangeString: - resourceData.mode == "http" + resourceData.mode == "http" || isInference ? "443,80" : resourceData["tcp-ports"], udpPortRangeString: - resourceData.mode == "http" + resourceData.mode == "http" || isInference ? "" : resourceData["udp-ports"], fullDomain: resourceData["full-domain"] || null, @@ -577,7 +633,8 @@ export async function updatePrivateResources( authDaemonMode: resourceData["auth-daemon"]?.mode || "native", authDaemonPort: resourceData["auth-daemon"]?.port || 22123, - status: resourceStatusFromSite + status: resourceStatusFromSite, + requiresExitNodeConnection: isInference }) .returning(); @@ -585,13 +642,32 @@ export async function updatePrivateResources( const siteResourceId = newResource.siteResourceId; - for (const site of allSites) { - await trx.insert(siteNetworks).values({ - siteId: site.siteId, - networkId: network.networkId - }); + if (network) { + for (const site of allSites) { + await trx.insert(siteNetworks).values({ + siteId: site.siteId, + networkId: network.networkId + }); + } } + await syncInferenceAiConfig({ + orgId, + trx, + mode: resourceData.mode, + scope: "site", + siteResourceId, + providers: resourceData["ai-providers"].map((p) => ({ + provider: p.provider, + accessMode: p["access-mode"], + enabled: p.enabled, + models: p.models.map((m) => ({ + model: m.model, + listType: m["list-type"] + })) + })) + }); + const [adminRole] = await trx .select() .from(roles) diff --git a/server/lib/blueprints/publicResources.ts b/server/lib/blueprints/publicResources.ts index f21bf0f26..9c9242984 100644 --- a/server/lib/blueprints/publicResources.ts +++ b/server/lib/blueprints/publicResources.ts @@ -56,6 +56,7 @@ import createHttpError from "http-errors"; import next from "next"; import { LimitId } from "../billing"; import { usageService } from "../billing/usageService"; +import { syncInferenceAiConfig } from "./aiProviders"; export type PublicResourcesResults = { proxyResource: Resource; @@ -287,7 +288,7 @@ export async function updatePublicResources( if (existingResource) { let domain; if ( - ["http", "ssh", "rdp", "vnc"].includes(resourceData.mode || "") + ["http", "ssh", "rdp", "vnc", "inference"].includes(resourceData.mode || "") ) { if (resourceData["full-domain"]?.startsWith("*.")) { const isLicensed = await isLicensedOrSubscribed( @@ -371,12 +372,12 @@ export async function updatePublicResources( name: resourceData.name || "Unnamed Resource", mode: resourceData.mode, - proxyPort: ["http", "ssh", "rdp", "vnc"].includes( + proxyPort: ["http", "ssh", "rdp", "vnc", "inference"].includes( resourceData.mode || "" ) ? null : resourceData["proxy-port"], - fullDomain: ["http", "ssh", "rdp", "vnc"].includes( + fullDomain: ["http", "ssh", "rdp", "vnc", "inference"].includes( resourceData.mode || "" ) ? resourceData["full-domain"] @@ -567,12 +568,13 @@ export async function updatePublicResources( .update(resources) .set({ name: resourceData.name || "Unnamed Resource", - proxyPort: ["http", "ssh", "rdp", "vnc"].includes( + mode: resourceData.mode, + proxyPort: ["http", "ssh", "rdp", "vnc", "inference"].includes( resourceData.mode || "" ) ? null : resourceData["proxy-port"], - fullDomain: ["http", "ssh", "rdp", "vnc"].includes( + fullDomain: ["http", "ssh", "rdp", "vnc", "inference"].includes( resourceData.mode || "" ) ? resourceData["full-domain"] @@ -674,6 +676,25 @@ export async function updatePublicResources( trx ); } + + await syncInferenceAiConfig({ + orgId, + trx, + mode: resourceData.mode || "", + scope: "public", + resourceId: existingResource.resourceId, + providers: (resourceData["ai-providers"] || []).map( + (p) => ({ + provider: p.provider, + accessMode: p["access-mode"], + enabled: p.enabled, + models: p.models.map((m) => ({ + model: m.model, + listType: m["list-type"] + })) + }) + ) + }); } const existingResourceTargets = await trx @@ -754,7 +775,7 @@ export async function updatePublicResources( : undefined), rewritePathType: targetData["rewrite-match"], priority: targetData.priority, - mode: resourceData.mode + mode: resourceData.mode as Target["mode"] }) .where(eq(targets.targetId, existingTarget.targetId)) .returning(); @@ -1059,7 +1080,7 @@ export async function updatePublicResources( let domain; if ( - ["http", "ssh", "rdp", "vnc"].includes(resourceData.mode || "") + ["http", "ssh", "rdp", "vnc", "inference"].includes(resourceData.mode || "") ) { if (resourceData["full-domain"]?.startsWith("*.")) { const isLicensed = await isLicensedOrSubscribed( @@ -1156,12 +1177,12 @@ export async function updatePublicResources( status: resourceStatusFromSite, name: resourceData.name || "Unnamed Resource", mode: resourceData.mode, - proxyPort: ["http", "ssh", "rdp", "vnc"].includes( + proxyPort: ["http", "ssh", "rdp", "vnc", "inference"].includes( resourceData.mode || "" ) ? null : resourceData["proxy-port"], - fullDomain: ["http", "ssh", "rdp", "vnc"].includes( + fullDomain: ["http", "ssh", "rdp", "vnc", "inference"].includes( resourceData.mode || "" ) ? resourceData["full-domain"] @@ -1218,6 +1239,23 @@ export async function updatePublicResources( resource = newResource; + await syncInferenceAiConfig({ + orgId, + trx, + mode: resourceData.mode || "", + scope: "public", + resourceId: newResource.resourceId, + providers: (resourceData["ai-providers"] || []).map((p) => ({ + provider: p.provider, + accessMode: p["access-mode"], + enabled: p.enabled, + models: p.models.map((m) => ({ + model: m.model, + listType: m["list-type"] + })) + })) + }); + await trx.insert(roleResources).values({ roleId: adminRole.roleId, resourceId: newResource.resourceId diff --git a/server/lib/blueprints/types.ts b/server/lib/blueprints/types.ts index e1952848f..099656ddc 100644 --- a/server/lib/blueprints/types.ts +++ b/server/lib/blueprints/types.ts @@ -183,6 +183,34 @@ export const HeaderSchema = z.object({ value: z.string().min(1) }); +export const AiProviderModelEntrySchema = z.object({ + model: z.string().min(1), + "list-type": z.enum(["allow", "block"]) +}); + +export const AiProviderAttachmentSchema = z + .object({ + provider: z.string().min(1), + "access-mode": z + .enum(["inherit", "select"]) + .optional() + .default("inherit"), + enabled: z.boolean().optional().default(true), + models: z.array(AiProviderModelEntrySchema).optional().default([]) + }) + .refine( + (provider) => { + if (provider.models.length === 0) { + return true; + } + return provider["access-mode"] === "select"; + }, + { + path: ["models"], + error: "'models' can only be set on a provider with access-mode 'select'" + } + ); + export const AuthDaemonSchema = z .object({ pam: z.enum(["passthrough", "push"]).optional().default("passthrough"), @@ -209,7 +237,9 @@ export const PublicResourceSchema = z protocol: z .enum(["http", "tcp", "udp", "ssh", "rdp", "vnc"]) .optional(), // this was the old one and is now DEPRECATED in favor of the mode - mode: z.enum(["http", "tcp", "udp", "ssh", "rdp", "vnc"]).optional(), + mode: z + .enum(["http", "tcp", "udp", "ssh", "rdp", "vnc", "inference"]) + .optional(), policy: z.string().optional(), ssl: z.boolean().optional(), scheme: z.enum(["http", "https"]).optional(), @@ -226,7 +256,8 @@ export const PublicResourceSchema = z "auth-daemon": AuthDaemonSchema.optional(), "proxy-protocol": z.boolean().optional(), "proxy-protocol-version": z.int().min(1).optional(), - labels: z.array(z.string().min(1)).optional() + labels: z.array(z.string().min(1)).optional(), + "ai-providers": z.array(AiProviderAttachmentSchema).optional() }) .refine( (resource) => { @@ -315,11 +346,13 @@ export const PublicResourceSchema = z return true; } - // If protocol/mode is http, ssh, rdp, or vnc, it must have a full-domain + // If protocol/mode is http, ssh, rdp, vnc, or inference, it must have a full-domain const effectiveProtocol = resource.mode ?? resource.protocol; if ( effectiveProtocol !== undefined && - ["http", "ssh", "rdp", "vnc"].includes(effectiveProtocol) + ["http", "ssh", "rdp", "vnc", "inference"].includes( + effectiveProtocol + ) ) { return ( resource["full-domain"] !== undefined && @@ -330,7 +363,43 @@ export const PublicResourceSchema = z }, { path: ["full-domain"], - error: "When protocol is 'http', 'ssh', 'rdp', or 'vnc', a 'full-domain' must be provided" + error: "When protocol is 'http', 'ssh', 'rdp', 'vnc', or 'inference', a 'full-domain' must be provided" + } + ) + .refine( + (resource) => { + if (isTargetsOnlyResource(resource)) { + return true; + } + + const effectiveMode = resource.mode ?? resource.protocol; + if (effectiveMode !== "inference") { + return true; + } + + return resource.targets.every((target) => target == null); + }, + { + path: ["targets"], + error: "When mode is 'inference', 'targets' must not be provided" + } + ) + .refine( + (resource) => { + if (isTargetsOnlyResource(resource)) { + return true; + } + + const effectiveMode = resource.mode ?? resource.protocol; + if (effectiveMode === "inference") { + return true; + } + + return (resource["ai-providers"]?.length ?? 0) === 0; + }, + { + path: ["ai-providers"], + error: "'ai-providers' can only be set when mode is 'inference'" } ) .refine( @@ -464,7 +533,7 @@ export function isTargetsOnlyResource(resource: any): boolean { export const PrivateResourceSchema = z .object({ name: z.string().min(1).max(255), - mode: z.enum(["host", "cidr", "http", "ssh"]), + mode: z.enum(["host", "cidr", "http", "ssh", "inference"]), site: z.string().optional(), // DEPRECATED IN FAVOR OF sites sites: z.array(z.string()).optional().default([]), // protocol: z.enum(["tcp", "udp"]).optional(), @@ -495,16 +564,17 @@ export const PrivateResourceSchema = z users: z.array(z.string()).optional().default([]), machines: z.array(z.string()).optional().default([]), labels: z.array(z.string().min(1)).optional().default([]), - "auth-daemon": AuthDaemonSchema.optional() + "auth-daemon": AuthDaemonSchema.optional(), + "ai-providers": z.array(AiProviderAttachmentSchema).optional().default([]) }) .refine( (data) => { - // destination is optional only for ssh+native; required for everything else + // destination is optional only for ssh+native or inference; required for everything else const isNativeSSH = data.mode === "ssh" && (data["auth-daemon"] === undefined || data["auth-daemon"].mode === "native"); - if (!isNativeSSH && !data.destination) { + if (data.mode !== "inference" && !isNativeSSH && !data.destination) { return false; } return true; @@ -512,7 +582,19 @@ export const PrivateResourceSchema = z { path: ["destination"], message: - "destination is required unless mode is 'ssh' with auth-daemon mode 'native'" + "destination is required unless mode is 'ssh' with auth-daemon mode 'native', or mode is 'inference'" + } + ) + .refine( + (data) => { + if (data.mode === "inference") { + return true; + } + return (data["ai-providers"]?.length ?? 0) === 0; + }, + { + path: ["ai-providers"], + error: "'ai-providers' can only be set when mode is 'inference'" } ) .refine(