From 653dd920ad25ef86ff8123963d5acfd326ca8f5d Mon Sep 17 00:00:00 2001 From: Owen Date: Wed, 12 Aug 2026 15:54:22 -0400 Subject: [PATCH] Allow remote header overrides and fix display issues with log --- messages/en-US.json | 1 + server/lib/readConfigFile.ts | 22 ++ server/lib/traefik/TraefikConfigManager.ts | 26 ++- server/lib/virtualApiKey.ts | 19 +- server/routers/aiGateway/pipeline.ts | 30 ++- .../routers/traefik/traefikConfigProvider.ts | 22 +- .../virtualApiKey/listMyVirtualApiKeys.ts | 18 +- src/app/[orgId]/settings/logs/ai/page.tsx | 3 - .../[orgId]/settings/logs/request/page.tsx | 2 + src/components/AiUsageAnalyticsData.tsx | 203 ++++++++---------- 10 files changed, 213 insertions(+), 133 deletions(-) diff --git a/messages/en-US.json b/messages/en-US.json index 0abe669ea..7630612ca 100644 --- a/messages/en-US.json +++ b/messages/en-US.json @@ -3437,6 +3437,7 @@ "validPassword": "Valid Password", "validEmail": "Valid email", "validSSO": "Valid SSO", + "validVirtualAPIKey": "Valid Virtual API Key", "view": "View", "configManaged": "Config Managed", "connectedClient": "Connected Client", diff --git a/server/lib/readConfigFile.ts b/server/lib/readConfigFile.ts index d5c52f13a..58d7b4006 100644 --- a/server/lib/readConfigFile.ts +++ b/server/lib/readConfigFile.ts @@ -111,6 +111,20 @@ export const configSchema = z }) .optional() .prefault({}), + remote_headers: z + .object({ + user_id: z.string().optional().default("Remote-User-Id"), + virtual_api_key_id: z + .string() + .optional() + .default("Remote-Virtual-Api-Key-Id"), + user: z.string().optional().default("Remote-User"), + email: z.string().optional().default("Remote-Email"), + name: z.string().optional().default("Remote-Name"), + role: z.string().optional().default("Remote-Role") + }) + .optional() + .prefault({}), resource_session_request_param: z .string() .optional() @@ -154,6 +168,14 @@ export const configSchema = z id: "P-Access-Token-Id", token: "P-Access-Token" }, + remote_headers: { + user_id: "Remote-User-Id", + virtual_api_key_id: "Remote-Virtual-Api-Key-Id", + user: "Remote-User", + email: "Remote-Email", + name: "Remote-Name", + role: "Remote-Role" + }, resource_session_request_param: "resource_session_request_param", dashboard_session_length_hours: 720, diff --git a/server/lib/traefik/TraefikConfigManager.ts b/server/lib/traefik/TraefikConfigManager.ts index a52c5232a..91bc37249 100644 --- a/server/lib/traefik/TraefikConfigManager.ts +++ b/server/lib/traefik/TraefikConfigManager.ts @@ -605,7 +605,31 @@ export class TraefikConfigManager { resourceSessionRequestParam: config.getRawConfig().server - .resource_session_request_param + .resource_session_request_param, + + remoteUserIdHeader: + config.getRawConfig().server.remote_headers + .user_id, + + remoteVirtualApiKeyIdHeader: + config.getRawConfig().server.remote_headers + .virtual_api_key_id, + + remoteUserHeader: + config.getRawConfig().server.remote_headers + .user, + + remoteEmailHeader: + config.getRawConfig().server.remote_headers + .email, + + remoteNameHeader: + config.getRawConfig().server.remote_headers + .name, + + remoteRoleHeader: + config.getRawConfig().server.remote_headers + .role } } }; diff --git a/server/lib/virtualApiKey.ts b/server/lib/virtualApiKey.ts index e4f7480db..471be89fd 100644 --- a/server/lib/virtualApiKey.ts +++ b/server/lib/virtualApiKey.ts @@ -5,6 +5,7 @@ import { import { db, resources, + User, virtualApiKeyResources, virtualApiKeys, type Transaction, @@ -148,12 +149,12 @@ async function selectUserVirtualApiKey( export async function getOrCreateUserVirtualApiKey(params: { orgId: string; - userId: string; + user: User; createdByUserId?: string | null; }): Promise<{ key: VirtualApiKey; secret: string }> { - const { orgId, userId, createdByUserId } = params; + const { orgId, user, createdByUserId } = params; - const existing = await selectUserVirtualApiKey(orgId, userId); + const existing = await selectUserVirtualApiKey(orgId, user.userId); if (existing) { return { key: existing, @@ -171,8 +172,8 @@ export async function getOrCreateUserVirtualApiKey(params: { virtualApiKeyId: minted.virtualApiKeyId, orgId, kind: "user", - userId, - name: null, + userId: user.userId, + name: `${user.name ?? user.username}'s API Key`, description: null, token: encryptVirtualApiKeyToken(minted.secret), lastChars: minted.lastChars, @@ -186,7 +187,7 @@ export async function getOrCreateUserVirtualApiKey(params: { return { key: created, secret: minted.secret }; } catch { - const raced = await selectUserVirtualApiKey(orgId, userId); + const raced = await selectUserVirtualApiKey(orgId, user.userId); if (raced) { return { key: raced, @@ -199,11 +200,11 @@ export async function getOrCreateUserVirtualApiKey(params: { export async function rotateUserVirtualApiKey(params: { orgId: string; - userId: string; + user: User; createdByUserId?: string | null; }): Promise<{ key: VirtualApiKey; secret: string }> { - const { orgId, userId, createdByUserId } = params; - const existing = await selectUserVirtualApiKey(orgId, userId); + const { orgId, user, createdByUserId } = params; + const existing = await selectUserVirtualApiKey(orgId, user.userId); if (!existing) { return getOrCreateUserVirtualApiKey(params); diff --git a/server/routers/aiGateway/pipeline.ts b/server/routers/aiGateway/pipeline.ts index 2d916a97a..06ca9cdae 100644 --- a/server/routers/aiGateway/pipeline.ts +++ b/server/routers/aiGateway/pipeline.ts @@ -180,15 +180,16 @@ export function applyRequestUserHeaders( if (!requestUser) { return; } - headers["Remote-User"] = requestUser.username; + const remoteHeaders = config.getRawConfig().server.remote_headers; + headers[remoteHeaders.user] = requestUser.username; if (requestUser.email) { - headers["Remote-Email"] = requestUser.email; + headers[remoteHeaders.email] = requestUser.email; } if (requestUser.name) { - headers["Remote-Name"] = requestUser.name; + headers[remoteHeaders.name] = requestUser.name; } if (requestUser.role) { - headers["Remote-Role"] = requestUser.role; + headers[remoteHeaders.role] = requestUser.role; } } @@ -236,14 +237,20 @@ async function resolveRequestUser( // Public inference: identity comes from Badger via Remote-* only when the // Traefik trust header proves the request passed verify-session (VAK). if (isAiGatewayTrustHeaderValid(req.headers as Record)) { + const remoteHeaders = config.getRawConfig().server.remote_headers; const virtualApiKeyId = - getRequestHeader(req, "remote-virtual-api-key-id") || null; - const userId = getRequestHeader(req, "remote-user-id"); + getRequestHeader(req, remoteHeaders.virtual_api_key_id) || null; + const userId = getRequestHeader(req, remoteHeaders.user_id); + logger.debug("+++++++AI gateway request identity from trust header", { + virtualApiKeyId, + userId + }); if (userId) { - const username = getRequestHeader(req, "remote-user") || userId; - const email = getRequestHeader(req, "remote-email"); - const name = getRequestHeader(req, "remote-name"); - const role = getRequestHeader(req, "remote-role"); + const username = + getRequestHeader(req, remoteHeaders.user) || userId; + const email = getRequestHeader(req, remoteHeaders.email); + const name = getRequestHeader(req, remoteHeaders.name); + const role = getRequestHeader(req, remoteHeaders.role); const orgRoles = orgId ? await getUserOrgRoles(userId, orgId) : []; return { @@ -742,6 +749,9 @@ export async function handleAiGatewayProxy( } logger.info(`AI gateway ${capability} request for host: ${host}`); + logger.debug("AI gateway request headers", { + headers: req.headers + }); const target = await resolveTarget(host); if (!target) { diff --git a/server/routers/traefik/traefikConfigProvider.ts b/server/routers/traefik/traefikConfigProvider.ts index e1847d2ed..02e05f5e0 100644 --- a/server/routers/traefik/traefikConfigProvider.ts +++ b/server/routers/traefik/traefikConfigProvider.ts @@ -68,7 +68,27 @@ export async function traefikConfigProvider( resourceSessionRequestParam: config.getRawConfig().server - .resource_session_request_param + .resource_session_request_param, + + remoteUserIdHeader: + config.getRawConfig().server.remote_headers + .user_id, + + remoteVirtualApiKeyIdHeader: + config.getRawConfig().server.remote_headers + .virtual_api_key_id, + + remoteUserHeader: + config.getRawConfig().server.remote_headers.user, + + remoteEmailHeader: + config.getRawConfig().server.remote_headers.email, + + remoteNameHeader: + config.getRawConfig().server.remote_headers.name, + + remoteRoleHeader: + config.getRawConfig().server.remote_headers.role } } }; diff --git a/server/routers/virtualApiKey/listMyVirtualApiKeys.ts b/server/routers/virtualApiKey/listMyVirtualApiKeys.ts index 5177a975a..46531560a 100644 --- a/server/routers/virtualApiKey/listMyVirtualApiKeys.ts +++ b/server/routers/virtualApiKey/listMyVirtualApiKeys.ts @@ -3,6 +3,7 @@ import { z } from "zod"; import { db, resources, + users, virtualApiKeyResources, virtualApiKeys, type VirtualApiKey @@ -152,9 +153,24 @@ export async function listMyVirtualApiKeys( resourceName = resource.name; } + const [user] = await db + .select() + .from(users) + .where(eq(users.userId, userId)) + .limit(1); + + if (!user) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `User with ID ${userId} not found` + ) + ); + } + const { key: userKeyRow } = await getOrCreateUserVirtualApiKey({ orgId, - userId, + user, createdByUserId: userId }); diff --git a/src/app/[orgId]/settings/logs/ai/page.tsx b/src/app/[orgId]/settings/logs/ai/page.tsx index 15c573337..8c077929b 100644 --- a/src/app/[orgId]/settings/logs/ai/page.tsx +++ b/src/app/[orgId]/settings/logs/ai/page.tsx @@ -577,9 +577,6 @@ export default function AiSessionLogsPage() { {row.usage && (
-
- {t("tokenUsage")} -
{t("promptTokens")} diff --git a/src/app/[orgId]/settings/logs/request/page.tsx b/src/app/[orgId]/settings/logs/request/page.tsx index a4d1dec37..32e4f4b86 100644 --- a/src/app/[orgId]/settings/logs/request/page.tsx +++ b/src/app/[orgId]/settings/logs/request/page.tsx @@ -247,6 +247,7 @@ export default function GeneralPage() { // 106 - Valid email // 107 - Valid SSO // 108 - Connected Client + // 109 - Valid Virtual API Key // 201 - Resource Not Found // 202 - Resource Blocked @@ -265,6 +266,7 @@ export default function GeneralPage() { 106: t("validEmail"), 107: t("validSSO"), 108: t("connectedClient"), + 109: t("validVirtualAPIKey"), 201: t("resourceNotFound"), 202: t("resourceBlocked"), 203: t("droppedByRule"), diff --git a/src/components/AiUsageAnalyticsData.tsx b/src/components/AiUsageAnalyticsData.tsx index a5d3db2b8..3cb35198e 100644 --- a/src/components/AiUsageAnalyticsData.tsx +++ b/src/components/AiUsageAnalyticsData.tsx @@ -161,120 +161,107 @@ export function AiUsageAnalyticsData(props: AiUsageAnalyticsDataProps) { return (
- -
- + + - + -
- - setFilter("providerId", v) - } - /> - setFilter("model", v)} - /> - - setFilter("resourceId", v) - } - /> - setFilter("roleId", v)} - /> - setFilter("userId", v)} - /> - - setFilter("virtualApiKeyId", v) - } - /> + setFilter("providerId", v)} + /> + setFilter("model", v)} + /> + setFilter("resourceId", v)} + /> + setFilter("roleId", v)} + /> + setFilter("userId", v)} + /> + setFilter("virtualApiKeyId", v)} + /> - {!isEmptySearchParams && ( - - )} -
-
-
+ {!isEmptySearchParams && ( -
+ )} + + @@ -301,7 +288,7 @@ type FilterSelectProps = { function FilterSelect(props: FilterSelectProps) { return ( -
+