Justin/free limit stop logging (#5487)

Author: chitalian

.claude/settings.local.json

@@ -105,6 +105,7 @@
       "WebFetch(domain:ai.google.dev)",
       "Bash(npx tsoa:*)",
       "Bash(python3:*)",
+      "Bash(lsof:*)",
       "Bash(git mv:*)",
       "Bash(npm run test:rate-limit:*)",
       "Bash(npx eslint:*)",

bifrost/lib/clients/jawnTypes/private.ts

@@ -2580,6 +2580,7 @@ Json: JsonObject;
     /** @enum {string} */
     BodyMappingType: "OPENAI" | "NO_MAPPING" | "RESPONSES";
     HeliconeMeta: {
+      freeLimitExceeded?: boolean;
       aiGatewayBodyMapping?: components["schemas"]["BodyMappingType"];
       providerModelId?: string;
       gatewayModel?: string;
@@ -2646,6 +2647,21 @@ Json: JsonObject;
     };
     Log: {
       response: {
+        model?: string;
+        /** Format: double */
+        reasoningTokens?: number;
+        /** Format: double */
+        completionAudioTokens?: number;
+        /** Format: double */
+        promptAudioTokens?: number;
+        /** Format: double */
+        promptCacheWriteTokens?: number;
+        /** Format: double */
+        promptCacheReadTokens?: number;
+        /** Format: double */
+        completionTokens?: number;
+        /** Format: double */
+        promptTokens?: number;
         /** Format: double */
         cost?: number;
         /** Format: double */

packages/llm-mapper/types.ts

@@ -271,6 +271,7 @@ type HeliconeMetadata = {
   promptVersion?: string | null;
   targetUrl?: string | null;
   requestReferrer?: string | null;
+  storageLocation?: string | null;
 };
 
 // UNORGANZIED

packages/llm-mapper/utils/getMappedContent.ts

@@ -104,6 +104,7 @@ const metaDataFromHeliconeRequest = (
     promptVersion: heliconeRequest.prompt_version ?? null,
     targetUrl: heliconeRequest.target_url ?? null,
     requestReferrer: heliconeRequest.request_referrer ?? null,
+    storageLocation: heliconeRequest.storage_location ?? null,
   };
 };
 

supabase/migrations/20260121000000_free_limit_exceeded_to_month.sql

@@ -0,0 +1,7 @@
+-- Add free_limit_exceeded column to track which month the free tier limit was exceeded
+-- Stores month in YYYY-MM format (e.g., "2026-01" = exceeded in January 2026)
+-- NULL = not exceeded or under limit
+-- When set, request/response bodies are not stored for non-PTB requests
+
+ALTER TABLE organization
+ADD COLUMN IF NOT EXISTS free_limit_exceeded TEXT DEFAULT NULL;

valhalla/jawn/src/lib/handlers/HandlerContext.ts

@@ -16,7 +16,7 @@ export class HandlerContext extends SetOnce {
   public legacyUsage: Usage;
   public usage?: ModelUsage;
   public costBreakdown?: CostBreakdown;
-  public storageLocation?: "s3" | "clickhouse";
+  public storageLocation?: "s3" | "clickhouse" | "not_stored_exceeded_free";
   public sizeBytes?: number;
   public rawLog: RawLog;
   public processedLog: ProcessedLog;
@@ -70,6 +70,16 @@ export type Log = {
     delayMs: number;
     cachedLatency?: number;
     cost?: number;
+    // Token usage (from Worker when body isn't stored)
+    promptTokens?: number;
+    completionTokens?: number;
+    promptCacheReadTokens?: number;
+    promptCacheWriteTokens?: number;
+    promptAudioTokens?: number;
+    completionAudioTokens?: number;
+    reasoningTokens?: number;
+    // Model (from Worker when body isn't stored)
+    model?: string;
   };
 };
 
@@ -80,6 +90,7 @@ export type Usage = {
   promptAudioTokens?: number;
   completionTokens?: number;
   completionAudioTokens?: number;
+  reasoningTokens?: number;
 
   // anthropic cache control
   promptCacheWrite5m?: number;
@@ -139,6 +150,9 @@ export type HeliconeMeta = {
   gatewayModel?: string; // registry format
   providerModelId?: string; // provider format
   aiGatewayBodyMapping?: BodyMappingType; // body mapping type
+
+  // Free tier limit
+  freeLimitExceeded?: boolean;
 };
 
 export type KafkaMessageContents = {

valhalla/jawn/src/lib/handlers/LoggingHandler.ts

@@ -43,7 +43,7 @@ type RequestRecord = {
   organizationId: string;
   requestBody: string;
   responseBody: string;
-  location: "s3" | "clickhouse";
+  location: "s3" | "clickhouse" | "not_stored_exceeded_free";
 };
 
 // Legacy type definitions for deleted tables
@@ -171,9 +171,16 @@ export class LoggingHandler extends AbstractLogHandler {
       });
 
       context.sizeBytes = size ?? 0;
-      // if we know size is def less than 10mb use clickhouse otherwise just stick to s3
-      context.storageLocation =
-        size && size <= S3_MIN_SIZE_THRESHOLD ? "clickhouse" : "s3";
+      // Determine storage location:
+      // 1. If free tier limit exceeded, bodies were not stored
+      // 2. If size is small enough, use clickhouse
+      // 3. Otherwise use s3
+      if (context.message.heliconeMeta.freeLimitExceeded) {
+        context.storageLocation = "not_stored_exceeded_free";
+      } else {
+        context.storageLocation =
+          size && size <= S3_MIN_SIZE_THRESHOLD ? "clickhouse" : "s3";
+      }
 
       const requestMapped = this.mapRequest(context);
       const responseMapped = this.mapResponse(context);
@@ -634,6 +641,13 @@ export class LoggingHandler extends AbstractLogHandler {
     responseText: string;
   } {
     try {
+      // If free tier limit exceeded, bodies were not stored
+      if (context.storageLocation === "not_stored_exceeded_free") {
+        return {
+          requestText: "",
+          responseText: "",
+        };
+      }
       if (context.storageLocation === "clickhouse") {
         return {
           requestText: JSON.stringify(context.processedLog.request.body),

valhalla/jawn/src/lib/handlers/RateLimitHandler.ts

@@ -8,6 +8,10 @@ import {
 import { RateLimitStore } from "../stores/RateLimitStore";
 import { AbstractLogHandler } from "./AbstractLogHandler";
 import { HandlerContext } from "./HandlerContext";
+import { dbQueryClickhouse, dbExecute } from "../shared/db/dbExecute";
+
+const FREE_TIER_LIMIT = 10_000;
+const FREE_TIER_CHECK_PROBABILITY = 0.01; // 1% of requests
 
 export class RateLimitHandler extends AbstractLogHandler {
   private rateLimitStore: RateLimitStore;
@@ -30,6 +34,17 @@ export class RateLimitHandler extends AbstractLogHandler {
     }
 
     try {
+      // Probabilistic free tier limit check (1% of requests)
+      // - If freeLimitExceeded is null: check with 1% probability
+      // - If freeLimitExceeded is current month: skip check (already exceeded this month)
+      // - If freeLimitExceeded is old month: check with 1% probability (may have reset)
+      if (
+        context.orgParams.tier === "free" &&
+        this.shouldCheckFreeTierLimit(context.orgParams.freeLimitExceeded)
+      ) {
+        await this.checkAndUpdateFreeTierLimit(context.orgParams.id);
+      }
+
       const { data: isRateLimited, error: rateLimitErr } = this.rateLimitEntry(
         context.orgParams.id,
         context.orgParams.percentLog
@@ -58,6 +73,79 @@ export class RateLimitHandler extends AbstractLogHandler {
     }
   }
 
+  private getCurrentMonth(): string {
+    return new Date().toISOString().slice(0, 7); // "YYYY-MM"
+  }
+
+  private shouldCheckFreeTierLimit(freeLimitExceeded: string | null): boolean {
+    if (freeLimitExceeded === null) {
+      // Not exceeded - check with 1% probability
+      return Math.random() < FREE_TIER_CHECK_PROBABILITY;
+    }
+
+    if (freeLimitExceeded === this.getCurrentMonth()) {
+      // Already exceeded this month - skip check entirely
+      return false;
+    }
+
+    // Old month - check with 1% probability to see if still over limit
+    return Math.random() < FREE_TIER_CHECK_PROBABILITY;
+  }
+
+  private async checkAndUpdateFreeTierLimit(orgId: string): Promise<void> {
+    try {
+      const count = await this.get30DayRequestCount(orgId);
+      if (count >= FREE_TIER_LIMIT) {
+        // Over limit - set to current month
+        await this.setFreeLimitExceeded(orgId, true);
+        console.log(
+          `[FreeTierLimit] Limit exceeded for org ${orgId}: ${count} requests in last 30 days`
+        );
+      } else {
+        // Under limit - clear the flag
+        await this.setFreeLimitExceeded(orgId, false);
+        console.log(
+          `[FreeTierLimit] Limit cleared for org ${orgId}: ${count} requests in last 30 days`
+        );
+      }
+    } catch (error) {
+      // Don't fail the request if the check fails
+      console.error(`Error checking free tier limit for org ${orgId}:`, error);
+    }
+  }
+
+  private async get30DayRequestCount(orgId: string): Promise<number> {
+    const { data, error } = await dbQueryClickhouse<{ count: number }>(
+      `SELECT COUNT(*) as count FROM request_response_rmt
+       WHERE organization_id = {val_0:String}
+       AND request_created_at >= now() - INTERVAL 30 DAY`,
+      [orgId]
+    );
+
+    if (error || !data || data.length === 0) {
+      console.error(`Error getting 30-day request count for org ${orgId}:`, error);
+      return 0;
+    }
+
+    return data[0].count ?? 0;
+  }
+
+  private async setFreeLimitExceeded(
+    orgId: string,
+    exceeded: boolean
+  ): Promise<void> {
+    // Store as month string (e.g., "2026-01") or null
+    const value = exceeded ? this.getCurrentMonth() : null;
+    const { error } = await dbExecute(
+      `UPDATE organization SET free_limit_exceeded = $1 WHERE id = $2`,
+      [value, orgId]
+    );
+
+    if (error) {
+      console.error(`Error setting free_limit_exceeded for org ${orgId}:`, error);
+    }
+  }
+
   public rateLimitEntry(
     orgId: string,
     percentLog: number

valhalla/jawn/src/lib/handlers/ResponseBodyHandler.ts

@@ -79,9 +79,10 @@ export class ResponseBodyHandler extends AbstractLogHandler {
         context.message.log.response.status =
           processedResponseBody.data.statusOverride;
       }
-      context.processedLog.response.model = getModelFromResponse(
-        processedResponseBody.data?.processedBody
-      );
+      // Get model from response body, or fall back to Worker-provided model when body isn't stored
+      context.processedLog.response.model =
+        getModelFromResponse(processedResponseBody.data?.processedBody) ||
+        context.message.log.response.model;
 
       const definedModel =
         calculateModel(
@@ -113,37 +114,55 @@ export class ResponseBodyHandler extends AbstractLogHandler {
       }
 
       // Set legacy usage values captured from body processors
+      // Fall back to Worker-provided tokens when body isn't stored (free tier limit exceeded)
       const legacyUsage =
         processedResponseBody.data?.usage ??
         processedResponseBody.data?.processedBody?.usage ??
         {};
-      context.legacyUsage.completionTokens = legacyUsage.completionTokens;
-      context.legacyUsage.promptTokens = legacyUsage.promptTokens;
-      context.legacyUsage.totalTokens = legacyUsage.totalTokens;
+      context.legacyUsage.completionTokens =
+        legacyUsage.completionTokens ??
+        context.message.log.response.completionTokens;
+      context.legacyUsage.promptTokens =
+        legacyUsage.promptTokens ?? context.message.log.response.promptTokens;
+      context.legacyUsage.totalTokens =
+        legacyUsage.totalTokens ??
+        ((context.legacyUsage.promptTokens ?? 0) +
+          (context.legacyUsage.completionTokens ?? 0) ||
+          undefined);
       context.legacyUsage.heliconeCalculated = legacyUsage.heliconeCalculated;
+      // Fall back to Worker-provided cache/audio tokens when body isn't stored
       context.legacyUsage.promptCacheWriteTokens =
-        legacyUsage.promptCacheWriteTokens;
+        legacyUsage.promptCacheWriteTokens ??
+        context.message.log.response.promptCacheWriteTokens;
       context.legacyUsage.promptCacheReadTokens =
-        legacyUsage.promptCacheReadTokens;
-      context.legacyUsage.promptAudioTokens = legacyUsage.promptAudioTokens;
+        legacyUsage.promptCacheReadTokens ??
+        context.message.log.response.promptCacheReadTokens;
+      context.legacyUsage.promptAudioTokens =
+        legacyUsage.promptAudioTokens ??
+        context.message.log.response.promptAudioTokens;
       context.legacyUsage.completionAudioTokens =
-        legacyUsage.completionAudioTokens;
+        legacyUsage.completionAudioTokens ??
+        context.message.log.response.completionAudioTokens;
+      context.legacyUsage.reasoningTokens =
+        legacyUsage.reasoningTokens ??
+        context.message.log.response.reasoningTokens;
       context.legacyUsage.promptCacheWrite5m = legacyUsage.promptCacheWrite5m;
       context.legacyUsage.promptCacheWrite1h = legacyUsage.promptCacheWrite1h;
       if (typeof legacyUsage.cost === "number" && legacyUsage.cost) {
         context.legacyUsage.cost = legacyUsage.cost;
       } else {
+        // Use context.legacyUsage which has Worker-provided tokens as fallback
         const cost = modelCost({
           model: context.processedLog.model ?? "",
           provider: context.message.log.request.provider ?? "",
-          sum_prompt_tokens: legacyUsage.promptTokens ?? 0,
-          prompt_cache_write_tokens: legacyUsage.promptCacheWriteTokens ?? 0,
-          prompt_cache_read_tokens: legacyUsage.promptCacheReadTokens ?? 0,
-          prompt_audio_tokens: legacyUsage.promptAudioTokens ?? 0,
-          sum_completion_tokens: legacyUsage.completionTokens ?? 0,
-          completion_audio_tokens: legacyUsage.completionAudioTokens ?? 0,
-          prompt_cache_write_5m: legacyUsage.promptCacheWrite5m ?? 0,
-          prompt_cache_write_1h: legacyUsage.promptCacheWrite1h ?? 0,
+          sum_prompt_tokens: context.legacyUsage.promptTokens ?? 0,
+          prompt_cache_write_tokens: context.legacyUsage.promptCacheWriteTokens ?? 0,
+          prompt_cache_read_tokens: context.legacyUsage.promptCacheReadTokens ?? 0,
+          prompt_audio_tokens: context.legacyUsage.promptAudioTokens ?? 0,
+          sum_completion_tokens: context.legacyUsage.completionTokens ?? 0,
+          completion_audio_tokens: context.legacyUsage.completionAudioTokens ?? 0,
+          prompt_cache_write_5m: context.legacyUsage.promptCacheWrite5m ?? 0,
+          prompt_cache_write_1h: context.legacyUsage.promptCacheWrite1h ?? 0,
         });
 
         context.legacyUsage.cost = cost;

valhalla/jawn/src/lib/handlers/S3ReaderHandler.ts

@@ -45,9 +45,16 @@ export class S3ReaderHandler extends AbstractLogHandler {
 
       if (content.error || !content.data) {
         if (content.error?.notFoundErr) {
-          // Not found is unrecoverable, we will have no request/response to log
-          // Do not process further, do not send to DLQ
-          return ok(`Content not found in S3: ${signedUrl.data}`);
+          // Content not found in S3 - this can happen when:
+          // 1. Free tier limit exceeded (bodies not stored)
+          // 2. Omit headers set (bodies not stored)
+          // Continue processing with empty bodies - metadata will still be logged
+          console.log(
+            `S3 content not found for request ${context.message.log.request.id}, continuing with empty bodies`
+          );
+          context.rawLog.rawRequestBody = undefined;
+          context.rawLog.rawResponseBody = undefined;
+          return await super.handle(context);
         }
         return err(
           `Error fetching content from S3: ${JSON.stringify(content.error)}`