feat(exa): add $10/month free allowance with credit billing for overages (#2169)

* feat(exa): add $10/month free allowance with credit billing for overages

Implement per-user monthly Exa usage tracking with a free tier + overage model:
- First $10/month is free for all authenticated users
- Beyond $10, usage is charged to the user's (or org's) Kilo credit balance
- Streaming disabled to guarantee cost tracking via costDollars in JSON responses
- Two-table storage strategy: pre-aggregated counter (O(1) lookups) + partitioned audit log

* feat(exa): store per-user free allowance on monthly usage row

Add free_allowance_microdollars column to exa_monthly_usage with lock-in
semantics: the allowance is set on the first request of the month (INSERT)
and not overwritten on subsequent requests (excluded from ON CONFLICT UPDATE).

This enables per-user tiered allowances in the future by modifying a single
pure function (getExaFreeAllowanceMicrodollars) without changing the billing
flow. The route handler now uses the stored allowance instead of the global
constant, and the 402 error message reflects the actual allowance.

* feat(exa): route exa API through global backend and use read replica for reads

* fix(exa): make recomputeBalance account for paid Exa usage

Add organization_id to exa_monthly_usage so charged amounts are tracked
per-context (personal vs org). Both recompute functions now include
total_charged_microdollars from exa_monthly_usage in cumulative usage,
preventing paid Exa charges from vanishing on balance recomputation.

Squashes branch migrations 0077+0078 into a single 0077.

* fix(exa): match readDb arg in getBalanceAndOrgSettings assertions

The route now passes readDb as a 3rd argument to getBalanceAndOrgSettings,
but two toHaveBeenCalledWith assertions only expected 2 arguments. The
failing assertion triggers Jest's pretty-printer on the real Drizzle
client (a Proxy-backed object), which hangs indefinitely.

* style: format unformatted files

* fix(exa): use per-request exa_usage_log for balance recomputation

Replace the exa_monthly_usage lump-sum approach with a chronological
merge-sort of exa_usage_log records into the usage stream. This gives
correct credit-expiration baselines when Exa charges are interleaved
with credit grants/expirations.

- Stop dropping old exa_usage_log partitions (retain indefinitely)
- Register exa-partition-maintenance cron in vercel.json
- Promote exa_usage_log insert from fire-and-forget to required
- Recompute functions merge-sort exa_usage_log with microdollar_usage

* chore(db): remove branch-local migration 0077 before merging main

* chore(db): regenerate exa migration as 0078 after merging main

* Prevent redirects to global app in other envs

* fix(web): use VERCEL_ENV instead of NODE_ENV for production rewrite check

* refactor(exa): use date-fns format instead of hand-rolled date helpers

* Remove stupid comment

* refactor: simplify mergeSortedByCreatedAt to concat+sort

* fix(exa): insert usage log before upserting counter to prevent free-request leak

If exa_usage_log insert fails (e.g. missing partition), the counter
was already incremented but deductFromBalance never ran — giving the
user a free paid request with no log row for recompute to recover from.
Reordering so the log insert happens first ensures a failed insert
leaves no side effects, and any later failure is recoverable.

* docs(exa): warn against inserting into microdollar_usage for personal billing

Recompute already picks up personal Exa charges from exa_usage_log,
so a microdollar_usage row would double-count.

* docs(exa): document that free allowance is intentionally per-user across contexts

Org usage counts toward the same free tier as personal usage. Once
exhausted, the charge goes to whichever context makes the request.
This prevents gaming via multiple orgs.

* Remove the exa plans

* chore(db): remove branch-local migration 0078 before merging main

* chore(db): regenerate exa migration as 0078 after merging main

Author: RSO

.plans/fix-recompute-exa-usage.md

