fix(execution): give per-user Redis budget keys a fixed window

The per-user byte budget refreshed its TTL on every accepted write, so for
any user who never went a full TTL without writing, the key never expired.
The per-execution data it accounted for kept expiring underneath it, so the
counter accrued bytes Redis had already dropped and drifted toward the
ceiling. On reaching it, every subsequent write for that user was rejected
until they stopped writing for a full TTL — and since a rejected write does
not refresh the TTL, it recovered on its own and then refilled.

User keys now get a fixed window: the TTL is set when the key is created and
never extended. Applied to all five Lua scripts that touch the key so the
writers cannot drift apart. Execution-scoped keys keep sliding — they are
refreshed on the same schedule as the data they account for, so their
counter and bytes stay in step.

Author: waleedlatif1

apps/sim/lib/execution/event-buffer.test.ts

@@ -68,6 +68,10 @@ function isResetScript(script: string): boolean {
   return script.includes('retained_bytes') && script.includes('replayStartEventId')
 }
 
+function countOccurrences(haystack: string, needle: string): number {
+  return haystack.split(needle).length - 1
+}
+
 describe('execution event buffer', () => {
   beforeEach(() => {
     vi.clearAllMocks()
@@ -417,6 +421,37 @@ describe('execution event buffer', () => {
     )
   })
 
+  it('never extends an existing user budget window while flushing events', async () => {
+    let flushScript = ''
+    mockRedis.eval.mockImplementation(async (script: string) => {
+      if (isFlushScript(script)) flushScript = script
+      return [1, false, 0]
+    })
+
+    const writer = createExecutionEventWriter('exec-1', { userId: 'user-1' })
+    await writer.writeTerminal(makeEvent('terminal'), 'complete')
+
+    expect(flushScript).not.toBe('')
+    const userKeyExpires = countOccurrences(flushScript, "redis.call('EXPIRE', KEYS[5]")
+    const userKeyTtlGuards = countOccurrences(flushScript, "redis.call('TTL', KEYS[5]) < 0")
+    expect(userKeyExpires).toBeGreaterThan(0)
+    expect(userKeyTtlGuards).toBe(userKeyExpires)
+  })
+
+  it('keeps sliding the execution budget window, which expires with its own data', async () => {
+    let flushScript = ''
+    mockRedis.eval.mockImplementation(async (script: string) => {
+      if (isFlushScript(script)) flushScript = script
+      return [1, false, 0]
+    })
+
+    const writer = createExecutionEventWriter('exec-1', { userId: 'user-1' })
+    await writer.writeTerminal(makeEvent('terminal'), 'complete')
+
+    expect(countOccurrences(flushScript, "redis.call('TTL', KEYS[4]) < 0")).toBe(0)
+    expect(countOccurrences(flushScript, "redis.call('EXPIRE', KEYS[4]")).toBeGreaterThan(0)
+  })
+
   it('reports pruned replay buffers before reading incomplete events', async () => {
     mockRedis.hgetall.mockResolvedValue({ status: 'active', earliestEventId: '10' })
 

apps/sim/lib/execution/event-buffer.ts

@@ -85,7 +85,9 @@ if net_bytes > 0 then
   redis.call('EXPIRE', KEYS[4], budget_ttl_seconds)
   if #KEYS >= 5 then
     redis.call('INCRBY', KEYS[5], net_bytes)
-    redis.call('EXPIRE', KEYS[5], budget_ttl_seconds)
+    if redis.call('TTL', KEYS[5]) < 0 then
+      redis.call('EXPIRE', KEYS[5], budget_ttl_seconds)
+    end
   end
 elseif net_bytes < 0 then
   local release_bytes = -net_bytes
@@ -99,17 +101,14 @@ elseif net_bytes < 0 then
     local user_next = redis.call('DECRBY', KEYS[5], release_bytes)
     if user_next <= 0 then
       redis.call('DEL', KEYS[5])
-    else
+    elseif redis.call('TTL', KEYS[5]) < 0 then
       redis.call('EXPIRE', KEYS[5], budget_ttl_seconds)
     end
   end
 else
   if redis.call('EXISTS', KEYS[4]) == 1 then
     redis.call('EXPIRE', KEYS[4], budget_ttl_seconds)
   end
-  if #KEYS >= 5 and redis.call('EXISTS', KEYS[5]) == 1 then
-    redis.call('EXPIRE', KEYS[5], budget_ttl_seconds)
-  end
 end
 for i = 9, #ARGV, 2 do
   redis.call('ZADD', KEYS[1], ARGV[i], ARGV[i + 1])
@@ -148,7 +147,7 @@ if retained_bytes > 0 then
     local user_next = redis.call('DECRBY', KEYS[4], retained_bytes)
     if user_next <= 0 then
       redis.call('DEL', KEYS[4])
-    else
+    elseif redis.call('TTL', KEYS[4]) < 0 then
       redis.call('EXPIRE', KEYS[4], tonumber(ARGV[4]))
     end
   end

apps/sim/lib/execution/redis-budget.server.test.ts

@@ -0,0 +1,76 @@
+/**
+ * @vitest-environment node
+ */
+import { describe, expect, it, vi } from 'vitest'
+import {
+  getExecutionRedisBudgetKeys,
+  reserveExecutionRedisBytes,
+} from '@/lib/execution/redis-budget.server'
+
+function countOccurrences(haystack: string, needle: string): number {
+  return haystack.split(needle).length - 1
+}
+
+async function captureReserveScript(userId?: string): Promise<string> {
+  let script = ''
+  const redis = {
+    eval: vi.fn(async (source: string) => {
+      script = source
+      return [1, 'ok', 0, 0]
+    }),
+  }
+
+  await reserveExecutionRedisBytes(redis as never, {
+    executionId: 'exec-1',
+    userId,
+    category: 'event_buffer',
+    operation: 'write_events',
+    bytes: 128,
+  })
+
+  return script
+}
+
+describe('reserveExecutionRedisBytes', () => {
+  it('scopes the reservation to the execution, and to the user when one is known', () => {
+    expect(
+      getExecutionRedisBudgetKeys({
+        executionId: 'exec-1',
+        category: 'event_buffer',
+        operation: 'write_events',
+        bytes: 1,
+      })
+    ).toEqual(['execution:redis-budget:execution:exec-1'])
+
+    expect(
+      getExecutionRedisBudgetKeys({
+        executionId: 'exec-1',
+        userId: 'user-1',
+        category: 'event_buffer',
+        operation: 'write_events',
+        bytes: 1,
+      })
+    ).toEqual(['execution:redis-budget:execution:exec-1', 'execution:redis-budget:user:user-1'])
+  })
+
+  /**
+   * A user key aggregates across every execution that user runs, so extending
+   * its TTL on each write keeps it alive indefinitely while the per-execution
+   * data it accounts for expires underneath it. The window must be fixed.
+   */
+  it('never extends an existing user budget window', async () => {
+    const script = await captureReserveScript('user-1')
+
+    const userKeyExpires = countOccurrences(script, "redis.call('EXPIRE', KEYS[2]")
+    const userKeyTtlGuards = countOccurrences(script, "redis.call('TTL', KEYS[2]) < 0")
+    expect(userKeyExpires).toBeGreaterThan(0)
+    expect(userKeyTtlGuards).toBe(userKeyExpires)
+  })
+
+  it('keeps sliding the execution budget window, which expires with its own data', async () => {
+    const script = await captureReserveScript('user-1')
+
+    expect(countOccurrences(script, "redis.call('TTL', KEYS[1]) < 0")).toBe(0)
+    expect(countOccurrences(script, "redis.call('EXPIRE', KEYS[1]")).toBeGreaterThan(0)
+  })
+})

apps/sim/lib/execution/redis-budget.server.ts

@@ -12,6 +12,22 @@ const MAX_EXECUTION_REDIS_BYTES = 64 * 1024 * 1024
 const MAX_USER_REDIS_BYTES = 256 * 1024 * 1024
 const REDIS_BUDGET_TTL_SECONDS = 60 * 60
 
+/**
+ * Execution and user budget keys expire differently on purpose.
+ *
+ * An execution key accounts for data that is refreshed on the same schedule as
+ * the key itself, so sliding its TTL on every write keeps the counter and the
+ * bytes it represents in step.
+ *
+ * A user key aggregates across every execution that user runs. Sliding its TTL
+ * on each write keeps it alive indefinitely for any user who stays active,
+ * while the per-execution data it accounts for keeps expiring underneath it —
+ * so the counter accrues bytes Redis has already dropped and eventually pins
+ * the user at their ceiling until they go a full TTL without writing. User
+ * keys therefore get a fixed window: the TTL is set when the key is created
+ * and never extended.
+ */
+
 const RESERVE_REDIS_BYTES_SCRIPT = `
 local bytes = tonumber(ARGV[1])
 local execution_limit = tonumber(ARGV[2])
@@ -32,7 +48,9 @@ redis.call('INCRBY', KEYS[1], bytes)
 redis.call('EXPIRE', KEYS[1], ttl_seconds)
 if #KEYS >= 2 then
   redis.call('INCRBY', KEYS[2], bytes)
-  redis.call('EXPIRE', KEYS[2], ttl_seconds)
+  if redis.call('TTL', KEYS[2]) < 0 then
+    redis.call('EXPIRE', KEYS[2], ttl_seconds)
+  end
 end
 return {1, 'ok', execution_current + bytes, user_current + bytes}
 `

apps/sim/lib/uploads/utils/user-file-base64.server.ts

@@ -59,7 +59,7 @@ if bytes and bytes > 0 then
     local user_next = redis.call('DECRBY', KEYS[4], bytes)
     if user_next <= 0 then
       redis.call('DEL', KEYS[4])
-    else
+    elseif redis.call('TTL', KEYS[4]) < 0 then
       redis.call('EXPIRE', KEYS[4], budget_ttl_seconds)
     end
   end
@@ -129,7 +129,7 @@ if #KEYS >= 4 then
       redis.call('DEL', KEYS[4])
     end
   end
-  if redis.call('EXISTS', KEYS[4]) == 1 then
+  if redis.call('EXISTS', KEYS[4]) == 1 and redis.call('TTL', KEYS[4]) < 0 then
     redis.call('EXPIRE', KEYS[4], budget_ttl_seconds)
   end
 end