Allow remote header overrides and fix display issues with log

This commit is contained in:
Owen
2026-08-12 15:54:22 -04:00
parent 47ae017f94
commit 653dd920ad
10 changed files with 213 additions and 133 deletions
+1
View File
@@ -3437,6 +3437,7 @@
"validPassword": "Valid Password", "validPassword": "Valid Password",
"validEmail": "Valid email", "validEmail": "Valid email",
"validSSO": "Valid SSO", "validSSO": "Valid SSO",
"validVirtualAPIKey": "Valid Virtual API Key",
"view": "View", "view": "View",
"configManaged": "Config Managed", "configManaged": "Config Managed",
"connectedClient": "Connected Client", "connectedClient": "Connected Client",
+22
View File
@@ -111,6 +111,20 @@ export const configSchema = z
}) })
.optional() .optional()
.prefault({}), .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 resource_session_request_param: z
.string() .string()
.optional() .optional()
@@ -154,6 +168,14 @@ export const configSchema = z
id: "P-Access-Token-Id", id: "P-Access-Token-Id",
token: "P-Access-Token" 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:
"resource_session_request_param", "resource_session_request_param",
dashboard_session_length_hours: 720, dashboard_session_length_hours: 720,
+25 -1
View File
@@ -605,7 +605,31 @@ export class TraefikConfigManager {
resourceSessionRequestParam: resourceSessionRequestParam:
config.getRawConfig().server 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
} }
} }
}; };
+10 -9
View File
@@ -5,6 +5,7 @@ import {
import { import {
db, db,
resources, resources,
User,
virtualApiKeyResources, virtualApiKeyResources,
virtualApiKeys, virtualApiKeys,
type Transaction, type Transaction,
@@ -148,12 +149,12 @@ async function selectUserVirtualApiKey(
export async function getOrCreateUserVirtualApiKey(params: { export async function getOrCreateUserVirtualApiKey(params: {
orgId: string; orgId: string;
userId: string; user: User;
createdByUserId?: string | null; createdByUserId?: string | null;
}): Promise<{ key: VirtualApiKey; secret: string }> { }): 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) { if (existing) {
return { return {
key: existing, key: existing,
@@ -171,8 +172,8 @@ export async function getOrCreateUserVirtualApiKey(params: {
virtualApiKeyId: minted.virtualApiKeyId, virtualApiKeyId: minted.virtualApiKeyId,
orgId, orgId,
kind: "user", kind: "user",
userId, userId: user.userId,
name: null, name: `${user.name ?? user.username}'s API Key`,
description: null, description: null,
token: encryptVirtualApiKeyToken(minted.secret), token: encryptVirtualApiKeyToken(minted.secret),
lastChars: minted.lastChars, lastChars: minted.lastChars,
@@ -186,7 +187,7 @@ export async function getOrCreateUserVirtualApiKey(params: {
return { key: created, secret: minted.secret }; return { key: created, secret: minted.secret };
} catch { } catch {
const raced = await selectUserVirtualApiKey(orgId, userId); const raced = await selectUserVirtualApiKey(orgId, user.userId);
if (raced) { if (raced) {
return { return {
key: raced, key: raced,
@@ -199,11 +200,11 @@ export async function getOrCreateUserVirtualApiKey(params: {
export async function rotateUserVirtualApiKey(params: { export async function rotateUserVirtualApiKey(params: {
orgId: string; orgId: string;
userId: string; user: User;
createdByUserId?: string | null; createdByUserId?: string | null;
}): Promise<{ key: VirtualApiKey; secret: string }> { }): 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) { if (!existing) {
return getOrCreateUserVirtualApiKey(params); return getOrCreateUserVirtualApiKey(params);
+20 -10
View File
@@ -180,15 +180,16 @@ export function applyRequestUserHeaders(
if (!requestUser) { if (!requestUser) {
return; return;
} }
headers["Remote-User"] = requestUser.username; const remoteHeaders = config.getRawConfig().server.remote_headers;
headers[remoteHeaders.user] = requestUser.username;
if (requestUser.email) { if (requestUser.email) {
headers["Remote-Email"] = requestUser.email; headers[remoteHeaders.email] = requestUser.email;
} }
if (requestUser.name) { if (requestUser.name) {
headers["Remote-Name"] = requestUser.name; headers[remoteHeaders.name] = requestUser.name;
} }
if (requestUser.role) { 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 // Public inference: identity comes from Badger via Remote-* only when the
// Traefik trust header proves the request passed verify-session (VAK). // Traefik trust header proves the request passed verify-session (VAK).
if (isAiGatewayTrustHeaderValid(req.headers as Record<string, string>)) { if (isAiGatewayTrustHeaderValid(req.headers as Record<string, string>)) {
const remoteHeaders = config.getRawConfig().server.remote_headers;
const virtualApiKeyId = const virtualApiKeyId =
getRequestHeader(req, "remote-virtual-api-key-id") || null; getRequestHeader(req, remoteHeaders.virtual_api_key_id) || null;
const userId = getRequestHeader(req, "remote-user-id"); const userId = getRequestHeader(req, remoteHeaders.user_id);
logger.debug("+++++++AI gateway request identity from trust header", {
virtualApiKeyId,
userId
});
if (userId) { if (userId) {
const username = getRequestHeader(req, "remote-user") || userId; const username =
const email = getRequestHeader(req, "remote-email"); getRequestHeader(req, remoteHeaders.user) || userId;
const name = getRequestHeader(req, "remote-name"); const email = getRequestHeader(req, remoteHeaders.email);
const role = getRequestHeader(req, "remote-role"); const name = getRequestHeader(req, remoteHeaders.name);
const role = getRequestHeader(req, remoteHeaders.role);
const orgRoles = orgId ? await getUserOrgRoles(userId, orgId) : []; const orgRoles = orgId ? await getUserOrgRoles(userId, orgId) : [];
return { return {
@@ -742,6 +749,9 @@ export async function handleAiGatewayProxy(
} }
logger.info(`AI gateway ${capability} request for host: ${host}`); logger.info(`AI gateway ${capability} request for host: ${host}`);
logger.debug("AI gateway request headers", {
headers: req.headers
});
const target = await resolveTarget(host); const target = await resolveTarget(host);
if (!target) { if (!target) {
@@ -68,7 +68,27 @@ export async function traefikConfigProvider(
resourceSessionRequestParam: resourceSessionRequestParam:
config.getRawConfig().server 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
} }
} }
}; };
@@ -3,6 +3,7 @@ import { z } from "zod";
import { import {
db, db,
resources, resources,
users,
virtualApiKeyResources, virtualApiKeyResources,
virtualApiKeys, virtualApiKeys,
type VirtualApiKey type VirtualApiKey
@@ -152,9 +153,24 @@ export async function listMyVirtualApiKeys(
resourceName = resource.name; 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({ const { key: userKeyRow } = await getOrCreateUserVirtualApiKey({
orgId, orgId,
userId, user,
createdByUserId: userId createdByUserId: userId
}); });
@@ -577,9 +577,6 @@ export default function AiSessionLogsPage() {
</div> </div>
{row.usage && ( {row.usage && (
<div> <div>
<div className="text-xs font-medium mb-2">
{t("tokenUsage")}
</div>
<div className="grid grid-cols-3 sm:grid-cols-6 gap-4 text-xs"> <div className="grid grid-cols-3 sm:grid-cols-6 gap-4 text-xs">
<div> <div>
<strong>{t("promptTokens")}</strong> <strong>{t("promptTokens")}</strong>
@@ -247,6 +247,7 @@ export default function GeneralPage() {
// 106 - Valid email // 106 - Valid email
// 107 - Valid SSO // 107 - Valid SSO
// 108 - Connected Client // 108 - Connected Client
// 109 - Valid Virtual API Key
// 201 - Resource Not Found // 201 - Resource Not Found
// 202 - Resource Blocked // 202 - Resource Blocked
@@ -265,6 +266,7 @@ export default function GeneralPage() {
106: t("validEmail"), 106: t("validEmail"),
107: t("validSSO"), 107: t("validSSO"),
108: t("connectedClient"), 108: t("connectedClient"),
109: t("validVirtualAPIKey"),
201: t("resourceNotFound"), 201: t("resourceNotFound"),
202: t("resourceBlocked"), 202: t("resourceBlocked"),
203: t("droppedByRule"), 203: t("droppedByRule"),
+95 -108
View File
@@ -161,120 +161,107 @@ export function AiUsageAnalyticsData(props: AiUsageAnalyticsDataProps) {
return ( return (
<div className="flex flex-col gap-5"> <div className="flex flex-col gap-5">
<Card> <Card>
<CardHeader className="flex flex-col gap-4"> <CardHeader className="flex flex-row flex-wrap items-end gap-x-3 gap-y-3">
<div className="flex flex-col lg:flex-row items-start lg:items-end w-full gap-2"> <DateRangePicker
<DateRangePicker startValue={{
startValue={{ date: dateRange.startDate,
date: dateRange.startDate, time: dateRange.startDate
time: dateRange.startDate ? getDateTime(dateRange.startDate)
? getDateTime(dateRange.startDate) : undefined
: undefined }}
}} endValue={{
endValue={{ date: dateRange.endDate,
date: dateRange.endDate, time: dateRange.endDate
time: dateRange.endDate ? getDateTime(dateRange.endDate)
? getDateTime(dateRange.endDate) : undefined
: undefined }}
}} onRangeChange={handleTimeRangeUpdate}
onRangeChange={handleTimeRangeUpdate} className="flex-wrap gap-2"
className="flex-wrap gap-2" />
/>
<Separator className="w-px h-6 self-end relative bottom-1.5 hidden lg:block" /> <Separator className="w-px h-6 self-end relative bottom-1.5 hidden lg:block" />
<div className="flex flex-wrap items-end gap-2"> <FilterSelect
<FilterSelect id="providerId"
id="providerId" label={t("aiUsageFilterProvider")}
label={t("aiUsageFilterProvider")} value={filters.providerId?.toString()}
value={filters.providerId?.toString()} options={providerOptions}
options={providerOptions} placeholder={t("aiUsageFilterAllProviders")}
placeholder={t("aiUsageFilterAllProviders")} onValueChange={(v) => setFilter("providerId", v)}
onValueChange={(v) => />
setFilter("providerId", v) <FilterSelect
} id="model"
/> label={t("aiUsageFilterModel")}
<FilterSelect value={filters.model}
id="model" options={modelOptions}
label={t("aiUsageFilterModel")} placeholder={t("aiUsageFilterAllModels")}
value={filters.model} onValueChange={(v) => setFilter("model", v)}
options={modelOptions} />
placeholder={t("aiUsageFilterAllModels")} <FilterSelect
onValueChange={(v) => setFilter("model", v)} id="resourceId"
/> label={t("aiUsageFilterResource")}
<FilterSelect value={filters.resourceId?.toString()}
id="resourceId" options={resourceOptions}
label={t("aiUsageFilterResource")} placeholder={t("aiUsageFilterAllResources")}
value={filters.resourceId?.toString()} onValueChange={(v) => setFilter("resourceId", v)}
options={resourceOptions} />
placeholder={t("aiUsageFilterAllResources")} <FilterSelect
onValueChange={(v) => id="roleId"
setFilter("resourceId", v) label={t("aiUsageFilterRole")}
} value={filters.roleId?.toString()}
/> options={roleOptions}
<FilterSelect placeholder={t("aiUsageFilterAllRoles")}
id="roleId" onValueChange={(v) => setFilter("roleId", v)}
label={t("aiUsageFilterRole")} />
value={filters.roleId?.toString()} <FilterSelect
options={roleOptions} id="userId"
placeholder={t("aiUsageFilterAllRoles")} label={t("aiUsageFilterUser")}
onValueChange={(v) => setFilter("roleId", v)} value={filters.userId}
/> options={userOptions}
<FilterSelect placeholder={t("aiUsageFilterAllUsers")}
id="userId" onValueChange={(v) => setFilter("userId", v)}
label={t("aiUsageFilterUser")} />
value={filters.userId} <FilterSelect
options={userOptions} id="virtualApiKeyId"
placeholder={t("aiUsageFilterAllUsers")} label={t("aiUsageFilterVirtualApiKey")}
onValueChange={(v) => setFilter("userId", v)} value={filters.virtualApiKeyId}
/> options={virtualApiKeyOptions}
<FilterSelect placeholder={t("aiUsageFilterAllVirtualApiKeys")}
id="virtualApiKeyId" onValueChange={(v) => setFilter("virtualApiKeyId", v)}
label={t("aiUsageFilterVirtualApiKey")} />
value={filters.virtualApiKeyId}
options={virtualApiKeyOptions}
placeholder={t(
"aiUsageFilterAllVirtualApiKeys"
)}
onValueChange={(v) =>
setFilter("virtualApiKeyId", v)
}
/>
{!isEmptySearchParams && ( {!isEmptySearchParams && (
<Button
variant="ghost"
onClick={() => router.replace(path)}
className="gap-2"
>
<XIcon className="size-4" />
{t("aiUsageResetFilters")}
</Button>
)}
</div>
</div>
<div className="flex justify-end">
<Button <Button
variant="outline" variant="ghost"
onClick={() => onClick={() => router.replace(path)}
queryClient.invalidateQueries({
queryKey: [
...AI_USAGE_ANALYTICS_QUERY_PREFIX,
props.orgId
]
})
}
disabled={isFetching}
className="gap-2" className="gap-2"
> >
<RefreshCw <XIcon className="size-4" />
className={cn( {t("aiUsageResetFilters")}
"size-4",
isFetching && "animate-spin"
)}
/>
{t("aiUsageRefresh")}
</Button> </Button>
</div> )}
<Button
variant="outline"
onClick={() =>
queryClient.invalidateQueries({
queryKey: [
...AI_USAGE_ANALYTICS_QUERY_PREFIX,
props.orgId
]
})
}
disabled={isFetching}
className="gap-2 ml-auto"
>
<RefreshCw
className={cn(
"size-4",
isFetching && "animate-spin"
)}
/>
{t("aiUsageRefresh")}
</Button>
</CardHeader> </CardHeader>
</Card> </Card>
@@ -301,7 +288,7 @@ type FilterSelectProps = {
function FilterSelect(props: FilterSelectProps) { function FilterSelect(props: FilterSelectProps) {
return ( return (
<div className="flex flex-col items-start gap-2 w-44"> <div className="flex flex-col items-start gap-2 w-36">
<Label htmlFor={props.id}>{props.label}</Label> <Label htmlFor={props.id}>{props.label}</Label>
<Select <Select
onValueChange={(newValue) => onValueChange={(newValue) =>