♻️ feat: Share Pull Request Lookups Across Replicas Through Redis

api/server/routes/__test-utils__/convos-route-mocks.js

@@ -150,6 +150,7 @@ module.exports = {
     createPullRequestSource: jest.fn(() => ({ find: jest.fn() })),
     createProxyAwareFetch: jest.fn(() => jest.fn()),
     createPullRequestLookup: jest.fn(() => jest.fn()),
+    createPullRequestLookupStore: jest.fn(() => undefined),
     createConversationPullRequestsHandler: jest.fn(() => jest.fn()),
     createConversationPullRequestHandler: jest.fn(
       () => (_req, res) => res.status(200).json({ pullRequest: null }),

api/server/routes/convos.js

@@ -23,6 +23,9 @@ const {
   createPullRequestSource,
   createProxyAwareFetch,
   createPullRequestLookup,
+  createPullRequestLookupStore,
+  keyvRedisClient,
+  cacheConfig,
   createConversationPullRequestsHandler,
   backgroundTaskRegistry,
   createSubagentThreadViewHandler,
@@ -175,6 +178,7 @@ const backgroundTaskCancelHandler = createBackgroundTaskCancelHandler({
 /** One lookup, so the header's single route and the sidebar's batch route share its cache. */
 const pullRequestLookup = createPullRequestLookup({
   source: createPullRequestSource({ fetchFn: createProxyAwareFetch() }),
+  store: createPullRequestLookupStore({ client: keyvRedisClient, config: cacheConfig }),
 });
 const conversationPullRequestHandler = createConversationPullRequestHandler({
   getConvoLaneGit: db.getConvoLaneGit,

librechat.example.yaml

@@ -977,6 +977,15 @@ endpoints:
   #     maxCandidatePullRequests: 10 # pull requests listed per state when matching a reused branch (1 to 100)
   #     maxCandidatePages: 1 # pages of those read per state, for branch names reused more often (1 to 10)
   #     maxHeadComparisons: 3 # candidates compared with the chat's last commit (0 to 20)
+  #     # Share looked-up pull requests and GitHub rate limit cooldowns between replicas, so they
+  #     # do not each spend the token on the same answer. Only takes effect with USE_REDIS. Failed
+  #     # lookups are never shared.
+  #     sharedCache: false
+  #     # Milliseconds one shared cache call may take before a lookup carries on without it (10 to 5000).
+  #     sharedCacheTimeoutMs: 250
+  #     # Redis keys the shared cache uses at most, for results and again for rate limits (64 to
+  #     # 1000000). Read from this file only.
+  #     sharedCacheSlots: 65536
   #   # "run_in_background" makes Code Interpreter tools eligible by default and enables per-tool MCP opt-in.
   #   # "tool_intents" enables live model-written labels for native tools and opted-in MCP tools.
   #   # (optional) Require user approval before matching tool calls. Disabled by default.

packages/api/src/cache/__tests__/pullRequests.cache_integration.spec.ts

@@ -0,0 +1,69 @@
+import type { PullRequestLookupStore } from '~/pulls/types';
+import { closeRedisClients } from './redisClients.helper';
+
+describe('pull request lookup store against Redis', () => {
+  const testPrefix = 'PullRequestStore-Integration-Test';
+  let redisClients: Awaited<typeof import('../redisClients')>;
+  let store: PullRequestLookupStore;
+
+  beforeAll(async () => {
+    process.env.USE_REDIS = process.env.USE_REDIS ?? 'true';
+    process.env.REDIS_URI = process.env.REDIS_URI ?? 'redis://127.0.0.1:6379';
+    process.env.REDIS_KEY_PREFIX = process.env.REDIS_KEY_PREFIX ?? testPrefix;
+    jest.resetModules();
+    redisClients = await import('../redisClients');
+    const { cacheConfig } = await import('../cacheConfig');
+    const { createPullRequestLookupStore } = await import('../pullRequests');
+    await redisClients.keyvRedisClientReady;
+    const created = createPullRequestLookupStore({
+      client: redisClients.keyvRedisClient,
+      config: cacheConfig,
+    });
+    if (created == null) throw new Error('Redis store is not available');
+    store = created;
+  });
+
+  afterEach(async () => {
+    const client = redisClients.keyvRedisClient;
+    if (client == null || !('scanIterator' in client)) return;
+    const { cacheConfig } = await import('../cacheConfig');
+    const match = `${cacheConfig.REDIS_KEY_PREFIX}${cacheConfig.GLOBAL_PREFIX_SEPARATOR}PULL_REQUESTS:*`;
+    for await (const page of client.scanIterator({ MATCH: match })) {
+      await Promise.all(page.map((key: string) => client.del(key)));
+    }
+  });
+
+  afterAll(async () => {
+    await closeRedisClients();
+  });
+
+  it('keeps the later value for one owner whatever order the writes arrive in', async () => {
+    await store.put('slot', { key: 'owner', order: 3_600_000 }, 60_000, 1_000);
+    await store.put('slot', { key: 'owner', order: 10_000 }, 60_000, 1_000);
+    await expect(store.get('slot', 1_000)).resolves.toEqual({ key: 'owner', order: 3_600_000 });
+    await store.put('slot', { key: 'owner', order: 7_200_000 }, 60_000, 1_000);
+    await expect(store.get('slot', 1_000)).resolves.toEqual({ key: 'owner', order: 7_200_000 });
+  });
+
+  it('lets another owner take a slot, as a fixed slot count requires', async () => {
+    await store.put('slot', { key: 'first', order: 3_600_000 }, 60_000, 1_000);
+    await store.put('slot', { key: 'second', order: 10 }, 60_000, 1_000);
+    await expect(store.get('slot', 1_000)).resolves.toEqual({ key: 'second', order: 10 });
+  });
+
+  it('replaces a value it cannot read', async () => {
+    const client = redisClients.keyvRedisClient;
+    if (client == null) throw new Error('Redis client is not initialized');
+    const { cacheConfig } = await import('../cacheConfig');
+    const key = `${cacheConfig.REDIS_KEY_PREFIX}${cacheConfig.GLOBAL_PREFIX_SEPARATOR}PULL_REQUESTS:slot`;
+    await client.set(key, 'not json');
+    await store.put('slot', { key: 'owner', order: 1 }, 60_000, 1_000);
+    await expect(store.get('slot', 1_000)).resolves.toEqual({ key: 'owner', order: 1 });
+  });
+
+  it('expires a value with its lifetime', async () => {
+    await store.put('expiring', { key: 'owner', order: 1 }, 50, 1_000);
+    await new Promise((resolve) => setTimeout(resolve, 120));
+    await expect(store.get('expiring', 1_000)).resolves.toBeUndefined();
+  });
+});

packages/api/src/cache/index.ts

@@ -4,5 +4,6 @@ export * from './keyvFiles';
 export { default as keyvMongo } from './keyvMongo';
 export * from './cacheFactory';
 export * from './principals';
+export * from './pullRequests';
 export * from './redisUtils';
 export * from './redisTelemetry';

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

@@ -0,0 +1,96 @@
+import { createServer } from 'node:net';
+import { CacheKeys } from 'librechat-data-provider';
+import { createClient, TimeoutError } from '@redis/client';
+import type { PullRequestStoreClient } from './pullRequests';
+import type { PullRequestStoreConfig } from './pullRequests';
+import { createPullRequestLookupStore } from './pullRequests';
+
+const base: PullRequestStoreConfig = {
+  USE_REDIS: true,
+  FORCED_IN_MEMORY_CACHE_NAMESPACES: [],
+  REDIS_KEY_PREFIX: '',
+  GLOBAL_PREFIX_SEPARATOR: '::',
+};
+
+function fakeClient() {
+  const values = new Map<string, string>();
+  const commands = {
+    get: jest.fn(async (key: string) => values.get(key) ?? null),
+    eval: jest.fn(async (_script: string, options: { keys: string[]; arguments: string[] }) => {
+      values.set(options.keys[0], options.arguments[0]);
+      return 1;
+    }),
+  };
+  const client = { withCommandOptions: jest.fn(() => commands) };
+  return { client, commands, values };
+}
+
+/** A port nothing listens on, so a client there stays disconnected. */
+async function closedPort(): Promise<number> {
+  const server = createServer();
+  await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
+  const address = server.address();
+  await new Promise<void>((resolve) => server.close(() => resolve()));
+  if (address == null || typeof address === 'string') throw new Error('no port');
+  return address.port;
+}
+
+describe('createPullRequestLookupStore', () => {
+  it('is absent without Redis, so nothing is shared', () => {
+    const config = { ...base, USE_REDIS: false };
+    expect(createPullRequestLookupStore({ client: fakeClient().client, config })).toBeUndefined();
+  });
+
+  it('is absent without a Redis client', () => {
+    expect(createPullRequestLookupStore({ client: null, config: base })).toBeUndefined();
+  });
+
+  it('is absent when the namespace is forced in memory', () => {
+    const config = { ...base, FORCED_IN_MEMORY_CACHE_NAMESPACES: [CacheKeys.PULL_REQUESTS] };
+    expect(createPullRequestLookupStore({ client: fakeClient().client, config })).toBeUndefined();
+  });
+
+  it('round-trips JSON under the namespace with the lifetime and the deadline', async () => {
+    const { client, commands } = fakeClient();
+    const config = { ...base, REDIS_KEY_PREFIX: 'tenant' };
+    const store = createPullRequestLookupStore({ client, config });
+    await store?.put('k', { key: 'owner', order: 2 }, 5_000.4, 40);
+    const key = `tenant::${CacheKeys.PULL_REQUESTS}:k`;
+    expect(commands.eval).toHaveBeenCalledWith(expect.any(String), {
+      keys: [key],
+      arguments: ['{"key":"owner","order":2}', '5001', 'owner', '2'],
+    });
+    expect(client.withCommandOptions).toHaveBeenLastCalledWith({ timeout: 40 });
+    await expect(store?.get('k', 30)).resolves.toEqual({ key: 'owner', order: 2 });
+    expect(client.withCommandOptions).toHaveBeenLastCalledWith({ timeout: 30 });
+    await expect(store?.get('missing', 30)).resolves.toBeUndefined();
+  });
+
+  it('rejects a failed command', async () => {
+    const { client, commands } = fakeClient();
+    commands.get.mockRejectedValueOnce(new Error('down'));
+    const store = createPullRequestLookupStore({ client, config: base });
+    await expect(store?.get('k', 30)).rejects.toThrow('down');
+  });
+
+  it('drops a command queued while Redis is unreachable at its deadline', async () => {
+    const redis = createClient({
+      socket: { host: '127.0.0.1', port: await closedPort(), reconnectStrategy: () => 60_000 },
+    });
+    redis.on('error', () => undefined);
+    /** Never settles: the client keeps waiting to reconnect until it is destroyed. */
+    void redis.connect().catch(() => undefined);
+    try {
+      /** The production client, typed as the store sees it, with no cast. */
+      const client: PullRequestStoreClient = redis;
+      const store = createPullRequestLookupStore({ client, config: base });
+      const started = Date.now();
+      await expect(store?.get('k', 50)).rejects.toBeInstanceOf(TimeoutError);
+      const value = { key: 'owner', order: 1 };
+      await expect(store?.put('k', value, 1_000, 50)).rejects.toBeInstanceOf(TimeoutError);
+      expect(Date.now() - started).toBeLessThan(2_000);
+    } finally {
+      redis.destroy();
+    }
+  });
+});

packages/api/src/cache/pullRequests.ts

@@ -0,0 +1,82 @@
+import { CacheKeys } from 'librechat-data-provider';
+import type { PullRequestLookupStore } from '~/pulls/types';
+import type { cacheConfig } from './cacheConfig';
+
+/** The commands the store sends; the single-node and the cluster client both provide them. */
+type StoreCommands = {
+  get(key: string): Promise<string | null>;
+  eval(script: string, options: { keys: string[]; arguments: string[] }): Promise<unknown>;
+};
+
+/**
+ * Writes ARGV[1] for ARGV[2] ms unless the key holds a value for the same owner (ARGV[3]) with an
+ * order at least ARGV[4]. One script, so the check and the write cannot interleave with another
+ * replica's: a slower writer never replaces a newer answer, a shorter rate limit never a longer.
+ */
+const PUT_SCRIPT = `
+local held = redis.call('GET', KEYS[1])
+if held then
+  local ok, value = pcall(cjson.decode, held)
+  if ok and type(value) == 'table' and value.key == ARGV[3]
+    and tonumber(value.order) and tonumber(value.order) >= tonumber(ARGV[4]) then
+    return 0
+  end
+end
+redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
+return 1
+`;
+
+export type PullRequestStoreClient = {
+  withCommandOptions(options: { timeout: number }): StoreCommands;
+};
+
+/** The deployment's cache settings the store follows: whether Redis is on, and its key prefix. */
+export type PullRequestStoreConfig = Pick<
+  typeof cacheConfig,
+  'USE_REDIS' | 'FORCED_IN_MEMORY_CACHE_NAMESPACES' | 'REDIS_KEY_PREFIX' | 'GLOBAL_PREFIX_SEPARATOR'
+>;
+
+/**
+ * The store pull request lookups share between replicas, or `undefined` when there is nothing to
+ * share it through: without Redis, or with the namespace forced in memory, each replica's own
+ * cache already is the whole cache.
+ *
+ * Every command carries the caller's timeout as its deadline. A command still unsent when it
+ * passes, such as one held in the offline queue while Redis is disconnected, is taken out of the
+ * queue and rejected, so a Redis outage does not pile up commands to replay on reconnect.
+ */
+export function createPullRequestLookupStore({
+  client,
+  config,
+}: {
+  client: PullRequestStoreClient | null;
+  config: PullRequestStoreConfig;
+}): PullRequestLookupStore | undefined {
+  if (
+    client == null ||
+    !config.USE_REDIS ||
+    config.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.PULL_REQUESTS)
+  ) {
+    return undefined;
+  }
+  const prefix = config.REDIS_KEY_PREFIX
+    ? `${config.REDIS_KEY_PREFIX}${config.GLOBAL_PREFIX_SEPARATOR}${CacheKeys.PULL_REQUESTS}:`
+    : `${CacheKeys.PULL_REQUESTS}:`;
+  return {
+    get: async (key, timeoutMs) => {
+      const raw = await client.withCommandOptions({ timeout: timeoutMs }).get(prefix + key);
+      return raw == null ? undefined : JSON.parse(raw);
+    },
+    put: async (key, value, ttlMs, timeoutMs) => {
+      await client.withCommandOptions({ timeout: timeoutMs }).eval(PUT_SCRIPT, {
+        keys: [prefix + key],
+        arguments: [
+          JSON.stringify(value),
+          String(Math.max(1, Math.ceil(ttlMs))),
+          value.key,
+          String(value.order),
+        ],
+      });
+    },
+  };
+}

packages/api/src/pulls/handler.spec.ts

@@ -95,6 +95,9 @@ describe('createConversationPullRequestHandler', () => {
       ttlMs: 30_000,
       cacheMaxEntries: 500,
       cacheMaxCredentials: 256,
+      sharedCache: false,
+      sharedCacheTimeoutMs: 250,
+      sharedCacheSlots: 65_536,
       allowedRepositories: ['o/r'],
       limits: {
         requestTimeoutMs: 10_000,
@@ -1115,6 +1118,43 @@ describe('createConversationPullRequestsHandler', () => {
       expect(hosts).toEqual(['github.com', 'ghe.example.com']);
     });
 
+    it.each([
+      ['sharing setting', { sharedCache: true }],
+      ['shared cache deadline', { sharedCache: true, sharedCacheTimeoutMs: 10 }],
+      ['cache capacity', { cacheMaxEntries: 50 }],
+    ])('does not join a lookup started with another %s', async (_label, changed) => {
+      const inputs: unknown[] = [];
+      const lookup = jest.fn(
+        (input: unknown) =>
+          new Promise(() => {
+            inputs.push(input);
+          }),
+      );
+      const settings = { ...enabled, batchTimeoutSeconds: 1 };
+      const getAppConfig = jest
+        .fn()
+        .mockResolvedValueOnce(configWith(settings))
+        .mockResolvedValueOnce(configWith(settings))
+        .mockResolvedValue(configWith({ ...settings, ...changed }));
+      const handler = createConversationPullRequestsHandler({
+        getConvosLaneGit: jest.fn(async () => [lane('a', 'o/r', 'feat/shared')]),
+        getAppConfig,
+        lookup: lookup as never,
+        env: { GH_TOKEN: 'ghp_secret' },
+      });
+      const call = () =>
+        handler(
+          { user: { id: 'u1' }, body: { conversationIds: ['a'] } } as unknown as ServerRequest,
+          { status: jest.fn().mockReturnThis(), json: jest.fn() } as unknown as Response,
+        );
+      void call();
+      await flush();
+      void call();
+      await flush();
+      expect(inputs).toHaveLength(2);
+      expect(inputs[1]).toMatchObject(changed);
+    });
+
     it('does not join a lookup started for another provider on the same host', async () => {
       const providers: Array<string | undefined> = [];
       const lookup = jest.fn(

packages/api/src/pulls/handler.ts

@@ -229,6 +229,9 @@ function lookupInput(settings: PullRequestSettings, token: string, lane: Eligibl
     ttlMs: (settings.cacheTtlSeconds ?? 30) * 1000,
     cacheMaxEntries: settings.cacheMaxEntries ?? 500,
     cacheMaxCredentials: settings.cacheMaxCredentials ?? 256,
+    sharedCache: settings.sharedCache === true,
+    sharedCacheTimeoutMs: settings.sharedCacheTimeoutMs ?? 250,
+    sharedCacheSlots: settings.sharedCacheSlots ?? 65_536,
     allowedRepositories: repositoryScope(settings),
     limits: {
       requestTimeoutMs: (settings.requestTimeoutSeconds ?? 10) * 1000,
@@ -476,6 +479,11 @@ export function createConversationPullRequestsHandler(deps: {
           input.ttlMs,
           input.limits,
           input.allowedRepositories,
+          /** Where the lookup caches and shares its answer, so each caller's own policy holds. */
+          input.cacheMaxEntries,
+          input.cacheMaxCredentials,
+          input.sharedCache,
+          input.sharedCacheTimeoutMs,
         ])}`;
         const joined = sharedLookups.join(shareKey, scope, limit, input);
         const result = await withDeadline(joined.promise, remaining);