@@ -0,0 +1,125 @@
+# Fix recomputeBalance to account for paid Exa usage
+
+## Problem
+
+`deductFromBalance()` in `exa-usage.ts` mutates the cached balance columns directly:
+
+- **Personal**: increments `kilocode_users.microdollars_used`
+- **Org**: calls `ingestOrganizationTokenUsage` which increments `organizations.microdollars_used`
+
+Neither path inserts into `microdollar_usage`. The recompute functions (`recomputeUserBalances.ts:75` and `recomputeOrganizationBalances.ts:57`) rebuild balances exclusively from `microdollar_usage`, so every billed Exa request vanishes the next time recompute runs.
+
+## Approach
+
+Rather than routing Exa through `microdollar_usage` (wrong data shape, pollutes LLM analytics), make the recompute functions also include charged Exa usage from `exa_usage_log`.
+
+`exa_usage_log` already stores `{cost_microdollars, created_at, charged_to_balance, kilo_user_id, organization_id}` — exactly the shape the merge-sort algorithm needs (`{cost, created_at}`).
+
+## Changes
+
+### 1. `recomputeUserBalances.ts` — include Exa charged records in `fetchUserBalanceData()`
+
+Add a second query alongside the existing `microdollar_usage` query:
+
+```ts
+import { exa_usage_log } from '@kilocode/db/schema';
+
+const exaUsageRecords = await db
+  .select({
+    cost: exa_usage_log.cost_microdollars,
+    created_at: exa_usage_log.created_at,
+  })
+  .from(exa_usage_log)
+  .where(
+    and(
+      eq(exa_usage_log.kilo_user_id, userId),
+      eq(exa_usage_log.charged_to_balance, true),
+      isNull(exa_usage_log.organization_id)
+    )
+  )
+  .orderBy(asc(exa_usage_log.created_at));
+```
+
+Then merge-sort the two sorted arrays before returning:
+
+```ts
+const usageRecords = mergeSortedByCreatedAt(llmUsageRecords, exaUsageRecords);
+return { user, usageRecords, creditTransactions };
+```
+
+`computeUserBalanceUpdates` and `applyUserBalanceUpdates` need zero changes — they already work on the generic `{cost, created_at}[]` shape.
+
+Update the docstring postcondition:
+
+```
+- microdollars_used = sum(microdollar_usage) + sum(exa charged usage)
+```
+
+### 2. `recomputeOrganizationBalances.ts` — same pattern
+
+Add a second query for org Exa charged records:
+
+```ts
+const exaUsageRecords = await db
+  .select({
+    cost: exa_usage_log.cost_microdollars,
+    created_at: exa_usage_log.created_at,
+  })
+  .from(exa_usage_log)
+  .where(
+    and(
+      eq(exa_usage_log.organization_id, args.organizationId),
+      eq(exa_usage_log.charged_to_balance, true)
+    )
+  )
+  .orderBy(asc(exa_usage_log.created_at));
+```
+
+Same merge-sort before the baseline computation loop.
+
+### 3. Add a shared `mergeSortedByCreatedAt` helper
+
+A small utility (either in a shared module or inline) that merges two sorted `{cost: number, created_at: string}[]` arrays:
+
+```ts
+function mergeSortedByCreatedAt(
+  a: { cost: number; created_at: string }[],
+  b: { cost: number; created_at: string }[]
+): { cost: number; created_at: string }[] {
+  const result = [];
+  let i = 0,
+    j = 0;
+  while (i < a.length && j < b.length) {
+    if (a[i].created_at <= b[j].created_at) result.push(a[i++]);
+    else result.push(b[j++]);
+  }
+  while (i < a.length) result.push(a[i++]);
+  while (j < b.length) result.push(b[j++]);
+  return result;
+}
+```
+
+Both recompute files can import this. If we don't want a new file, it can be defined as a local function in each file (they're short enough).
+
+### 4. Tests
+
+**`recomputeUserBalances.test.ts`**:
+
+- Add a test that inserts `exa_usage_log` rows with `charged_to_balance = true` alongside normal `microdollar_usage` rows, then verifies `recomputeUserBalances` includes both in `microdollars_used`.
+- Add a pure test for `computeUserBalanceUpdates` with a pre-merged usage array that includes both LLM and Exa records interleaved by time, verifying baselines are computed correctly.
+
+**`recomputeOrganizationBalances.test.ts`**:
+
+- Same pattern — insert Exa charged usage for an org and verify recompute includes it.
+
+## Not in scope
+
+**Reliability of `exa_usage_log` inserts**: The audit log insert is currently fire-and-forget (`try/catch` that swallows errors at `exa-usage.ts:106-118`). If it fails, the balance deduction happens but no log row exists, so recompute would miss it. This is a pre-existing design tradeoff (partition might not exist). The risk is low because:
+
+- Partition maintenance runs monthly and creates partitions ahead of time
+- `exa_monthly_usage` (the counter) is always reliably written and could serve as a cross-check
+
+If we want to tighten this later, the options are:
+
+1. Make log inserts required when `charged_to_balance = true` (rethrow on failure)
+2. Add a cross-check in recompute: compare `sum(exa_usage_log.cost WHERE charged)` vs `sum(exa_monthly_usage.total_charged)` and log a warning on mismatch

