diff --git a/graphile/graphile-llm/src/metering.ts b/graphile/graphile-llm/src/metering.ts index a5f02a500..520c20f0f 100644 --- a/graphile/graphile-llm/src/metering.ts +++ b/graphile/graphile-llm/src/metering.ts @@ -41,6 +41,8 @@ export interface MeteringContext { billing: BillingConfig; /** Entity ID to meter against (from JWT claims) */ entityId: string; + /** Per-request correlation ID (from request.id pgSetting) */ + requestId: string | null; } export interface MeteringOptions { @@ -185,6 +187,7 @@ export async function meteredEmbed( // Record actual usage (input_chars as the metered amount) ctx.withPgClient(ctx.pgSettings, async (pgClient) => { await recordUsage(pgClient, ctx.billing, ctx.entityId, meterSlug, text.length, { + request_id: ctx.requestId, input_chars: text.length, dims: result.length, latency_ms: latencyMs, @@ -271,6 +274,7 @@ export async function meteredChat( const inputChars = messages.reduce((sum, m) => sum + m.content.length, 0); ctx.withPgClient(ctx.pgSettings, async (pgClient) => { await recordUsage(pgClient, ctx.billing, ctx.entityId, meterSlug, inputChars + result.length, { + request_id: ctx.requestId, input_chars: inputChars, output_chars: result.length, messages_count: messages.length, diff --git a/graphile/graphile-llm/src/plugins/metering-plugin.ts b/graphile/graphile-llm/src/plugins/metering-plugin.ts index 4c959a8fd..1c332fa71 100644 --- a/graphile/graphile-llm/src/plugins/metering-plugin.ts +++ b/graphile/graphile-llm/src/plugins/metering-plugin.ts @@ -66,6 +66,7 @@ async function buildMeteringContext( const pgSettings: Record = graphqlContext?.pgSettings ?? {}; const entityId = resolveEntityId(pgSettings); const databaseId = pgSettings['jwt.claims.database_id'] ?? null; + const requestId = pgSettings['request.id'] ?? null; if (!entityId || !databaseId) return null; const withPgClient: WithPgClient | undefined = graphqlContext?.withPgClient; @@ -88,6 +89,7 @@ async function buildMeteringContext( pgSettings, billing: billingConfig, entityId, + requestId, }; } diff --git a/graphql/server/src/middleware/graphile.ts b/graphql/server/src/middleware/graphile.ts index 139740210..2b2ebb000 100644 --- a/graphql/server/src/middleware/graphile.ts +++ b/graphql/server/src/middleware/graphile.ts @@ -275,15 +275,24 @@ const buildPreset = ( pgSettings['default_transaction_read_only'] = 'on'; } + if (req.requestId) { + pgSettings['request.id'] = req.requestId; + } + return { pgSettings }; } } + const anonSettings: Record = { + role: anonRole, + ...context, + }; + if (req?.requestId) { + anonSettings['request.id'] = req.requestId; + } + return { - pgSettings: { - role: anonRole, - ...context, - }, + pgSettings: anonSettings, }; }, }, diff --git a/graphql/server/src/middleware/llm-api.ts b/graphql/server/src/middleware/llm-api.ts index 644d8427f..531664d84 100644 --- a/graphql/server/src/middleware/llm-api.ts +++ b/graphql/server/src/middleware/llm-api.ts @@ -114,6 +114,9 @@ function getPgSettings(req: Request): Record { if (req.databaseId) { settings['jwt.claims.database_id'] = req.databaseId; } + if (req.requestId) { + settings['request.id'] = req.requestId; + } return settings; }