🎫 perf: Limit Redis Script Caching to Atomic Claims (#16116)

* refactor(redis): import script cache proposal from #16069

Squashed copy of PR #16069 at 7e994194c2, rebased onto current dev without changing the original branch.

Co-authored-by: Marco Beretta <81851188+berry-13@users.noreply.github.com>

* refactor(redis): cache independent claims without queueing streams

---------

Co-authored-by: Lia <lia@librechat.ai>
Co-authored-by: Marco Beretta <81851188+berry-13@users.noreply.github.com>

e2e/playwright.config.redis.ts

@@ -25,5 +25,6 @@ export default defineConfig({
     /thread-fold\.spec\.ts/,
     /tool-approvals\.spec\.ts/,
     /usage\.spec\.ts/,
+    /scenarios\/redis-stream\.spec\.ts/,
   ],
 });

e2e/specs/mock/scenarios/redis-stream.spec.ts

@@ -0,0 +1,45 @@
+import { expect, test } from '@playwright/test';
+import {
+  MOCK_ENDPOINTS,
+  NEW_CHAT_PATH,
+  messagesView,
+  selectMockEndpoint,
+  sendMessageAndWaitForCompletion,
+} from '../helpers';
+
+const ORDERED_PIECE_COUNT = 64;
+
+function orderedPieces(): string[] {
+  return Array.from(
+    { length: ORDERED_PIECE_COUNT },
+    (_, index) => `piece-${String(index).padStart(3, '0')}`,
+  );
+}
+
+test.describe('Redis-backed stream delivery', () => {
+  test('renders a complete ordered reply across a reload @scenario:redis-stream-renders-ordered-reply', async ({
+    page,
+  }) => {
+    test.setTimeout(60000);
+    const label = `redis-${Date.now()}`;
+    const expected = `E2E ordered reply ${label} ${orderedPieces().join(' ')}`;
+
+    await page.goto(NEW_CHAT_PATH, { timeout: 10000 });
+    await selectMockEndpoint(page, MOCK_ENDPOINTS[0]);
+
+    const response = await sendMessageAndWaitForCompletion(page, `E2E_ORDERED_REPLY:${label}`);
+    expect(response.ok()).toBeTruthy();
+
+    const assistantContent = messagesView(page)
+      .locator('.message-render')
+      .last()
+      .locator('.message-content');
+    await expect(assistantContent).toContainText('piece-010', { timeout: 30000 });
+    await expect(assistantContent).toHaveText(expected, { timeout: 30000 });
+
+    await page.reload({ timeout: 10000 });
+    await expect(
+      messagesView(page).locator('.message-render').last().locator('.message-content'),
+    ).toHaveText(expected, { timeout: 30000 });
+  });
+});

packages/api/src/cache/redisScript.spec.ts

