link the usage with the session to display together

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,

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,

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
         };
     });
 }

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;