imrpoved logging (#5448)

Author: chitalian

worker/src/lib/ai-gateway/AttemptBuilder.ts

@@ -170,7 +170,7 @@ export class AttemptBuilder {
     }
 
     const providerData = providerDataResult.data;
-      
+
     // Build BYOK and PTB attempts in parallel since they're independent
     const [byokAttempts, ptbAttempts] = await Promise.all([
       this.buildByokAttempts(

worker/src/lib/durable-objects/Wallet.ts

@@ -535,7 +535,11 @@ export class Wallet extends DurableObject<Env> {
     orgId: string,
     escrowId: string,
     actualCost: number
-  ): { clickhouseLastCheckedAt: number; remainingBalance: number } {
+  ): {
+    clickhouseLastCheckedAt: number;
+    remainingBalance: number;
+    staleEscrowsCleared?: number;
+  } {
     const actualCostScaled = actualCost * SCALE_FACTOR;
     if (actualCostScaled < 0) {
       throw new Error("actualCost cannot be negative");
@@ -595,15 +599,31 @@ export class Wallet extends DurableObject<Env> {
         (totalCreditsPurchased - totalDebits - totalEscrow) / SCALE_FACTOR;
 
       // Probabilistic cleanup: 1% chance to clean stale escrows
+      let staleEscrowsCleared: number | undefined;
       if (Math.random() < 0.01) {
         const oneHourAgo = Date.now() - 60 * 60 * 1000;
-        this.ctx.storage.sql.exec(
-          "DELETE FROM escrows WHERE created_at < ?",
-          oneHourAgo
-        );
+        // Count before delete
+        const countBefore = this.ctx.storage.sql
+          .exec<{ count: number }>(
+            "SELECT COUNT(*) as count FROM escrows WHERE created_at < ?",
+            oneHourAgo
+          )
+          .one().count;
+
+        if (countBefore > 0) {
+          this.ctx.storage.sql.exec(
+            "DELETE FROM escrows WHERE created_at < ?",
+            oneHourAgo
+          );
+          staleEscrowsCleared = countBefore;
+        }
       }
 
-      return { clickhouseLastCheckedAt: result.checked_at, remainingBalance };
+      return {
+        clickhouseLastCheckedAt: result.checked_at,
+        remainingBalance,
+        staleEscrowsCleared,
+      };
     });
   }
 

worker/src/lib/managers/WalletManager.ts

@@ -69,7 +69,7 @@ export class WalletManager {
       }
       const escrowId = escrowResult.data.reservedEscrowId;
 
-      const { clickhouseLastCheckedAt, remainingBalance } =
+      const { clickhouseLastCheckedAt, remainingBalance, staleEscrowsCleared } =
         await this.walletStub.finalizeEscrow(
           organizationId,
           escrowId,
@@ -80,6 +80,22 @@ export class WalletManager {
       const walletKVSync = new WalletKVSync(this.env.WALLET_KV, organizationId);
       await walletKVSync.storeWalletState(remainingBalance);
 
+      // Log stale escrow cleanup to DataDog
+      if (staleEscrowsCleared !== undefined && staleEscrowsCleared > 0) {
+        const tracer = createDataDogTracer(this.env);
+        tracer.startTrace(
+          "ptb.stale_escrows_cleared",
+          "escrow_cleanup",
+          {
+            org_id: organizationId,
+            escrows_cleared: String(staleEscrowsCleared),
+          },
+          true // forceSample - always log cleanup events
+        );
+        tracer.finishTrace();
+        this.ctx.waitUntil(tracer.sendTrace());
+      }
+
       if (
         cost === undefined &&
         statusCode >= 200 &&