From a19c074f80a0dc6b4af292b63f98d75a73ff053d Mon Sep 17 00:00:00 2001
From: Mark IJbema <mark@kilocode.ai>
Date: Tue, 19 May 2026 16:06:06 +0200
Subject: [PATCH 1/2] refactor(gateway): return inserted usage row id and
 timestamp

`logMicrodollarUsage` and `processTokenData` now surface the
`{ usageId, createdAt }` of the row they inserted. Existing callers
ignore the return value, so behavior is unchanged. Lets follow-up
writes (e.g. side-table attribution) key onto the same usage row
without re-deriving the id.
---
 apps/web/src/lib/ai-gateway/processUsage.ts | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/apps/web/src/lib/ai-gateway/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts
index 35005f0aed7..a4c82232103 100644
--- a/apps/web/src/lib/ai-gateway/processUsage.ts
+++ b/apps/web/src/lib/ai-gateway/processUsage.ts
@@ -245,7 +245,7 @@ export function toInsertableDbUsageRecord(
 export async function logMicrodollarUsage(
   usageStats: MicrodollarUsageStats,
   usageContext: MicrodollarUsageContext
-) {
+): Promise<{ usageId: string; createdAt: string }> {
   usageContext.status_code = usageStats.status_code;
   const contextInfo = extractUsageContextInfo(usageContext);
   const { core, metadata } = toInsertableDbUsageRecord(usageStats, contextInfo);
@@ -256,6 +256,8 @@ export async function logMicrodollarUsage(
     usageContext.prior_microdollar_usage,
     usageContext.posthog_distinct_id ?? null
   );
+
+  return { usageId: core.id, createdAt: core.created_at };
 }
 
 async function saveUsageRelatedData(
@@ -909,14 +911,14 @@ export function calculateKiloExclusiveCost_mUsd(
 export async function processTokenData(
   usageStats: MicrodollarUsageStats | null,
   usageContext: MicrodollarUsageContext
-) {
+): Promise<{ usageId: string; createdAt: string } | null> {
   if (!usageStats) {
     captureMessage('SUSPICIOUS: No usage information', {
       level: 'error',
       tags: { source: 'usage_processing' },
       extra: { usageContext },
     });
-    return;
+    return null;
   }
 
   const timer = createTimer();
@@ -987,7 +989,7 @@ export async function processTokenData(
     usageStats.cacheDiscount_mUsd = 0;
   }
 
-  await logMicrodollarUsage(usageStats, usageContext);
+  return logMicrodollarUsage(usageStats, usageContext);
 }
 
 function useAnthropicStyleTokenCounting(requestedModel: string, provider: ProviderId) {

From 0fcea99163c1954e16ed0c6d473b0a503a60feb6 Mon Sep 17 00:00:00 2001
From: Mark IJbema <mark@kilocode.ai>
Date: Tue, 19 May 2026 16:25:46 +0200
Subject: [PATCH 2/2] fix(gateway): return null from logMicrodollarUsage when
 insert fails

Per kilo-code-bot review on PR #3336: `insertUsageRecord` catches all
exceptions and returns null without re-throwing, so a swallowed DB
failure would have caused this PR's new `{ usageId, createdAt }`
return to claim a row that was never persisted.

Thread the insert result through `saveUsageRelatedData` and return
`null` from `logMicrodollarUsage` / `processTokenData` when the
underlying insert returned null. Existing callers ignore the return,
so behavior is unchanged for them. Future callers (the experiment
attribution write in #3325) will see `null` and skip the dependent
FK insert instead of dangling on a missing row.
---
 apps/web/src/lib/ai-gateway/processUsage.ts | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/apps/web/src/lib/ai-gateway/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts
index a4c82232103..e123e473b3d 100644
--- a/apps/web/src/lib/ai-gateway/processUsage.ts
+++ b/apps/web/src/lib/ai-gateway/processUsage.ts
@@ -245,19 +245,22 @@ export function toInsertableDbUsageRecord(
 export async function logMicrodollarUsage(
   usageStats: MicrodollarUsageStats,
   usageContext: MicrodollarUsageContext
-): Promise<{ usageId: string; createdAt: string }> {
+): Promise<{ usageId: string; createdAt: string } | null> {
   usageContext.status_code = usageStats.status_code;
   const contextInfo = extractUsageContextInfo(usageContext);
   const { core, metadata } = toInsertableDbUsageRecord(usageStats, contextInfo);
 
-  await saveUsageRelatedData(
+  const inserted = await saveUsageRelatedData(
     core,
     metadata,
     usageContext.prior_microdollar_usage,
     usageContext.posthog_distinct_id ?? null
   );
 
-  return { usageId: core.id, createdAt: core.created_at };
+  // `insertUsageRecord` swallows DB errors and returns null; surface that
+  // failure to callers so dependent FK writes don't dangle on a row that
+  // was never persisted.
+  return inserted ? { usageId: core.id, createdAt: core.created_at } : null;
 }
 
 async function saveUsageRelatedData(
@@ -265,7 +268,7 @@ async function saveUsageRelatedData(
   metadataFields: UsageMetaData,
   prior_microdollar_usage: number,
   posthog_distinct_id: string | null
-) {
+): Promise<BalanceUpdateResult> {
   const isFirst = await isFirstUsage(coreUsageFields, prior_microdollar_usage);
   if (isFirst && posthog_distinct_id)
     await sendFirstUsageEvent(coreUsageFields, posthog_distinct_id);
@@ -279,6 +282,7 @@ async function saveUsageRelatedData(
     );
   }
   await ingestOrganizationTokenUsage(coreUsageFields);
+  return balanceUpdateResult;
 }
 
 async function isFirstUsage(