From f7aca854171a6b325e40be9de72ee4299e2f44f0 Mon Sep 17 00:00:00 2001 From: Owen Date: Tue, 11 Aug 2026 16:14:46 -0400 Subject: [PATCH] link the usage with the session to display together --- messages/en-US.json | 3 ++ server/db/pg/schema/schema.ts | 10 +++- server/db/sqlite/schema/schema.ts | 10 +++- server/lib/aiBudgetEnforcement.ts | 6 +++ server/routers/aiGateway/logAiSession.ts | 4 +- server/routers/aiGateway/pipeline.ts | 11 ++++- server/routers/auditLogs/queryAiSessionLog.ts | 49 ++++++++++++++++++- server/routers/auditLogs/types.ts | 10 ++++ src/app/[orgId]/settings/logs/ai/page.tsx | 34 ++++++++++++- 9 files changed, 129 insertions(+), 8 deletions(-) diff --git a/messages/en-US.json b/messages/en-US.json index d5ecdf780..44b1bf6a3 100644 --- a/messages/en-US.json +++ b/messages/en-US.json @@ -3411,6 +3411,9 @@ "aiSessionLogTruncated": "This session was truncated before storage and may be incomplete.", "aiSessionViewRaw": "View Raw JSON", "aiSessionViewChat": "View Chat", + "tokens": "Tokens", + "cost": "Cost", + "estimated": "estimated", "requestAnalyticsDescription": "View detailed request analytics for resources in this organization", "logRetentionRequestLabel": "HTTP Request Log Retention", "logRetentionRequestDescription": "How long to retain request logs", diff --git a/server/db/pg/schema/schema.ts b/server/db/pg/schema/schema.ts index cbb1e60cb..afaccc480 100644 --- a/server/db/pg/schema/schema.ts +++ b/server/db/pg/schema/schema.ts @@ -1831,6 +1831,13 @@ export const aiUsageRecords = pgTable( userId: varchar("userId").references(() => users.userId, { onDelete: "set null" }), + // Links this usage record back to the aiSessionLog row for the same + // request (aiSessionLog.sessionId), so token/cost usage can be shown + // alongside the session transcript. Not a DB-level FK - aiSessionLog + // lives in the separate logs database. Nullable because the session + // log may be disabled (retention set to 0) while usage tracking + // stays on. + sessionId: varchar("sessionId"), requestedModel: varchar("requestedModel").notNull(), promptTokens: integer("promptTokens").notNull().default(0), cacheReadTokens: integer("cacheReadTokens").notNull().default(0), @@ -1862,7 +1869,8 @@ export const aiUsageRecords = pgTable( t.orgId, t.userId, t.createdAt - ) + ), + index("idx_ai_usage_records_session").on(t.sessionId) ] ); diff --git a/server/db/sqlite/schema/schema.ts b/server/db/sqlite/schema/schema.ts index 8934638d7..bf3be9f87 100644 --- a/server/db/sqlite/schema/schema.ts +++ b/server/db/sqlite/schema/schema.ts @@ -1819,6 +1819,13 @@ export const aiUsageRecords = sqliteTable( userId: text("userId").references(() => users.userId, { onDelete: "set null" }), + // Links this usage record back to the aiSessionLog row for the same + // request (aiSessionLog.sessionId), so token/cost usage can be shown + // alongside the session transcript. Not a DB-level FK - aiSessionLog + // lives in the separate logs database. Nullable because the session + // log may be disabled (retention set to 0) while usage tracking + // stays on. + sessionId: text("sessionId"), requestedModel: text("requestedModel").notNull(), promptTokens: integer("promptTokens").notNull().default(0), cacheReadTokens: integer("cacheReadTokens").notNull().default(0), @@ -1852,7 +1859,8 @@ export const aiUsageRecords = sqliteTable( t.orgId, t.userId, t.createdAt - ) + ), + index("idx_ai_usage_records_session").on(t.sessionId) ] ); diff --git a/server/lib/aiBudgetEnforcement.ts b/server/lib/aiBudgetEnforcement.ts index 4177af9b7..d73f5ba59 100644 --- a/server/lib/aiBudgetEnforcement.ts +++ b/server/lib/aiBudgetEnforcement.ts @@ -428,6 +428,11 @@ export type UsageRecordInput = { usage: AiUsage; costUsd: number | null; createdAt?: number; + // Same id as the aiSessionLog row logged for this request, so the two + // can be joined to show token/cost usage alongside the session + // transcript. Undefined when the session wasn't logged (e.g. session + // log retention disabled for the org). + sessionId?: string; }; type AiUsageRecordInsert = InferInsertModel; @@ -539,6 +544,7 @@ export async function recordUsage(input: UsageRecordInput): Promise { resourceId: input.resourceId, siteResourceId: input.siteResourceId, userId: input.userId, + sessionId: input.sessionId, requestedModel: input.requestedModel, promptTokens: usage.promptTokens, cacheReadTokens: usage.cacheReadTokens, diff --git a/server/routers/aiGateway/logAiSession.ts b/server/routers/aiGateway/logAiSession.ts index 92723bed4..ee05517f5 100644 --- a/server/routers/aiGateway/logAiSession.ts +++ b/server/routers/aiGateway/logAiSession.ts @@ -1,4 +1,3 @@ -import { randomUUID } from "crypto"; import { logsDb, db, orgs, aiSessionLog, type AiProvider } from "@server/db"; import type { InferInsertModel } from "drizzle-orm"; import logger from "@server/logger"; @@ -169,6 +168,7 @@ function truncateBody(value: string): { value: string; truncated: boolean } { } export function logAiSession(data: { + sessionId: string; capability: AiCapability; provider: AiProvider; requestedModel: string | undefined; @@ -230,7 +230,7 @@ export function logAiSession(data: { } sessionLogBuffer.push({ - sessionId: randomUUID(), + sessionId: data.sessionId, orgId: sanitizeString(data.orgId), providerId: data.provider.providerId, capability: data.capability, diff --git a/server/routers/aiGateway/pipeline.ts b/server/routers/aiGateway/pipeline.ts index 579b51019..7648e2355 100644 --- a/server/routers/aiGateway/pipeline.ts +++ b/server/routers/aiGateway/pipeline.ts @@ -1,3 +1,4 @@ +import { randomUUID } from "crypto"; import { Request, Response } from "express"; import { and, eq, inArray } from "drizzle-orm"; import { @@ -594,6 +595,12 @@ export function recordAiGatewayCompletion(args: { const pricing = getModelPricing(provider.type as AiProviderType, model); const cost = calculateAiCost(pricing, usage); + // Shared by the usage record and the session log so the two can be + // joined later to show token/cost usage alongside the transcript - + // generated up front since neither buffered insert's row id is known + // until its next batch flush. + const sessionId = randomUUID(); + logger.info("AI gateway request usage", { capability, providerId: provider.providerId, @@ -618,7 +625,8 @@ export function recordAiGatewayCompletion(args: { userId: requestUserId, requestedModel: model ?? "unknown", usage, - costUsd: cost?.totalCost ?? null + costUsd: cost?.totalCost ?? null, + sessionId }); if (budgets.length > 0) { @@ -635,6 +643,7 @@ export function recordAiGatewayCompletion(args: { } logAiSession({ + sessionId, capability, provider, requestedModel, diff --git a/server/routers/auditLogs/queryAiSessionLog.ts b/server/routers/auditLogs/queryAiSessionLog.ts index 2412f8980..d61bdf5c8 100644 --- a/server/routers/auditLogs/queryAiSessionLog.ts +++ b/server/routers/auditLogs/queryAiSessionLog.ts @@ -2,6 +2,7 @@ import { logsDb, aiSessionLog, aiProviders, + aiUsageRecords, resources, siteResources, users, @@ -253,6 +254,51 @@ async function enrichWithDetails( } } + const usageMap = new Map< + string, + { + promptTokens: number; + cacheReadTokens: number; + cacheWriteTokens: number; + completionTokens: number; + reasoningTokens: number; + totalTokens: number; + costUsd: number | null; + estimated: boolean; + } + >(); + const sessionIds = logs.map((log) => log.sessionId); + if (sessionIds.length > 0) { + const usageDetails = await primaryDb + .select({ + sessionId: aiUsageRecords.sessionId, + promptTokens: aiUsageRecords.promptTokens, + cacheReadTokens: aiUsageRecords.cacheReadTokens, + cacheWriteTokens: aiUsageRecords.cacheWriteTokens, + completionTokens: aiUsageRecords.completionTokens, + reasoningTokens: aiUsageRecords.reasoningTokens, + totalTokens: aiUsageRecords.totalTokens, + costUsd: aiUsageRecords.costUsd, + estimated: aiUsageRecords.estimated + }) + .from(aiUsageRecords) + .where(inArray(aiUsageRecords.sessionId, sessionIds)); + + for (const u of usageDetails) { + if (!u.sessionId) continue; + usageMap.set(u.sessionId, { + promptTokens: u.promptTokens, + cacheReadTokens: u.cacheReadTokens, + cacheWriteTokens: u.cacheWriteTokens, + completionTokens: u.completionTokens, + reasoningTokens: u.reasoningTokens, + totalTokens: u.totalTokens, + costUsd: u.costUsd, + estimated: u.estimated + }); + } + } + return logs.map((log) => { const provider = providerMap.get(log.providerId); @@ -281,7 +327,8 @@ async function enrichWithDetails( providerType: provider?.type ?? null, resourceName, resourceNiceId, - userEmail: log.userId ? (userMap.get(log.userId) ?? null) : null + userEmail: log.userId ? (userMap.get(log.userId) ?? null) : null, + usage: usageMap.get(log.sessionId) ?? null }; }); } diff --git a/server/routers/auditLogs/types.ts b/server/routers/auditLogs/types.ts index 42b93904e..b63acbf30 100644 --- a/server/routers/auditLogs/types.ts +++ b/server/routers/auditLogs/types.ts @@ -119,6 +119,16 @@ export type QueryAiSessionLogResponse = { truncated: boolean; statusCode: number | null; createdAt: number; + usage: { + promptTokens: number; + cacheReadTokens: number; + cacheWriteTokens: number; + completionTokens: number; + reasoningTokens: number; + totalTokens: number; + costUsd: number | null; + estimated: boolean; + } | null; }[]; pagination: { total: number; diff --git a/src/app/[orgId]/settings/logs/ai/page.tsx b/src/app/[orgId]/settings/logs/ai/page.tsx index 67d18b94c..7330db111 100644 --- a/src/app/[orgId]/settings/logs/ai/page.tsx +++ b/src/app/[orgId]/settings/logs/ai/page.tsx @@ -459,7 +459,7 @@ export default function AiSessionLogsPage() { const renderExpandedRow = (row: any) => { return (
-
+
{t("aiSessionId")}

@@ -472,6 +472,26 @@ export default function AiSessionLogsPage() { {row.statusCode ?? "N/A"}

+
+ {t("tokens")} +

+ {row.usage + ? `${row.usage.totalTokens.toLocaleString()}${ + row.usage.estimated + ? ` (${t("estimated")})` + : "" + }` + : "N/A"} +

+
+
+ {t("cost")} +

+ {row.usage && row.usage.costUsd != null + ? `$${row.usage.costUsd.toFixed(4)}` + : "N/A"} +

+