apps/web/next.config.mjs

@@ -45,12 +45,16 @@ const nextConfig = {
     // Uses beforeFiles to ensure the rewrite happens BEFORE filesystem routes are checked
     // See: https://nextjs.org/docs/app/api-reference/config/next-config-js/rewrites
     const globalApiRewrites =
-      process.env.GLOBAL_KILO_BACKEND !== 'true'
+      process.env.VERCEL_ENV === 'production' && process.env.GLOBAL_KILO_BACKEND !== 'true'
         ? [
             {
               source: '/api/fim/completions',
               destination: 'https://global-api.kilo.ai/api/fim/completions',
             },
+            {
+              source: '/api/exa/:path*',
+              destination: 'https://global-api.kilo.ai/api/exa/:path*',
+            },
             {
               source: '/api/marketplace/:path*',
               destination: 'https://global-api.kilo.ai/api/marketplace/:path*',

apps/web/src/app/api/cron/exa-partition-maintenance/route.ts

@@ -0,0 +1,59 @@
+import { NextResponse } from 'next/server';
+import { captureException } from '@sentry/nextjs';
+import { db } from '@/lib/drizzle';
+import { CRON_SECRET } from '@/lib/config.server';
+import { sql } from 'drizzle-orm';
+import { format } from 'date-fns';
+
+if (!CRON_SECRET) {
+  throw new Error('CRON_SECRET is not configured in environment variables');
+}
+
+/**
+ * Exa Usage Log Partition Maintenance
+ *
+ * Run monthly. Creates the next two months' partitions (idempotent).
+ * Old partitions are retained indefinitely — the recompute balance
+ * functions depend on the full exa_usage_log history.
+ */
+export async function GET(request: Request) {
+  const authHeader = request.headers.get('authorization');
+  if (authHeader !== `Bearer ${CRON_SECRET}`) {
+    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
+  }
+
+  const now = new Date();
+  const created: string[] = [];
+  const errors: string[] = [];
+
+  // Create partitions for the current month and the next 2 months
+  for (let offset = 0; offset <= 2; offset++) {
+    const target = new Date(now.getFullYear(), now.getMonth() + offset, 1);
+    const nextMonth = new Date(target.getFullYear(), target.getMonth() + 1, 1);
+    const name = `exa_usage_log_${format(target, 'yyyy_MM')}`;
+
+    try {
+      await db.execute(
+        sql.raw(
+          `CREATE TABLE IF NOT EXISTS "${name}" PARTITION OF "exa_usage_log" FOR VALUES FROM ('${format(target, 'yyyy-MM-dd')}') TO ('${format(nextMonth, 'yyyy-MM-dd')}')`
+        )
+      );
+      created.push(name);
+    } catch (error) {
+      const msg = `Failed to create partition ${name}: ${error instanceof Error ? error.message : String(error)}`;
+      console.error(`[exa-partition-maintenance] ${msg}`);
+      captureException(error, { tags: { source: 'exa-partition-maintenance', partition: name } });
+      errors.push(msg);
+    }
+  }
+
+  console.log(
+    `[exa-partition-maintenance] created=[${created.join(', ')}] errors=${errors.length}`
+  );
+
+  return NextResponse.json({
+    success: errors.length === 0,
+    created,
+    errors,
+  });
+}