@@ -0,0 +1,84 @@
+import Redis from 'ioredis';
+import { createHash } from 'node:crypto';
+import { evalScript } from './redisScript';
+
+function createClient() {
+  const client = new Redis({ lazyConnect: true });
+  const evalsha = jest.spyOn(client, 'evalsha').mockResolvedValue(1);
+  const evalCommand = jest.spyOn(client, 'eval').mockResolvedValue(1);
+  return { client, evalsha, evalCommand };
+}
+
+describe('independent Redis script execution', () => {
+  test('uses the exact SHA and arguments without an EVAL on success', async () => {
+    const { client, evalsha, evalCommand } = createClient();
+    await expect(evalScript(client, 'return ARGV[1]', 1, '{key}', 'value')).resolves.toBe(1);
+    expect(evalsha).toHaveBeenCalledWith(
+      createHash('sha1').update('return ARGV[1]').digest('hex'),
+      1,
+      '{key}',
+      'value',
+    );
+    expect(evalCommand).not.toHaveBeenCalled();
+  });
+
+  test('falls back once for both cold and previously successful scripts', async () => {
+    const { client, evalsha, evalCommand } = createClient();
+    evalsha.mockRejectedValueOnce(new Error('NOSCRIPT No matching script'));
+    await expect(evalScript(client, 'return 1', 0)).resolves.toBe(1);
+    await expect(evalScript(client, 'return 1', 0)).resolves.toBe(1);
+    evalsha.mockRejectedValueOnce(new Error('NOSCRIPT No matching script'));
+    await expect(evalScript(client, 'return 1', 0)).resolves.toBe(1);
+    expect(evalsha).toHaveBeenCalledTimes(3);
+    expect(evalCommand).toHaveBeenCalledTimes(2);
+  });
+
+  test.each([
+    "NOPERM this user has no permissions to run the 'evalsha' command",
+    "ERR unknown command 'evalsha', with args beginning with: 'sha'",
+  ])('memoizes command-level EVAL-only compatibility: %s', async (message) => {
+    const { client, evalsha, evalCommand } = createClient();
+    evalsha.mockRejectedValueOnce(new Error(message));
+    await expect(evalScript(client, 'return 1', 0)).resolves.toBe(1);
+    await expect(evalScript(client, 'return 2', 0)).resolves.toBe(1);
+    expect(evalsha).toHaveBeenCalledTimes(1);
+    expect(evalCommand).toHaveBeenCalledTimes(2);
+  });
+
+  test.each([
+    'READONLY replica cannot accept writes',
+    'Connection is closed',
+    'NOPERM this user has no permissions to access one of the keys used as arguments',
+    'ERR Error running script: NOSCRIPT failure inside script',
+    'ERR Error running script: NOPERM EVALSHA inside script',
+  ])('does not replay an ambiguous or script-runtime failure: %s', async (message) => {
+    const { client, evalsha, evalCommand } = createClient();
+    evalsha.mockRejectedValueOnce(new Error(message));
+    await expect(evalScript(client, 'return 1', 0)).rejects.toThrow(message);
+    expect(evalCommand).not.toHaveBeenCalled();
+  });
+
+  test('propagates a failed EVAL fallback', async () => {
+    const { client, evalsha, evalCommand } = createClient();
+    evalsha.mockRejectedValueOnce(new Error('NOSCRIPT No matching script'));
+    evalCommand.mockRejectedValueOnce(new Error('EVAL failed'));
+    await expect(evalScript(client, 'return 1', 0)).rejects.toThrow('EVAL failed');
+    expect(evalCommand).toHaveBeenCalledTimes(1);
+  });
+
+  test('dispatches concurrent warm calls without waiting for a predecessor', async () => {
+    const { client, evalsha } = createClient();
+    let release!: () => void;
+    evalsha.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          release = () => resolve(1);
+        }),
+    );
+    const first = evalScript(client, 'return 1', 1, '{same}key');
+    await expect(evalScript(client, 'return 1', 1, '{same}key')).resolves.toBe(1);
+    expect(evalsha).toHaveBeenCalledTimes(2);
+    release();
+    await expect(first).resolves.toBe(1);
+  });
+});

packages/api/src/cache/redisScript.ts

@@ -0,0 +1,64 @@
+import { createHash } from 'node:crypto';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import type { Redis, Cluster } from 'ioredis';
+
+export type RedisScriptArg = string | number | Buffer;
+export type RedisScriptResult = string | number | boolean | null | undefined | RedisScriptResult[];
+export type RedisScriptClient = Pick<Redis | Cluster, 'eval' | 'evalsha'>;
+
+const scriptShas = new Map<string, string>();
+const evalOnlyClients = new WeakSet<object>();
+const fallbackContext = new AsyncLocalStorage<boolean>();
+
+function isNoScriptError(error: unknown): boolean {
+  return error instanceof Error && error.message.startsWith('NOSCRIPT ');
+}
+
+function isUnsupportedEvalsha(error: unknown): boolean {
+  if (!(error instanceof Error)) {
+    return false;
+  }
+  return /^(?:NOPERM .*|ERR unknown command )['"]?evalsha['"]?(?:\s|,|$)/i.test(error.message);
+}
+
+export function isEvalshaFallbackInProgress(error: unknown): boolean {
+  return (
+    fallbackContext.getStore() === true && (isNoScriptError(error) || isUnsupportedEvalsha(error))
+  );
+}
+
+/**
+ * For independent operations only: a NOSCRIPT fallback can execute after later commands.
+ * Callers must await prerequisites and results before dependent work. Ordered stream
+ * writes/publications use direct EVAL instead. Only static script bodies belong here.
+ * Network, script-runtime, and other ambiguous failures are never retried by this helper.
+ */
+export async function evalScript(
+  client: RedisScriptClient,
+  script: string,
+  numberOfKeys: number,
+  ...args: RedisScriptArg[]
+): Promise<RedisScriptResult> {
+  if (evalOnlyClients.has(client)) {
+    return (await client.eval(script, numberOfKeys, ...args)) as RedisScriptResult;
+  }
+  let sha = scriptShas.get(script);
+  if (sha == null) {
+    sha = createHash('sha1').update(script).digest('hex');
+    scriptShas.set(script, sha);
+  }
+  try {
+    return (await fallbackContext.run(true, () =>
+      client.evalsha(sha, numberOfKeys, ...args),
+    )) as RedisScriptResult;
+  } catch (error) {
+    const unsupported = isUnsupportedEvalsha(error);
+    if (!isNoScriptError(error) && !unsupported) {
+      throw error;
+    }
+    if (unsupported) {
+      evalOnlyClients.add(client);
+    }
+    return (await client.eval(script, numberOfKeys, ...args)) as RedisScriptResult;
+  }
+}

packages/api/src/cache/redisTelemetry.spec.ts

@@ -11,6 +11,7 @@ import {
   runWithRedisRequestTelemetry,
 } from './redisTelemetry';
 import { isMetricsConfigured, recordRedisOperation } from '~/app/metrics';
+import { evalScript } from './redisScript';
 
 jest.mock('~/app/metrics', () => ({
   isMetricsConfigured: jest.fn(() => false),
@@ -120,6 +121,92 @@ describe('redisTelemetry', () => {
     expect(mockRecordRedisOperation).toHaveBeenCalledTimes(1);
   });
 
+  it('does not count an expected EVALSHA NOSCRIPT miss as a Redis error', async () => {
+    const span = createSpan();
+    const telemetry = createRedisRequestTelemetry(span as unknown as Span);
+    const evalsha = jest.fn().mockRejectedValue(new Error('NOSCRIPT No matching script'));
+    const evalCommand = jest.fn().mockResolvedValue(1);
+    const redis = instrumentIORedisClient(
+      { evalsha, eval: evalCommand },
+      RedisUseCases.GENERATION_STREAM,
+    );
+
+    await runWithRedisRequestTelemetry(telemetry, async () => {
+      await expect(evalScript(redis, 'return 1', 0)).resolves.toBe(1);
+    });
+    finishRedisRequestTelemetry(telemetry);
+
+    expect(telemetry.errors).toBe(0);
+    expect(mockRecordRedisOperation).toHaveBeenCalledWith(
+      'ioredis',
+      RedisUseCases.GENERATION_STREAM,
+      'evalsha',
+      'success',
+      expect.any(Number),
+    );
+    expect(mockRecordRedisOperation).toHaveBeenCalledWith(
+      'ioredis',
+      RedisUseCases.GENERATION_STREAM,
+      'eval',
+      'success',
+      expect.any(Number),
+    );
+  });
+
+  it('counts a key permission failure as a Redis error', async () => {
+    const span = createSpan();
+    const telemetry = createRedisRequestTelemetry(span as unknown as Span);
+    const permissionError = new Error(
+      'NOPERM this user has no permissions to access one of the keys used as arguments',
+    );
+    const evalsha = jest.fn().mockResolvedValueOnce(1).mockRejectedValueOnce(permissionError);
+    const redis = instrumentIORedisClient(
+      { evalsha, eval: jest.fn() },
+      RedisUseCases.GENERATION_STREAM,
+    );
+
+    await runWithRedisRequestTelemetry(telemetry, async () => {
+      await expect(evalScript(redis, 'return 1', 0)).resolves.toBe(1);
+      await expect(evalScript(redis, 'return 1', 0)).rejects.toBe(permissionError);
+    });
+    finishRedisRequestTelemetry(telemetry);
+
+    expect(telemetry.errors).toBe(1);
+    expect(mockRecordRedisOperation).toHaveBeenLastCalledWith(
+      'ioredis',
+      RedisUseCases.GENERATION_STREAM,
+      'evalsha',
+      'error',
+      expect.any(Number),
+    );
+  });
+
+  it('counts a direct EVALSHA miss as a Redis error', async () => {
+    const span = createSpan();
+    const telemetry = createRedisRequestTelemetry(span as unknown as Span);
+    const redis = instrumentIORedisClient(
+      {
+        evalsha: jest.fn().mockRejectedValue(new Error('NOSCRIPT No matching script')),
+        eval: jest.fn(),
+      },
+      RedisUseCases.GENERATION_STREAM,
+    );
+
+    await runWithRedisRequestTelemetry(telemetry, async () => {
+      await expect(redis.evalsha('sha', 0)).rejects.toThrow('NOSCRIPT');
+    });
+    finishRedisRequestTelemetry(telemetry);
+
+    expect(telemetry.errors).toBe(1);
+    expect(mockRecordRedisOperation).toHaveBeenCalledWith(
+      'ioredis',
+      RedisUseCases.GENERATION_STREAM,
+      'evalsha',
+      'error',
+      expect.any(Number),
+    );
+  });
+
   it('records resolved ioredis pipeline command errors', async () => {
     const span = createSpan();
     const telemetry = createRedisRequestTelemetry(span as unknown as Span);
@@ -147,7 +234,6 @@ describe('redisTelemetry', () => {
       expect.any(Number),
     );
   });
-
   it('preserves the ioredis client constructor', () => {
     class FakeRedisClient {}
 

packages/api/src/cache/redisTelemetry.ts

@@ -8,6 +8,7 @@ import {
   type RedisOperationStatus,
 } from '~/app/metrics';
 
+import { isEvalshaFallbackInProgress } from './redisScript';
 const REDIS_CACHE_METHODS = [
   'clear',
   'delete',
@@ -226,7 +227,8 @@ export async function observeRedisOperation<T>(
     }
     return result;
   } catch (error) {
-    status = 'error';
+    status =
+      redisOperation === 'evalsha' && isEvalshaFallbackInProgress(error) ? 'success' : 'error';
     throw error;
   } finally {
     const durationSeconds = Number(process.hrtime.bigint() - startedAt) / 1_000_000_000;

packages/api/src/stream/__tests__/scriptRecovery.stream_integration.spec.ts

@@ -0,0 +1,196 @@
+import { trace } from '@opentelemetry/api';
+import type { Redis, Cluster } from 'ioredis';
+import {
+  createRedisRequestTelemetry,
+  finishRedisRequestTelemetry,
+  instrumentIORedisClient,
+  RedisUseCases,
+  runWithRedisRequestTelemetry,
+} from '~/cache/redisTelemetry';
+import { clearRedisTestPrefix, createRedisTestClient } from './helpers/redis';
+import { RedisJobStore } from '../implementations/RedisJobStore';
+import { evalScript } from '~/cache/redisScript';
+
+const describeRedis = process.env.USE_REDIS === 'true' ? describe : describe.skip;
+
+describeRedis('Redis script cache recovery', () => {
+  const keyPrefix = 'ScriptRecovery-Integration-Test:';
+  let redis: Redis | Cluster;
+  let originalWindow: string | undefined;
+
+  /** Invalidate only server state: the caller must recover without any test-only state reset. */
+  async function invalidateServerScripts(): Promise<void> {
+    const nodes = (redis as Cluster).isCluster
+      ? (redis as Cluster).nodes('master')
+      : [redis as Redis];
+    await Promise.all(nodes.map((node) => node.script('FLUSH')));
+  }
+
+  beforeAll(async () => {
+    originalWindow = process.env.STREAM_DELTA_COALESCE_MS;
+    process.env.STREAM_DELTA_COALESCE_MS = '25';
+    redis = createRedisTestClient(keyPrefix);
+    await redis.connect();
+  });
+
+  afterEach(async () => {
+    await clearRedisTestPrefix(redis, keyPrefix);
+  });
+
+  afterAll(async () => {
+    redis.disconnect();
+    if (originalWindow === undefined) {
+      delete process.env.STREAM_DELTA_COALESCE_MS;
+    } else {
+      process.env.STREAM_DELTA_COALESCE_MS = originalWindow;
+    }
+  });
+
+  test.each(['pending', 'already-flushing'])(
+    'preserves durable order after partial warm-up (%s)',
+    async (mode) => {
+      const store = new RedisJobStore(redis);
+      const streamId = `recovery-${mode}`;
+      const event = (text: string) => ({ event: 'on_message_delta', data: { text } });
+      try {
+        await redis.hset(`stream:{${streamId}}:job`, 'createdAt', '100', 'status', 'running');
+        const warmBatch = store.appendChunk(streamId, event('warm-batch'), 100, undefined, {
+          coalesce: true,
+        });
+        await store.flushPendingAppends(streamId);
+        expect(await warmBatch).toBe(true);
+        expect(await store.appendChunk(streamId, event('warm-direct'), 100)).toBe(true);
+
+        await invalidateServerScripts();
+        expect(await store.appendChunk(streamId, event('reload-direct-only'), 100)).toBe(true);
+        await redis.del(`stream:{${streamId}}:chunks`);
+
+        const delta = store.appendChunk(streamId, event('earlier-delta'), 100, undefined, {
+          coalesce: true,
+        });
+        const flushing =
+          mode === 'already-flushing' ? store.flushPendingAppends(streamId) : Promise.resolve();
+        const control = { event: 'on_pending_action', data: { text: 'later-control' } };
+        const appended = store.appendChunk(streamId, control, 100);
+        expect(await delta).toBe(true);
+        expect(await appended).toBe(true);
+        await flushing;
+
+        const entries = await redis.xrange(`stream:{${streamId}}:chunks`, '-', '+');
+        expect(entries.map(([, fields]) => fields[1])).toEqual([
+          JSON.stringify(event('earlier-delta')),
+          JSON.stringify(control),
+        ]);
+      } finally {
+        await store.destroy();
+      }
+    },
+  );
+
+  test('recovers a previously successful SHA in one fallback without recording an error', async () => {
+    const client = instrumentIORedisClient(redis, RedisUseCases.GENERATION_STREAM);
+    const script = 'return ARGV[1]';
+    const key = '{recovery-metrics}:key';
+    await expect(evalScript(client, script, 1, key, 'warm')).resolves.toBe('warm');
+    await invalidateServerScripts();
+    const evalsha = jest.spyOn(redis, 'evalsha');
+    const evalCommand = jest.spyOn(redis, 'eval');
+    const span = trace.getTracer('script-recovery-test').startSpan('recovery');
+    const telemetry = createRedisRequestTelemetry(span);
+
+    await runWithRedisRequestTelemetry(telemetry, async () => {
+      await expect(evalScript(client, script, 1, key, 'recovered')).resolves.toBe('recovered');
+    });
+    finishRedisRequestTelemetry(telemetry);
+    span.end();
+
+    expect(telemetry.errors).toBe(0);
+    expect(evalsha).toHaveBeenCalledTimes(1);
+    expect(evalCommand).toHaveBeenCalledTimes(1);
+    await expect(evalScript(client, script, 1, key, 'warm-again')).resolves.toBe('warm-again');
+    expect(evalsha).toHaveBeenCalledTimes(2);
+    expect(evalCommand).toHaveBeenCalledTimes(1);
+  });
+
+  test('records a failed EVAL recovery and does not fail an independent concurrent call', async () => {
+    const client = instrumentIORedisClient(redis, RedisUseCases.GENERATION_STREAM);
+    const script =
+      'if ARGV[1] == "fail" then return redis.error_reply("ERR recovery failed") end return ARGV[1]';
+    const key = '{recovery-failure}:key';
+    await expect(evalScript(client, script, 1, key, 'warm')).resolves.toBe('warm');
+    await invalidateServerScripts();
+    const span = trace.getTracer('script-recovery-test').startSpan('failed-recovery');
+    const telemetry = createRedisRequestTelemetry(span);
+
+    await runWithRedisRequestTelemetry(telemetry, async () => {
+      const failed = evalScript(client, script, 1, key, 'fail');
+      const successor = evalScript(client, script, 1, key, 'successor');
+      await expect(failed).rejects.toThrow('ERR recovery failed');
+      await expect(successor).resolves.toBe('successor');
+    });
+    finishRedisRequestTelemetry(telemetry);
+    span.end();
+    expect(telemetry.errors).toBe(1);
+  });
+
+  test.each([false, true])(
+    'atomic claim races retain exactly one winner (cache flushed: %s)',
+    async (flush) => {
+      const store = new RedisJobStore(redis);
+      const claim = (index: number) => ({
+        streamId: `claim-${index}`,
+        conversationId: 'conversation',
+        claimedAt: 100,
+        claimToken: `token-${index}`,
+      });
+      try {
+        await store.claimIdempotencyKey('{claim-race}:warm', claim(0), 60);
+        if (flush) {
+          await invalidateServerScripts();
+        }
+        const key = `{claim-race}:contended-${flush}`;
+        const results = await Promise.all(
+          Array.from({ length: 32 }, (_, index) =>
+            store.claimIdempotencyKey(key, claim(index), 60),
+          ),
+        );
+        const winners = results.filter((result) => result.claimed);
+        expect(winners).toHaveLength(1);
+        expect(winners[0].existing).toEqual(
+          expect.objectContaining({ claimToken: expect.any(String) }),
+        );
+        expect(
+          results.every(
+            (result) => result.existing?.claimToken === winners[0].existing?.claimToken,
+          ),
+        ).toBe(true);
+        expect(await store.getIdempotencyClaim(key)).toEqual(winners[0].existing);
+        await store.releaseIdempotencyKey(key, winners[0].existing);
+        expect(await store.hasIdempotencyKey(key)).toBe(false);
+      } finally {
+        await store.destroy();
+      }
+    },
+  );
+
+  test('dispatches a same-stream burst without a response-dependent queue', async () => {
+    const store = new RedisJobStore(redis);
+    const streamId = 'pipelined-burst';
+    try {
+      await redis.hset(`stream:{${streamId}}:job`, 'createdAt', '100', 'status', 'running');
+      const evalCommand = jest.spyOn(redis, 'eval');
+      const evalsha = jest.spyOn(redis, 'evalsha');
+      const events = Array.from({ length: 32 }, (_, index) => ({ event: 'delta', data: index }));
+      const pending = events.map((event) => store.appendChunk(streamId, event, 100));
+      expect(evalCommand).toHaveBeenCalledTimes(32);
+      expect(evalsha).not.toHaveBeenCalled();
+      expect(await Promise.all(pending)).toEqual(events.map(() => true));
+      const entries = await redis.xrange(`stream:{${streamId}}:chunks`, '-', '+');
+      expect(entries.map(([, fields]) => fields[1])).toEqual(
+        events.map((event) => JSON.stringify(event)),
+      );
+    } finally {
+      await store.destroy();
+    }
+  });
+});

packages/api/src/stream/implementations/RedisJobStore.ts

@@ -49,6 +49,7 @@ import {
 import { instrumentIORedisClient, RedisUseCases } from '~/cache/redisTelemetry';
 import { RecoveredSteerPayloadMismatchError } from '~/stream/SteerRecovery';
 import { createCheckpointNamespace } from '~/stream/checkpoints';
+import { evalScript } from '~/cache/redisScript';
 
 const CLIENT_REQUEST_ID_PATTERN = /^[A-Za-z0-9:_-]{1,128}$/;
 
@@ -2982,7 +2983,10 @@ export class RedisJobStore implements IJobStoreV2 {
     value: IdempotencyClaimValue,
     ttlSeconds: number,
   ): Promise<IdempotencyClaimResult> {
-    const result = await this.redis.eval(
+    // Contenders need one atomic winner, not dispatch-order fairness. Callers await
+    // this claim before dependent writes, so a cache miss may safely delay it.
+    const result = await evalScript(
+      this.redis,
       IDEMPOTENCY_CLAIM_LUA,
       1,
       KEYS.idempotency(key),