From 014e232542b47bd5aa2724b33825759329a36c41 Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:27:40 +0200 Subject: [PATCH 1/5] feat: Share Pull Request Lookups Across Replicas Through Redis --- .../__test-utils__/convos-route-mocks.js | 1 + api/server/routes/convos.js | 2 + librechat.example.yaml | 5 + packages/api/src/cache/index.ts | 1 + packages/api/src/cache/pullRequests.spec.ts | 49 +++++ packages/api/src/cache/pullRequests.ts | 28 +++ packages/api/src/pulls/handler.spec.ts | 2 + packages/api/src/pulls/handler.ts | 2 + packages/api/src/pulls/lookup.spec.ts | 207 +++++++++++++++++- packages/api/src/pulls/lookup.ts | 166 ++++++++++++-- packages/api/src/pulls/types.ts | 14 ++ packages/data-provider/src/config.spec.ts | 13 ++ packages/data-provider/src/config.ts | 12 + 13 files changed, 481 insertions(+), 21 deletions(-) create mode 100644 packages/api/src/cache/pullRequests.spec.ts create mode 100644 packages/api/src/cache/pullRequests.ts diff --git a/api/server/routes/__test-utils__/convos-route-mocks.js b/api/server/routes/__test-utils__/convos-route-mocks.js index 4b9e7a6df38..06eae310d8c 100644 --- a/api/server/routes/__test-utils__/convos-route-mocks.js +++ b/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 }), diff --git a/api/server/routes/convos.js b/api/server/routes/convos.js index 795d983f593..147f2d07bc5 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -23,6 +23,7 @@ const { createPullRequestSource, createProxyAwareFetch, createPullRequestLookup, + createPullRequestLookupStore, createConversationPullRequestsHandler, backgroundTaskRegistry, createSubagentThreadViewHandler, @@ -175,6 +176,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(), }); const conversationPullRequestHandler = createConversationPullRequestHandler({ getConvoLaneGit: db.getConvoLaneGit, diff --git a/librechat.example.yaml b/librechat.example.yaml index 884fdcac373..cb5359cb42d 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -977,6 +977,11 @@ 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. + # sharedCache: false + # # Milliseconds one shared cache call may take before a lookup carries on without it (10 to 5000). + # sharedCacheTimeoutMs: 250 # # "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. diff --git a/packages/api/src/cache/index.ts b/packages/api/src/cache/index.ts index a390f878098..473be450b1c 100644 --- a/packages/api/src/cache/index.ts +++ b/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'; diff --git a/packages/api/src/cache/pullRequests.spec.ts b/packages/api/src/cache/pullRequests.spec.ts new file mode 100644 index 00000000000..3193811d28f --- /dev/null +++ b/packages/api/src/cache/pullRequests.spec.ts @@ -0,0 +1,49 @@ +import { CacheKeys } from 'librechat-data-provider'; +import { createPullRequestLookupStore } from './pullRequests'; +import { standardCache } from './cacheFactory'; +import { cacheConfig } from './cacheConfig'; + +jest.mock('./cacheConfig', () => ({ + cacheConfig: { USE_REDIS: true, FORCED_IN_MEMORY_CACHE_NAMESPACES: [] as string[] }, +})); +jest.mock('./cacheFactory', () => ({ standardCache: jest.fn() })); + +const config = cacheConfig as { USE_REDIS: boolean; FORCED_IN_MEMORY_CACHE_NAMESPACES: string[] }; + +describe('createPullRequestLookupStore', () => { + beforeEach(() => { + config.USE_REDIS = true; + config.FORCED_IN_MEMORY_CACHE_NAMESPACES = []; + jest.mocked(standardCache).mockReset(); + }); + + it('is absent without Redis, so nothing is shared', () => { + config.USE_REDIS = false; + expect(createPullRequestLookupStore()).toBeUndefined(); + expect(standardCache).not.toHaveBeenCalled(); + }); + + it('is absent when the namespace is forced in memory', () => { + config.FORCED_IN_MEMORY_CACHE_NAMESPACES = [CacheKeys.PULL_REQUESTS]; + expect(createPullRequestLookupStore()).toBeUndefined(); + }); + + it('reads and writes the Redis cache with the given lifetime, and rejects its failures', async () => { + const cache = { + get: jest.fn().mockResolvedValue({ until: 1 }), + set: jest.fn().mockResolvedValue(true), + }; + jest + .mocked(standardCache) + .mockReturnValue(cache as unknown as ReturnType); + const store = createPullRequestLookupStore(); + expect(standardCache).toHaveBeenCalledWith(CacheKeys.PULL_REQUESTS, undefined, undefined, { + throwOnErrors: true, + }); + await expect(store?.get('k')).resolves.toEqual({ until: 1 }); + await expect(store?.set('k', { until: 2 }, 5_000)).resolves.toBeUndefined(); + expect(cache.set).toHaveBeenCalledWith('k', { until: 2 }, 5_000); + cache.get.mockRejectedValueOnce(new Error('down')); + await expect(store?.get('k')).rejects.toThrow('down'); + }); +}); diff --git a/packages/api/src/cache/pullRequests.ts b/packages/api/src/cache/pullRequests.ts new file mode 100644 index 00000000000..5ad3aa9a192 --- /dev/null +++ b/packages/api/src/cache/pullRequests.ts @@ -0,0 +1,28 @@ +import { CacheKeys } from 'librechat-data-provider'; +import type { PullRequestLookupStore } from '~/pulls/types'; +import { standardCache } from './cacheFactory'; +import { cacheConfig } from './cacheConfig'; + +/** + * 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. + */ +export function createPullRequestLookupStore(): PullRequestLookupStore | undefined { + if ( + !cacheConfig.USE_REDIS || + cacheConfig.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.PULL_REQUESTS) + ) { + return undefined; + } + /** Rejects instead of resolving a failed call as a miss, so the lookup can log it. */ + const cache = standardCache(CacheKeys.PULL_REQUESTS, undefined, undefined, { + throwOnErrors: true, + }); + return { + get: (key) => cache.get(key), + set: async (key, value, ttlMs) => { + await cache.set(key, value, ttlMs); + }, + }; +} diff --git a/packages/api/src/pulls/handler.spec.ts b/packages/api/src/pulls/handler.spec.ts index 80c28de9534..6adc31a85d9 100644 --- a/packages/api/src/pulls/handler.spec.ts +++ b/packages/api/src/pulls/handler.spec.ts @@ -95,6 +95,8 @@ describe('createConversationPullRequestHandler', () => { ttlMs: 30_000, cacheMaxEntries: 500, cacheMaxCredentials: 256, + sharedCache: false, + sharedCacheTimeoutMs: 250, allowedRepositories: ['o/r'], limits: { requestTimeoutMs: 10_000, diff --git a/packages/api/src/pulls/handler.ts b/packages/api/src/pulls/handler.ts index 4f7a190dc26..f0f18492f59 100644 --- a/packages/api/src/pulls/handler.ts +++ b/packages/api/src/pulls/handler.ts @@ -229,6 +229,8 @@ 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, allowedRepositories: repositoryScope(settings), limits: { requestTimeoutMs: (settings.requestTimeoutSeconds ?? 10) * 1000, diff --git a/packages/api/src/pulls/lookup.spec.ts b/packages/api/src/pulls/lookup.spec.ts index 76bd6e409c5..940c3382727 100644 --- a/packages/api/src/pulls/lookup.spec.ts +++ b/packages/api/src/pulls/lookup.spec.ts @@ -1,5 +1,5 @@ import type { TConversationPullRequest } from 'librechat-data-provider'; -import type { PullRequestSource } from './types'; +import type { PullRequestLookupStore, PullRequestSource } from './types'; import { createPullRequestLookup } from './lookup'; import { PullRequestSourceError } from './types'; @@ -475,3 +475,208 @@ describe('the repositories the credential may be used for', () => { expect(find).toHaveBeenCalledTimes(2); }); }); + +describe('the shared store', () => { + const token = 'ghp_secret_token'; + const shared = { ...input, token, sharedCache: true }; + + /** A Map with lifetimes, standing in for Redis between two replicas. */ + function createStore(clock: () => number) { + const values = new Map(); + const store = { + get: jest.fn(async (key: string) => { + const held = values.get(key); + return held != null && held.until > clock() + ? JSON.parse(JSON.stringify(held.value)) + : undefined; + }), + set: jest.fn(async (key: string, value: unknown, ttlMs: number) => { + values.set(key, { value, until: clock() + ttlMs }); + }), + }; + return { store, values }; + } + + /** Background writes settle on the next turns of the event loop. */ + const settle = () => new Promise((resolve) => setImmediate(resolve)); + + it("serves one replica's answer to another without asking GitHub again", async () => { + const { store } = createStore(Date.now); + const first = jest.fn().mockResolvedValue(value); + const second = jest.fn().mockResolvedValue(value); + const a = createPullRequestLookup({ source: { find: first }, store }); + const b = createPullRequestLookup({ source: { find: second }, store }); + await expect(a(shared)).resolves.toEqual({ ok: true, value }); + await settle(); + await expect(b(shared)).resolves.toEqual({ ok: true, value }); + expect(second).not.toHaveBeenCalled(); + /** Now in the second replica's own cache, so the store is not asked again. */ + store.get.mockClear(); + await b(shared); + expect(store.get).not.toHaveBeenCalled(); + }); + + it('keeps a shared answer no longer than its lifetime', async () => { + let clock = 0; + const { store } = createStore(() => clock); + const a = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + now: () => clock, + }); + const find = jest.fn().mockResolvedValue(null); + const b = createPullRequestLookup({ source: { find }, store, now: () => clock }); + await a(shared); + await settle(); + clock = 30_001; + await expect(b(shared)).resolves.toEqual({ ok: true, value: null }); + expect(find).toHaveBeenCalledTimes(1); + }); + + it('puts every replica into the cooldown one of them was rate limited into', async () => { + const { store } = createStore(Date.now); + const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); + const find = jest.fn().mockResolvedValue(value); + const a = createPullRequestLookup({ source: { find: limited }, store }); + const b = createPullRequestLookup({ source: { find }, store }); + await a(shared); + await settle(); + /** Another branch, so it is the cooldown and not a cached failure that answers. */ + await expect(b({ ...shared, branch: 'other' })).resolves.toEqual({ + ok: false, + error: { code: 'RATE_LIMITED' }, + }); + expect(find).not.toHaveBeenCalled(); + const cooldown = store.set.mock.calls.find(([key]) => key.includes('cooldown')); + expect(cooldown?.[2]).toBe(60_000); + }); + + it('does not share a cooldown with lookups on another host', async () => { + const { store } = createStore(Date.now); + const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); + const find = jest.fn().mockResolvedValue(value); + const a = createPullRequestLookup({ source: { find: limited }, store }); + const b = createPullRequestLookup({ source: { find }, store }); + await a(shared); + await settle(); + await expect(b({ ...shared, host: 'ghe.example.com' })).resolves.toEqual({ ok: true, value }); + }); + + it('never puts the token in a stored key or value', async () => { + const { store } = createStore(Date.now); + const find = jest + .fn() + .mockResolvedValueOnce(value) + .mockRejectedValueOnce(new PullRequestSourceError('RATE_LIMITED')); + const lookup = createPullRequestLookup({ source: { find }, store }); + await lookup(shared); + await lookup({ ...shared, branch: 'b'.repeat(5_000) }); + await settle(); + expect(store.set).toHaveBeenCalledTimes(3); + for (const [key, stored] of store.set.mock.calls) { + expect(key).toMatch(/^pr:v1:/); + expect(key.length).toBeLessThan(100); + expect(key).not.toContain(token); + expect(JSON.stringify(stored)).not.toContain(token); + } + }); + + it('answers from GitHub when the store fails', async () => { + const store: PullRequestLookupStore = { + get: jest.fn().mockRejectedValue(new Error('ECONNREFUSED')), + set: jest.fn().mockRejectedValue(new Error('ECONNREFUSED')), + }; + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ source: { find }, store }); + await expect(lookup(shared)).resolves.toEqual({ ok: true, value }); + await settle(); + await lookup(shared); + expect(find).toHaveBeenCalledTimes(1); + }); + + it('does not wait on a store that never answers for longer than its timeout', async () => { + const hang = () => new Promise(() => undefined); + const store: PullRequestLookupStore = { get: jest.fn(hang), set: jest.fn(hang) }; + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ source: { find }, store }); + const started = Date.now(); + await expect(lookup({ ...shared, sharedCacheTimeoutMs: 20 })).resolves.toEqual({ + ok: true, + value, + }); + expect(Date.now() - started).toBeLessThan(1_000); + expect(find).toHaveBeenCalledTimes(1); + }); + + it.each([ + ['nothing', undefined], + ['a string', 'cached'], + ['an entry without an expiry', { result: { ok: true, value } }], + ['an expiry that is not a number', { result: { ok: true, value }, expiresAt: 'later' }], + ['a result without ok', { result: { value }, expiresAt: Number.MAX_SAFE_INTEGER }], + [ + 'an unknown error code', + { result: { ok: false, error: { code: 'TEAPOT' } }, expiresAt: Number.MAX_SAFE_INTEGER }, + ], + [ + 'a pull request missing fields', + { result: { ok: true, value: { number: 1 } }, expiresAt: Number.MAX_SAFE_INTEGER }, + ], + [ + 'a pull request with an unknown state', + { + result: { ok: true, value: { ...value, state: 'draft' } }, + expiresAt: Number.MAX_SAFE_INTEGER, + }, + ], + [ + 'a pull request page that is not https', + { + result: { ok: true, value: { ...value, url: 'javascript:alert(1)' } }, + expiresAt: Number.MAX_SAFE_INTEGER, + }, + ], + [ + 'a pull request from an unknown provider', + { + result: { ok: true, value: { ...value, provider: 'gitea' } }, + expiresAt: Number.MAX_SAFE_INTEGER, + }, + ], + ['an entry that has expired', { result: { ok: true, value }, expiresAt: 0 }], + ])('ignores %s in the store', async (_label, stored) => { + const store: PullRequestLookupStore = { + get: jest.fn(async (key: string) => (key.includes('cooldown') ? { until: 'soon' } : stored)), + set: jest.fn(async () => undefined), + }; + const find = jest.fn().mockResolvedValue(null); + const lookup = createPullRequestLookup({ source: { find }, store }); + await expect(lookup(shared)).resolves.toEqual({ ok: true, value: null }); + expect(find).toHaveBeenCalledTimes(1); + }); + + it('is never touched when the shared cache is off', async () => { + const { store } = createStore(Date.now); + const find = jest + .fn() + .mockResolvedValueOnce(value) + .mockRejectedValueOnce(new PullRequestSourceError('RATE_LIMITED')); + const lookup = createPullRequestLookup({ source: { find }, store }); + await lookup(input); + await lookup({ ...input, sharedCache: false, branch: 'b' }); + await settle(); + expect(store.get).not.toHaveBeenCalled(); + expect(store.set).not.toHaveBeenCalled(); + }); + + it('does not split the cache by whether it is shared', async () => { + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ + source: { find }, + store: createStore(Date.now).store, + }); + await lookup(input); + await lookup({ ...input, sharedCache: true }); + expect(find).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/api/src/pulls/lookup.ts b/packages/api/src/pulls/lookup.ts index 8e98f5622d8..53e150e0894 100644 --- a/packages/api/src/pulls/lookup.ts +++ b/packages/api/src/pulls/lookup.ts @@ -1,7 +1,18 @@ import { createHash } from 'node:crypto'; import { logger } from '@librechat/data-schemas'; -import type { PullRequestLookup, PullRequestLookupResult, PullRequestSource } from './types'; -import { PullRequestSourceError } from './types'; +import { + PULL_REQUEST_CHECKS, + PULL_REQUEST_STATES, + PULL_REQUEST_MERGEABLE, + PULL_REQUEST_PROVIDERS, +} from 'librechat-data-provider'; +import type { + PullRequestSource, + PullRequestLookup, + PullRequestLookupStore, + PullRequestLookupResult, +} from './types'; +import { PULL_REQUEST_ERROR_CODES, PullRequestSourceError } from './types'; import { getSafeErrorMetadata } from '~/utils'; /** A failed lookup is remembered briefly so a failing upstream is not hammered by every poll. */ @@ -13,13 +24,83 @@ const MAX_COOLDOWN_MS = 60 * 60_000; const DEFAULT_MAX_ENTRIES = 500; /** Credentials with a cache partition at once unless configured; the least recently used goes. */ const DEFAULT_MAX_PARTITIONS = 256; +/** A shared store call slower than this is abandoned; Redis answers in single milliseconds. */ +const DEFAULT_SHARED_TIMEOUT_MS = 250; +/** Versioned, so a change to the stored shape is a miss for replicas still on the old one. */ +const STORE_PREFIX = 'pr:v1:'; type Entry = { result: PullRequestLookupResult; expiresAt: number }; +type Cooldown = { until: number }; /** Keys carry a digest, never the token, so no credential is held in a map key or a heap dump. */ const scopeOf = (token: string): string => createHash('sha256').update(token).digest('hex').slice(0, 32); +/** Repository and branch are user-influenced and unbounded, so the stored key is a fixed size. */ +const storeKeyOf = (key: string): string => + `${STORE_PREFIX}result:${createHash('sha256').update(key).digest('hex')}`; + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null && !Array.isArray(value); + +const isOneOf = (options: readonly T[], value: unknown): value is T => + typeof value === 'string' && (options as readonly string[]).includes(value); + +const isTime = (value: unknown): value is number => + typeof value === 'number' && Number.isFinite(value); + +/** Another replica, or another version of one, wrote this; anything not shaped right is a miss. */ +function isEntry(value: unknown): value is Entry { + if (!isRecord(value) || !isTime(value.expiresAt) || !isRecord(value.result)) return false; + const { result } = value; + if (result.ok === false) { + return isRecord(result.error) && isOneOf(PULL_REQUEST_ERROR_CODES, result.error.code); + } + if (result.ok !== true) return false; + const pr = result.value; + if (pr === null) return true; + return ( + isRecord(pr) && + Number.isInteger(pr.number) && + typeof pr.title === 'string' && + typeof pr.url === 'string' && + /** The page becomes a link in the card, so only https ever comes back out of the store. */ + pr.url.startsWith('https://') && + (pr.provider === undefined || isOneOf(PULL_REQUEST_PROVIDERS, pr.provider)) && + isTime(pr.additions) && + isTime(pr.deletions) && + typeof pr.isDraft === 'boolean' && + isOneOf(PULL_REQUEST_STATES, pr.state) && + isOneOf(PULL_REQUEST_MERGEABLE, pr.mergeable) && + isOneOf(PULL_REQUEST_CHECKS, pr.checks) + ); +} + +const isCooldown = (value: unknown): value is Cooldown => isRecord(value) && isTime(value.until); + +/** + * Runs one shared store call, bounded by `timeoutMs`. The store is an optimization, so a failure + * or a slow answer is logged and reads as `undefined`, and the lookup carries on in memory. + */ +async function bounded(call: () => Promise, timeoutMs: number): Promise { + let timer: NodeJS.Timeout | undefined; + const timeout = new Promise((resolve) => { + timer = setTimeout(() => { + logger.debug('[PullRequests] Shared cache timed out', { timeoutMs }); + resolve(undefined); + }, timeoutMs); + timer.unref?.(); + }); + try { + return await Promise.race([call(), timeout]); + } catch (error) { + logger.debug('[PullRequests] Shared cache unavailable', getSafeErrorMetadata(error)); + return undefined; + } finally { + clearTimeout(timer); + } +} + /** * Caches lookups per credential, repository, branch, recorded commit and lookup policy, and shares * one in-flight request between concurrent callers with all of those in common. Each credential @@ -28,13 +109,19 @@ const scopeOf = (token: string): string => * scopes both the cache and the in-flight map. A rate limit applies to the credential, not the * branch, so it starts a cooldown for every branch under that credential. Callers must have * authorized the conversation before asking. + * + * With a `store` and `sharedCache` on the input, results and cooldowns are also written to the + * store and read from it on a local miss, so replicas share them. In-flight coalescing stays per + * process, and the in-memory cache stays in front of the store. */ export function createPullRequestLookup({ source, + store, now = Date.now, maxEntries = DEFAULT_MAX_ENTRIES, }: { source: PullRequestSource; + store?: PullRequestLookupStore; now?: () => number; maxEntries?: number; }): PullRequestLookup { @@ -80,17 +167,15 @@ export function createPullRequestLookup({ function remember( scope: string, key: string, - result: PullRequestLookupResult, - ttlMs: number, + entry: Entry, capacity: number, credentials: number, ): void { - const lifetime = result.ok ? ttlMs : Math.min(ttlMs, FAILURE_TTL_MS); const entries = partitionOf(scope, credentials); const bound = Math.max(capacities.get(scope) ?? 0, capacity); capacities.set(scope, bound); entries.delete(key); - entries.set(key, { result, expiresAt: now() + lifetime }); + entries.set(key, entry); while (entries.size > bound) { const oldest = entries.keys().next(); if (oldest.done) break; @@ -98,14 +183,25 @@ export function createPullRequestLookup({ } } - function startCooldown(scope: string, retryAfterMs?: number): void { - const wait = Math.min(MAX_COOLDOWN_MS, Math.max(MIN_COOLDOWN_MS, retryAfterMs ?? 0)); - cooldowns.set(scope, now() + wait); - for (const [other, until] of cooldowns) { - if (until <= now()) cooldowns.delete(other); + function startCooldown(scope: string, until: number): void { + cooldowns.set(scope, until); + for (const [other, end] of cooldowns) { + if (end <= now()) cooldowns.delete(other); } } + /** Writes in the background: the answer is already known, so the caller never waits on it. */ + function share( + target: PullRequestLookupStore, + key: string, + value: Entry | Cooldown, + ttlMs: number, + timeoutMs: number, + ): void { + if (ttlMs <= 0) return; + void bounded(() => target.set(key, value, ttlMs), timeoutMs); + } + return async ({ repo, branch, @@ -119,6 +215,8 @@ export function createPullRequestLookup({ allowedRepositories, cacheMaxEntries, cacheMaxCredentials, + sharedCache, + sharedCacheTimeoutMs = DEFAULT_SHARED_TIMEOUT_MS, }) => { const scope = scopeOf(token); /** Everything that changes the answer or how long it may be reused is part of the key. */ @@ -149,7 +247,32 @@ export function createPullRequestLookup({ cooldowns.delete(limitScope); } + /** Only when this caller's settings ask for it; the store is there whenever Redis is. */ + const shared = sharedCache === true ? store : undefined; + const lifetimeOf = (result: PullRequestLookupResult): number => + result.ok ? ttlMs : Math.min(ttlMs, FAILURE_TTL_MS); + const capacity = cacheMaxEntries ?? maxEntries; + const credentials = cacheMaxCredentials ?? DEFAULT_MAX_PARTITIONS; + const storeKey = storeKeyOf(key); + const cooldownKey = `${STORE_PREFIX}cooldown:${createHash('sha256').update(limitScope).digest('hex')}`; + const run = (async (): Promise => { + if (shared != null) { + const [stored, cooldown] = await Promise.all([ + bounded(() => shared.get(storeKey), sharedCacheTimeoutMs), + bounded(() => shared.get(cooldownKey), sharedCacheTimeoutMs), + ]); + if (isEntry(stored) && stored.expiresAt > now()) { + /** Never kept longer than this replica would have kept its own answer. */ + const expiresAt = Math.min(stored.expiresAt, now() + lifetimeOf(stored.result)); + remember(scope, key, { result: stored.result, expiresAt }, capacity, credentials); + return stored.result; + } + if (isCooldown(cooldown) && cooldown.until > now()) { + startCooldown(limitScope, Math.min(cooldown.until, now() + MAX_COOLDOWN_MS)); + return rateLimited; + } + } let result: PullRequestLookupResult; try { result = { @@ -172,18 +295,21 @@ export function createPullRequestLookup({ } const code = error instanceof PullRequestSourceError ? error.code : 'UPSTREAM_ERROR'; if (error instanceof PullRequestSourceError && error.code === 'RATE_LIMITED') { - startCooldown(limitScope, error.retryAfterMs); + const wait = Math.min( + MAX_COOLDOWN_MS, + Math.max(MIN_COOLDOWN_MS, error.retryAfterMs ?? 0), + ); + startCooldown(limitScope, now() + wait); + if (shared != null) { + share(shared, cooldownKey, { until: now() + wait }, wait, sharedCacheTimeoutMs); + } } result = { ok: false, error: { code } }; } - remember( - scope, - key, - result, - ttlMs, - cacheMaxEntries ?? maxEntries, - cacheMaxCredentials ?? DEFAULT_MAX_PARTITIONS, - ); + const lifetime = lifetimeOf(result); + const entry: Entry = { result, expiresAt: now() + lifetime }; + remember(scope, key, entry, capacity, credentials); + if (shared != null) share(shared, storeKey, entry, lifetime, sharedCacheTimeoutMs); return result; })().finally(() => inflight.delete(key)); inflight.set(key, run); diff --git a/packages/api/src/pulls/types.ts b/packages/api/src/pulls/types.ts index e0755adcec6..c5b5f07e732 100644 --- a/packages/api/src/pulls/types.ts +++ b/packages/api/src/pulls/types.ts @@ -75,6 +75,20 @@ export type PullRequestLookupInput = PullRequestFindInput & { cacheMaxEntries?: number; /** Credentials cached at once; the lookup's own default if omitted. */ cacheMaxCredentials?: number; + /** Read and write the lookup's shared store, when it has one. It never changes the answer, so it + * is not part of any cache key. */ + sharedCache?: boolean; + /** Longest one shared store call may take before the lookup carries on without it. */ + sharedCacheTimeoutMs?: number; +}; + +/** + * Storage the lookup shares with other replicas, so one answer or one rate limit from GitHub serves + * every replica. Values must round-trip as JSON; `ttlMs` is how long the store may keep one. + */ +export type PullRequestLookupStore = { + get(key: string): Promise; + set(key: string, value: unknown, ttlMs: number): Promise; }; export type PullRequestLookup = (input: PullRequestLookupInput) => Promise; diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index 034d3620900..79e26f7cee3 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -1150,9 +1150,17 @@ describe('agent pull request config', () => { cacheMaxCredentials: 256, maxConcurrentLookups: 4, batchTimeoutSeconds: 20, + sharedCache: false, + sharedCacheTimeoutMs: 250, }); }); + it('shares the cache between replicas only when asked to, and only as a boolean', () => { + expect(parse({ sharedCache: true }).success).toBe(true); + expect(parse({ sharedCache: false }).success).toBe(true); + expect(parse({ sharedCache: 'true' }).success).toBe(false); + }); + const enabled = { enabled: true, token: '${GITHUB_PULL_REQUEST_TOKEN}', @@ -1258,6 +1266,9 @@ describe('agent pull request config', () => { ['batchTimeoutSeconds', 0], ['batchTimeoutSeconds', 121], ['batchTimeoutSeconds', 1.5], + ['sharedCacheTimeoutMs', 9], + ['sharedCacheTimeoutMs', 5001], + ['sharedCacheTimeoutMs', 1.5], ])('rejects %s of %s', (field, value) => { expect(parse({ [field]: value }).success).toBe(false); }); @@ -1274,9 +1285,11 @@ describe('agent pull request config', () => { cacheMaxEntries: 100_000, maxConcurrentLookups: 16, batchTimeoutSeconds: 120, + sharedCacheTimeoutMs: 5000, }).success, ).toBe(true); expect(parse({ maxHeadComparisons: 0 }).success).toBe(true); + expect(parse({ sharedCacheTimeoutMs: 10 }).success).toBe(true); }); it.each([4, 3601, 1.5])('rejects a cache lifetime of %s seconds', (cacheTtlSeconds) => { diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index ecb9f888d02..fea2d9e8805 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1960,6 +1960,14 @@ export const agentsEndpointSchema = baseEndpointSchema maxCandidatePages: z.number().int().min(1).max(10).optional().default(1), /** Candidates compared with that commit before the search gives up. Each is one request. */ maxHeadComparisons: z.number().int().min(0).max(20).optional().default(3), + /** Share looked-up pull requests and GitHub rate limit cooldowns between replicas + * through Redis, so each replica does not spend the token on the same answer. Takes + * effect only when the deployment runs with Redis (`USE_REDIS`); each replica keeps its + * own in-memory cache either way. */ + sharedCache: z.boolean().optional().default(false), + /** Longest one shared cache read or write may take, in milliseconds, before a lookup + * carries on without it, so a slow Redis cannot hold the header request. */ + sharedCacheTimeoutMs: z.number().int().min(10).max(5000).optional().default(250), }) .optional(), /** Conversational background-task delivery policy. Automatic completion wakeups are @@ -4698,6 +4706,10 @@ export enum CacheKeys { * Key for the per-model reasoning effort cache of OpenRouter endpoints. */ REASONING_CAPABILITIES = 'REASONING_CAPABILITIES', + /** + * Key for pull request lookups and GitHub rate limit cooldowns shared across replicas. + */ + PULL_REQUESTS = 'PULL_REQUESTS', /** * Key for the app config namespace. */ From 5fe6e29dc0238cdc02ffbb4a30f18ae94f9f9392 Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:31:40 +0200 Subject: [PATCH 2/5] fix: Bound the Shared Pull Request Cache and Drop Commands at Their Deadline --- librechat.example.yaml | 3 +- packages/api/src/cache/pullRequests.spec.ts | 100 +++++++++++++++----- packages/api/src/cache/pullRequests.ts | 39 ++++++-- packages/api/src/pulls/lookup.spec.ts | 82 +++++++++++++++- packages/api/src/pulls/lookup.ts | 61 ++++++++---- packages/api/src/pulls/types.ts | 7 +- packages/data-provider/src/config.ts | 7 +- 7 files changed, 237 insertions(+), 62 deletions(-) diff --git a/librechat.example.yaml b/librechat.example.yaml index cb5359cb42d..d7eeb50b119 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -978,7 +978,8 @@ endpoints: # 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. + # # do not each spend the token on the same answer. Only takes effect with USE_REDIS. Redis + # # holds at most cacheMaxEntries x cacheMaxCredentials results, one per slot. # sharedCache: false # # Milliseconds one shared cache call may take before a lookup carries on without it (10 to 5000). # sharedCacheTimeoutMs: 250 diff --git a/packages/api/src/cache/pullRequests.spec.ts b/packages/api/src/cache/pullRequests.spec.ts index 3193811d28f..949f185f573 100644 --- a/packages/api/src/cache/pullRequests.spec.ts +++ b/packages/api/src/cache/pullRequests.spec.ts @@ -1,49 +1,101 @@ +import { createServer } from 'node:net'; import { CacheKeys } from 'librechat-data-provider'; +import { createClient, TimeoutError } from '@redis/client'; +import type { PullRequestStoreClient } from './pullRequests'; import { createPullRequestLookupStore } from './pullRequests'; -import { standardCache } from './cacheFactory'; import { cacheConfig } from './cacheConfig'; jest.mock('./cacheConfig', () => ({ - cacheConfig: { USE_REDIS: true, FORCED_IN_MEMORY_CACHE_NAMESPACES: [] as string[] }, + cacheConfig: { + USE_REDIS: true, + FORCED_IN_MEMORY_CACHE_NAMESPACES: [] as string[], + REDIS_KEY_PREFIX: '', + GLOBAL_PREFIX_SEPARATOR: '::', + }, })); -jest.mock('./cacheFactory', () => ({ standardCache: jest.fn() })); +jest.mock('./redisClients', () => ({ keyvRedisClient: null })); -const config = cacheConfig as { USE_REDIS: boolean; FORCED_IN_MEMORY_CACHE_NAMESPACES: string[] }; +const config = cacheConfig as { + USE_REDIS: boolean; + FORCED_IN_MEMORY_CACHE_NAMESPACES: string[]; + REDIS_KEY_PREFIX: string; +}; + +function fakeClient() { + const values = new Map(); + const commands = { + get: jest.fn(async (key: string) => values.get(key) ?? null), + set: jest.fn(async (key: string, value: string) => values.set(key, value)), + }; + 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 { + const server = createServer(); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + await new Promise((resolve) => server.close(() => resolve())); + if (address == null || typeof address === 'string') throw new Error('no port'); + return address.port; +} describe('createPullRequestLookupStore', () => { beforeEach(() => { config.USE_REDIS = true; config.FORCED_IN_MEMORY_CACHE_NAMESPACES = []; - jest.mocked(standardCache).mockReset(); + config.REDIS_KEY_PREFIX = ''; }); it('is absent without Redis, so nothing is shared', () => { config.USE_REDIS = false; - expect(createPullRequestLookupStore()).toBeUndefined(); - expect(standardCache).not.toHaveBeenCalled(); + expect(createPullRequestLookupStore(fakeClient().client)).toBeUndefined(); + }); + + it('is absent without a Redis client', () => { + expect(createPullRequestLookupStore(null)).toBeUndefined(); }); it('is absent when the namespace is forced in memory', () => { config.FORCED_IN_MEMORY_CACHE_NAMESPACES = [CacheKeys.PULL_REQUESTS]; - expect(createPullRequestLookupStore()).toBeUndefined(); + expect(createPullRequestLookupStore(fakeClient().client)).toBeUndefined(); + }); + + it('round-trips JSON under the namespace with the lifetime and the deadline', async () => { + config.REDIS_KEY_PREFIX = 'tenant'; + const { client, commands } = fakeClient(); + const store = createPullRequestLookupStore(client); + await store?.set('k', { until: 2 }, 5_000.4, 40); + const key = `tenant::${CacheKeys.PULL_REQUESTS}:k`; + expect(commands.set).toHaveBeenCalledWith(key, '{"until":2}', { PX: 5_001 }); + expect(client.withCommandOptions).toHaveBeenLastCalledWith({ timeout: 40 }); + await expect(store?.get('k', 30)).resolves.toEqual({ until: 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')); + await expect(createPullRequestLookupStore(client)?.get('k', 30)).rejects.toThrow('down'); }); - it('reads and writes the Redis cache with the given lifetime, and rejects its failures', async () => { - const cache = { - get: jest.fn().mockResolvedValue({ until: 1 }), - set: jest.fn().mockResolvedValue(true), - }; - jest - .mocked(standardCache) - .mockReturnValue(cache as unknown as ReturnType); - const store = createPullRequestLookupStore(); - expect(standardCache).toHaveBeenCalledWith(CacheKeys.PULL_REQUESTS, undefined, undefined, { - throwOnErrors: true, + 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 }, }); - await expect(store?.get('k')).resolves.toEqual({ until: 1 }); - await expect(store?.set('k', { until: 2 }, 5_000)).resolves.toBeUndefined(); - expect(cache.set).toHaveBeenCalledWith('k', { until: 2 }, 5_000); - cache.get.mockRejectedValueOnce(new Error('down')); - await expect(store?.get('k')).rejects.toThrow('down'); + redis.on('error', () => undefined); + /** Never settles: the client keeps waiting to reconnect until it is destroyed. */ + void redis.connect().catch(() => undefined); + try { + const store = createPullRequestLookupStore(redis as unknown as PullRequestStoreClient); + const started = Date.now(); + await expect(store?.get('k', 50)).rejects.toBeInstanceOf(TimeoutError); + await expect(store?.set('k', { until: 1 }, 1_000, 50)).rejects.toBeInstanceOf(TimeoutError); + expect(Date.now() - started).toBeLessThan(2_000); + } finally { + redis.destroy(); + } }); }); diff --git a/packages/api/src/cache/pullRequests.ts b/packages/api/src/cache/pullRequests.ts index 5ad3aa9a192..bf86717768e 100644 --- a/packages/api/src/cache/pullRequests.ts +++ b/packages/api/src/cache/pullRequests.ts @@ -1,28 +1,49 @@ import { CacheKeys } from 'librechat-data-provider'; import type { PullRequestLookupStore } from '~/pulls/types'; -import { standardCache } from './cacheFactory'; +import { keyvRedisClient } from './redisClients'; import { cacheConfig } from './cacheConfig'; +/** The commands the store sends; the single-node and the cluster client both provide them. */ +type StoreCommands = { + get(key: string): Promise; + set(key: string, value: string, options: { PX: number }): Promise; +}; + +export type PullRequestStoreClient = { + withCommandOptions(options: { timeout: number }): StoreCommands; +}; + /** * 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(): PullRequestLookupStore | undefined { +export function createPullRequestLookupStore( + client: PullRequestStoreClient | null = keyvRedisClient, +): PullRequestLookupStore | undefined { if ( + client == null || !cacheConfig.USE_REDIS || cacheConfig.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.PULL_REQUESTS) ) { return undefined; } - /** Rejects instead of resolving a failed call as a miss, so the lookup can log it. */ - const cache = standardCache(CacheKeys.PULL_REQUESTS, undefined, undefined, { - throwOnErrors: true, - }); + const prefix = cacheConfig.REDIS_KEY_PREFIX + ? `${cacheConfig.REDIS_KEY_PREFIX}${cacheConfig.GLOBAL_PREFIX_SEPARATOR}${CacheKeys.PULL_REQUESTS}:` + : `${CacheKeys.PULL_REQUESTS}:`; return { - get: (key) => cache.get(key), - set: async (key, value, ttlMs) => { - await cache.set(key, value, ttlMs); + get: async (key, timeoutMs) => { + const raw = await client.withCommandOptions({ timeout: timeoutMs }).get(prefix + key); + return raw == null ? undefined : JSON.parse(raw); + }, + set: async (key, value, ttlMs, timeoutMs) => { + await client + .withCommandOptions({ timeout: timeoutMs }) + .set(prefix + key, JSON.stringify(value), { PX: Math.max(1, Math.ceil(ttlMs)) }); }, }; } diff --git a/packages/api/src/pulls/lookup.spec.ts b/packages/api/src/pulls/lookup.spec.ts index 940c3382727..dcb8560b22b 100644 --- a/packages/api/src/pulls/lookup.spec.ts +++ b/packages/api/src/pulls/lookup.spec.ts @@ -476,6 +476,9 @@ describe('the repositories the credential may be used for', () => { }); }); +const isPlainObject = (value: unknown): value is Record => + typeof value === 'object' && value !== null; + describe('the shared store', () => { const token = 'ghp_secret_token'; const shared = { ...input, token, sharedCache: true }; @@ -484,13 +487,13 @@ describe('the shared store', () => { function createStore(clock: () => number) { const values = new Map(); const store = { - get: jest.fn(async (key: string) => { + get: jest.fn(async (key: string, _timeoutMs: number) => { const held = values.get(key); return held != null && held.until > clock() ? JSON.parse(JSON.stringify(held.value)) : undefined; }), - set: jest.fn(async (key: string, value: unknown, ttlMs: number) => { + set: jest.fn(async (key: string, value: unknown, ttlMs: number, _timeoutMs: number) => { values.set(key, { value, until: clock() + ttlMs }); }), }; @@ -608,9 +611,76 @@ describe('the shared store', () => { expect(find).toHaveBeenCalledTimes(1); }); + /** The key digest a lookup with `shared` writes, so a planted entry differs only as labelled. */ + async function sharedKey(): Promise { + const { store } = createStore(Date.now); + await createPullRequestLookup({ source: { find: jest.fn().mockResolvedValue(value) }, store })( + shared, + ); + await settle(); + return (store.set.mock.calls[0][1] as { key: string }).key; + } + + it('holds no more results than every partition together, however many branches are looked up', async () => { + const { store, values } = createStore(Date.now); + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ source: { find }, store }); + for (let i = 0; i < 200; i++) { + await lookup({ ...shared, branch: `b${i}`, cacheMaxEntries: 10, cacheMaxCredentials: 2 }); + } + await settle(); + expect(store.set).toHaveBeenCalledTimes(200); + expect(values.size).toBeLessThanOrEqual(20); + }); + + it('holds one rate limit cooldown per credential slot', async () => { + const { store, values } = createStore(Date.now); + const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); + const lookup = createPullRequestLookup({ source: { find: limited }, store }); + for (let i = 0; i < 50; i++) { + await lookup({ ...shared, token: `t${i}`, cacheMaxCredentials: 4 }); + } + await settle(); + expect([...values.keys()].filter((key) => key.includes('cooldown')).length).toBeLessThanOrEqual( + 4, + ); + }); + + it('gives the store its timeout as the deadline of every call', async () => { + const { store } = createStore(Date.now); + const lookup = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + }); + await lookup({ ...shared, sharedCacheTimeoutMs: 75 }); + await settle(); + expect(store.get.mock.calls.every((call) => call[1] === 75)).toBe(true); + expect(store.set.mock.calls[0][3]).toBe(75); + }); + + it('serves a well-formed entry stored under its own key', async () => { + const key = await sharedKey(); + const store: PullRequestLookupStore = { + get: jest.fn(async (slot: string) => + slot.includes('cooldown') + ? undefined + : { result: { ok: true, value }, expiresAt: Number.MAX_SAFE_INTEGER, key }, + ), + set: jest.fn(async () => undefined), + }; + const find = jest.fn().mockResolvedValue(null); + const lookup = createPullRequestLookup({ source: { find }, store }); + await expect(lookup(shared)).resolves.toEqual({ ok: true, value }); + expect(find).not.toHaveBeenCalled(); + }); + it.each([ ['nothing', undefined], ['a string', 'cached'], + [ + 'an answer for another key in the same slot', + { result: { ok: true, value }, expiresAt: Number.MAX_SAFE_INTEGER, key: 'other' }, + ], ['an entry without an expiry', { result: { ok: true, value } }], ['an expiry that is not a number', { result: { ok: true, value }, expiresAt: 'later' }], ['a result without ok', { result: { value }, expiresAt: Number.MAX_SAFE_INTEGER }], @@ -644,9 +714,13 @@ describe('the shared store', () => { }, ], ['an entry that has expired', { result: { ok: true, value }, expiresAt: 0 }], - ])('ignores %s in the store', async (_label, stored) => { + ])('ignores %s in the store', async (_label, planted) => { + const key = await sharedKey(); + const stored = isPlainObject(planted) && !('key' in planted) ? { ...planted, key } : planted; const store: PullRequestLookupStore = { - get: jest.fn(async (key: string) => (key.includes('cooldown') ? { until: 'soon' } : stored)), + get: jest.fn(async (slot: string) => + slot.includes('cooldown') ? { until: 'soon' } : stored, + ), set: jest.fn(async () => undefined), }; const find = jest.fn().mockResolvedValue(null); diff --git a/packages/api/src/pulls/lookup.ts b/packages/api/src/pulls/lookup.ts index 53e150e0894..379bc2eb623 100644 --- a/packages/api/src/pulls/lookup.ts +++ b/packages/api/src/pulls/lookup.ts @@ -30,15 +30,23 @@ const DEFAULT_SHARED_TIMEOUT_MS = 250; const STORE_PREFIX = 'pr:v1:'; type Entry = { result: PullRequestLookupResult; expiresAt: number }; -type Cooldown = { until: number }; +/** `key` is the digest of the full key, since a slot is shared by every key that maps to it. */ +type Stored = Entry & { key: string }; +type Cooldown = { key: string; until: number }; /** Keys carry a digest, never the token, so no credential is held in a map key or a heap dump. */ const scopeOf = (token: string): string => createHash('sha256').update(token).digest('hex').slice(0, 32); -/** Repository and branch are user-influenced and unbounded, so the stored key is a fixed size. */ -const storeKeyOf = (key: string): string => - `${STORE_PREFIX}result:${createHash('sha256').update(key).digest('hex')}`; +const digestOf = (key: string): string => createHash('sha256').update(key).digest('hex'); + +/** + * The shared store is a fixed number of slots, so it holds no more than its slot count however + * many repositories and branches are looked up: an answer goes in the slot its key's digest picks, + * replacing whatever was there, and a read whose slot holds another key is a miss. + */ +const slotOf = (kind: 'result' | 'cooldown', digest: string, slots: number): string => + `${STORE_PREFIX}${kind}:${Number.parseInt(digest.slice(0, 12), 16) % slots}`; const isRecord = (value: unknown): value is Record => typeof value === 'object' && value !== null && !Array.isArray(value); @@ -50,8 +58,9 @@ const isTime = (value: unknown): value is number => typeof value === 'number' && Number.isFinite(value); /** Another replica, or another version of one, wrote this; anything not shaped right is a miss. */ -function isEntry(value: unknown): value is Entry { - if (!isRecord(value) || !isTime(value.expiresAt) || !isRecord(value.result)) return false; +function isStored(value: unknown, key: string): value is Stored { + if (!isRecord(value) || value.key !== key) return false; + if (!isTime(value.expiresAt) || !isRecord(value.result)) return false; const { result } = value; if (result.ok === false) { return isRecord(result.error) && isOneOf(PULL_REQUEST_ERROR_CODES, result.error.code); @@ -76,11 +85,13 @@ function isEntry(value: unknown): value is Entry { ); } -const isCooldown = (value: unknown): value is Cooldown => isRecord(value) && isTime(value.until); +const isCooldown = (value: unknown, key: string): value is Cooldown => + isRecord(value) && value.key === key && isTime(value.until); /** * Runs one shared store call, bounded by `timeoutMs`. The store is an optimization, so a failure - * or a slow answer is logged and reads as `undefined`, and the lookup carries on in memory. + * or a slow answer is logged and reads as `undefined`, and the lookup carries on in memory. The + * store gets the same deadline and drops its own command at it; the race covers one that does not. */ async function bounded(call: () => Promise, timeoutMs: number): Promise { let timer: NodeJS.Timeout | undefined; @@ -194,12 +205,12 @@ export function createPullRequestLookup({ function share( target: PullRequestLookupStore, key: string, - value: Entry | Cooldown, + value: Stored | Cooldown, ttlMs: number, timeoutMs: number, ): void { if (ttlMs <= 0) return; - void bounded(() => target.set(key, value, ttlMs), timeoutMs); + void bounded(() => target.set(key, value, ttlMs, timeoutMs), timeoutMs); } return async ({ @@ -253,22 +264,25 @@ export function createPullRequestLookup({ result.ok ? ttlMs : Math.min(ttlMs, FAILURE_TTL_MS); const capacity = cacheMaxEntries ?? maxEntries; const credentials = cacheMaxCredentials ?? DEFAULT_MAX_PARTITIONS; - const storeKey = storeKeyOf(key); - const cooldownKey = `${STORE_PREFIX}cooldown:${createHash('sha256').update(limitScope).digest('hex')}`; + /** As many results as every replica's partitions hold together, and a cooldown per credential. */ + const resultDigest = digestOf(key); + const cooldownDigest = digestOf(limitScope); + const storeKey = slotOf('result', resultDigest, capacity * credentials); + const cooldownKey = slotOf('cooldown', cooldownDigest, credentials); const run = (async (): Promise => { if (shared != null) { const [stored, cooldown] = await Promise.all([ - bounded(() => shared.get(storeKey), sharedCacheTimeoutMs), - bounded(() => shared.get(cooldownKey), sharedCacheTimeoutMs), + bounded(() => shared.get(storeKey, sharedCacheTimeoutMs), sharedCacheTimeoutMs), + bounded(() => shared.get(cooldownKey, sharedCacheTimeoutMs), sharedCacheTimeoutMs), ]); - if (isEntry(stored) && stored.expiresAt > now()) { + if (isStored(stored, resultDigest) && stored.expiresAt > now()) { /** Never kept longer than this replica would have kept its own answer. */ const expiresAt = Math.min(stored.expiresAt, now() + lifetimeOf(stored.result)); remember(scope, key, { result: stored.result, expiresAt }, capacity, credentials); return stored.result; } - if (isCooldown(cooldown) && cooldown.until > now()) { + if (isCooldown(cooldown, cooldownDigest) && cooldown.until > now()) { startCooldown(limitScope, Math.min(cooldown.until, now() + MAX_COOLDOWN_MS)); return rateLimited; } @@ -299,9 +313,16 @@ export function createPullRequestLookup({ MAX_COOLDOWN_MS, Math.max(MIN_COOLDOWN_MS, error.retryAfterMs ?? 0), ); - startCooldown(limitScope, now() + wait); + const end = now() + wait; + startCooldown(limitScope, end); if (shared != null) { - share(shared, cooldownKey, { until: now() + wait }, wait, sharedCacheTimeoutMs); + share( + shared, + cooldownKey, + { key: cooldownDigest, until: end }, + wait, + sharedCacheTimeoutMs, + ); } } result = { ok: false, error: { code } }; @@ -309,7 +330,9 @@ export function createPullRequestLookup({ const lifetime = lifetimeOf(result); const entry: Entry = { result, expiresAt: now() + lifetime }; remember(scope, key, entry, capacity, credentials); - if (shared != null) share(shared, storeKey, entry, lifetime, sharedCacheTimeoutMs); + if (shared != null) { + share(shared, storeKey, { ...entry, key: resultDigest }, lifetime, sharedCacheTimeoutMs); + } return result; })().finally(() => inflight.delete(key)); inflight.set(key, run); diff --git a/packages/api/src/pulls/types.ts b/packages/api/src/pulls/types.ts index c5b5f07e732..1efd436ab37 100644 --- a/packages/api/src/pulls/types.ts +++ b/packages/api/src/pulls/types.ts @@ -84,11 +84,12 @@ export type PullRequestLookupInput = PullRequestFindInput & { /** * Storage the lookup shares with other replicas, so one answer or one rate limit from GitHub serves - * every replica. Values must round-trip as JSON; `ttlMs` is how long the store may keep one. + * every replica. Values must round-trip as JSON; `ttlMs` is how long the store may keep one. A call + * still unsent at `timeoutMs` must be dropped, not left queued, and reject. */ export type PullRequestLookupStore = { - get(key: string): Promise; - set(key: string, value: unknown, ttlMs: number): Promise; + get(key: string, timeoutMs: number): Promise; + set(key: string, value: unknown, ttlMs: number, timeoutMs: number): Promise; }; export type PullRequestLookup = (input: PullRequestLookupInput) => Promise; diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index fea2d9e8805..2eb98a6675b 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1963,10 +1963,13 @@ export const agentsEndpointSchema = baseEndpointSchema /** Share looked-up pull requests and GitHub rate limit cooldowns between replicas * through Redis, so each replica does not spend the token on the same answer. Takes * effect only when the deployment runs with Redis (`USE_REDIS`); each replica keeps its - * own in-memory cache either way. */ + * own in-memory cache either way. Redis holds at most `cacheMaxEntries` times + * `cacheMaxCredentials` results and `cacheMaxCredentials` cooldowns: each goes in a + * slot picked from its key, replacing what was there. */ sharedCache: z.boolean().optional().default(false), /** Longest one shared cache read or write may take, in milliseconds, before a lookup - * carries on without it, so a slow Redis cannot hold the header request. */ + * carries on without it, so a slow Redis cannot hold the header request. A command + * still queued at that point, as during a disconnect, is dropped rather than sent later. */ sharedCacheTimeoutMs: z.number().int().min(10).max(5000).optional().default(250), }) .optional(), From e08be110e1ba35e1b6a07658355fe7a58437f673 Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:36:54 +0200 Subject: [PATCH 3/5] fix: Pass the Cache Config to the Pull Request Store and Check It Before a Local Cooldown --- api/server/routes/convos.js | 4 +- packages/api/src/cache/pullRequests.spec.ts | 47 ++++++++------------- packages/api/src/cache/pullRequests.ts | 27 ++++++++---- packages/api/src/pulls/lookup.spec.ts | 39 +++++++++++++++++ packages/api/src/pulls/lookup.ts | 21 +++++---- 5 files changed, 90 insertions(+), 48 deletions(-) diff --git a/api/server/routes/convos.js b/api/server/routes/convos.js index 147f2d07bc5..4ec0007c4e1 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -24,6 +24,8 @@ const { createProxyAwareFetch, createPullRequestLookup, createPullRequestLookupStore, + keyvRedisClient, + cacheConfig, createConversationPullRequestsHandler, backgroundTaskRegistry, createSubagentThreadViewHandler, @@ -176,7 +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(), + store: createPullRequestLookupStore({ client: keyvRedisClient, config: cacheConfig }), }); const conversationPullRequestHandler = createConversationPullRequestHandler({ getConvoLaneGit: db.getConvoLaneGit, diff --git a/packages/api/src/cache/pullRequests.spec.ts b/packages/api/src/cache/pullRequests.spec.ts index 949f185f573..302cdc3e378 100644 --- a/packages/api/src/cache/pullRequests.spec.ts +++ b/packages/api/src/cache/pullRequests.spec.ts @@ -2,23 +2,14 @@ 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'; -import { cacheConfig } from './cacheConfig'; -jest.mock('./cacheConfig', () => ({ - cacheConfig: { - USE_REDIS: true, - FORCED_IN_MEMORY_CACHE_NAMESPACES: [] as string[], - REDIS_KEY_PREFIX: '', - GLOBAL_PREFIX_SEPARATOR: '::', - }, -})); -jest.mock('./redisClients', () => ({ keyvRedisClient: null })); - -const config = cacheConfig as { - USE_REDIS: boolean; - FORCED_IN_MEMORY_CACHE_NAMESPACES: string[]; - REDIS_KEY_PREFIX: string; +const base: PullRequestStoreConfig = { + USE_REDIS: true, + FORCED_IN_MEMORY_CACHE_NAMESPACES: [], + REDIS_KEY_PREFIX: '', + GLOBAL_PREFIX_SEPARATOR: '::', }; function fakeClient() { @@ -42,30 +33,24 @@ async function closedPort(): Promise { } describe('createPullRequestLookupStore', () => { - beforeEach(() => { - config.USE_REDIS = true; - config.FORCED_IN_MEMORY_CACHE_NAMESPACES = []; - config.REDIS_KEY_PREFIX = ''; - }); - it('is absent without Redis, so nothing is shared', () => { - config.USE_REDIS = false; - expect(createPullRequestLookupStore(fakeClient().client)).toBeUndefined(); + const config = { ...base, USE_REDIS: false }; + expect(createPullRequestLookupStore({ client: fakeClient().client, config })).toBeUndefined(); }); it('is absent without a Redis client', () => { - expect(createPullRequestLookupStore(null)).toBeUndefined(); + expect(createPullRequestLookupStore({ client: null, config: base })).toBeUndefined(); }); it('is absent when the namespace is forced in memory', () => { - config.FORCED_IN_MEMORY_CACHE_NAMESPACES = [CacheKeys.PULL_REQUESTS]; - expect(createPullRequestLookupStore(fakeClient().client)).toBeUndefined(); + 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 () => { - config.REDIS_KEY_PREFIX = 'tenant'; const { client, commands } = fakeClient(); - const store = createPullRequestLookupStore(client); + const config = { ...base, REDIS_KEY_PREFIX: 'tenant' }; + const store = createPullRequestLookupStore({ client, config }); await store?.set('k', { until: 2 }, 5_000.4, 40); const key = `tenant::${CacheKeys.PULL_REQUESTS}:k`; expect(commands.set).toHaveBeenCalledWith(key, '{"until":2}', { PX: 5_001 }); @@ -78,7 +63,8 @@ describe('createPullRequestLookupStore', () => { it('rejects a failed command', async () => { const { client, commands } = fakeClient(); commands.get.mockRejectedValueOnce(new Error('down')); - await expect(createPullRequestLookupStore(client)?.get('k', 30)).rejects.toThrow('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 () => { @@ -89,7 +75,8 @@ describe('createPullRequestLookupStore', () => { /** Never settles: the client keeps waiting to reconnect until it is destroyed. */ void redis.connect().catch(() => undefined); try { - const store = createPullRequestLookupStore(redis as unknown as PullRequestStoreClient); + const client = redis as unknown as PullRequestStoreClient; + const store = createPullRequestLookupStore({ client, config: base }); const started = Date.now(); await expect(store?.get('k', 50)).rejects.toBeInstanceOf(TimeoutError); await expect(store?.set('k', { until: 1 }, 1_000, 50)).rejects.toBeInstanceOf(TimeoutError); diff --git a/packages/api/src/cache/pullRequests.ts b/packages/api/src/cache/pullRequests.ts index bf86717768e..9c00cd3f1e4 100644 --- a/packages/api/src/cache/pullRequests.ts +++ b/packages/api/src/cache/pullRequests.ts @@ -1,7 +1,6 @@ import { CacheKeys } from 'librechat-data-provider'; import type { PullRequestLookupStore } from '~/pulls/types'; -import { keyvRedisClient } from './redisClients'; -import { cacheConfig } from './cacheConfig'; +import type { cacheConfig } from './cacheConfig'; /** The commands the store sends; the single-node and the cluster client both provide them. */ type StoreCommands = { @@ -13,6 +12,12 @@ 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 @@ -22,18 +27,22 @@ export type PullRequestStoreClient = { * 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: PullRequestStoreClient | null = keyvRedisClient, -): PullRequestLookupStore | undefined { +export function createPullRequestLookupStore({ + client, + config, +}: { + client: PullRequestStoreClient | null; + config: PullRequestStoreConfig; +}): PullRequestLookupStore | undefined { if ( client == null || - !cacheConfig.USE_REDIS || - cacheConfig.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.PULL_REQUESTS) + !config.USE_REDIS || + config.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.PULL_REQUESTS) ) { return undefined; } - const prefix = cacheConfig.REDIS_KEY_PREFIX - ? `${cacheConfig.REDIS_KEY_PREFIX}${cacheConfig.GLOBAL_PREFIX_SEPARATOR}${CacheKeys.PULL_REQUESTS}:` + 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) => { diff --git a/packages/api/src/pulls/lookup.spec.ts b/packages/api/src/pulls/lookup.spec.ts index dcb8560b22b..2b8d44f5614 100644 --- a/packages/api/src/pulls/lookup.spec.ts +++ b/packages/api/src/pulls/lookup.spec.ts @@ -743,6 +743,45 @@ describe('the shared store', () => { expect(store.set).not.toHaveBeenCalled(); }); + it('serves a shared answer while this replica is cooling down', async () => { + const { store } = createStore(Date.now); + const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); + const cooled = createPullRequestLookup({ source: { find: limited }, store }); + const other = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + }); + await other({ ...shared, branch: 'other' }); + await settle(); + await cooled({ ...shared, branch: 'limited' }); + await settle(); + await expect(cooled({ ...shared, branch: 'other' })).resolves.toEqual({ ok: true, value }); + /** Nothing stored for this branch, so the cooldown still answers without asking GitHub. */ + await expect(cooled({ ...shared, branch: 'third' })).resolves.toEqual({ + ok: false, + error: { code: 'RATE_LIMITED' }, + }); + expect(limited).toHaveBeenCalledTimes(1); + }); + + it('does not join a lookup running under another sharing policy', async () => { + const { store } = createStore(Date.now); + let release: (pr: typeof value) => void = () => undefined; + const answer = new Promise((resolve) => (release = resolve)); + const find = jest.fn(() => answer); + const lookup = createPullRequestLookup({ source: { find }, store }); + const local = lookup({ ...input, token, sharedCache: false }); + const fast = lookup({ ...shared, sharedCacheTimeoutMs: 10 }); + const slow = lookup({ ...shared, sharedCacheTimeoutMs: 5_000 }); + await settle(); + /** The local caller never reached the store; each shared deadline asked it on its own. */ + expect(store.get.mock.calls.map((call) => call[1]).sort((x, y) => x - y)).toEqual([ + 10, 10, 5_000, 5_000, + ]); + release(value); + await Promise.all([local, fast, slow]); + }); + it('does not split the cache by whether it is shared', async () => { const find = jest.fn().mockResolvedValue(value); const lookup = createPullRequestLookup({ diff --git a/packages/api/src/pulls/lookup.ts b/packages/api/src/pulls/lookup.ts index 379bc2eb623..0360d2e2003 100644 --- a/packages/api/src/pulls/lookup.ts +++ b/packages/api/src/pulls/lookup.ts @@ -248,18 +248,22 @@ export function createPullRequestLookup({ const key = `${scope}\0${provider ?? ''}\0${host ?? ''}\0${repo}#${branch}\0${head ?? ''}\0${pullRequest ?? ''}\0${policy}`; const cached = touch(scope)?.get(key); if (cached != null && cached.expiresAt > now()) return cached.result; - const pending = inflight.get(key); + /** Only when this caller's settings ask for it; the store is there whenever Redis is. */ + const shared = sharedCache === true ? store : undefined; + /** A caller joins only work run under its own sharing policy, so its toggle and deadline hold. */ + const flight = shared != null ? `${key}\0shared\0${sharedCacheTimeoutMs}` : key; + const pending = inflight.get(flight); if (pending != null) return pending; /** A rate limit belongs to the credential on one provider host; elsewhere it is not limited. */ const limitScope = `${scope}\0${provider ?? ''}\0${host ?? ''}`; + let cooling = false; const until = cooldowns.get(limitScope); if (until != null) { - if (until > now()) return rateLimited; - cooldowns.delete(limitScope); + if (until > now()) cooling = true; + else cooldowns.delete(limitScope); } - - /** Only when this caller's settings ask for it; the store is there whenever Redis is. */ - const shared = sharedCache === true ? store : undefined; + /** Another replica may already hold the answer, so only a lookup without the store stops here. */ + if (cooling && shared == null) return rateLimited; const lifetimeOf = (result: PullRequestLookupResult): number => result.ok ? ttlMs : Math.min(ttlMs, FAILURE_TTL_MS); const capacity = cacheMaxEntries ?? maxEntries; @@ -282,6 +286,7 @@ export function createPullRequestLookup({ remember(scope, key, { result: stored.result, expiresAt }, capacity, credentials); return stored.result; } + if (cooling) return rateLimited; if (isCooldown(cooldown, cooldownDigest) && cooldown.until > now()) { startCooldown(limitScope, Math.min(cooldown.until, now() + MAX_COOLDOWN_MS)); return rateLimited; @@ -334,8 +339,8 @@ export function createPullRequestLookup({ share(shared, storeKey, { ...entry, key: resultDigest }, lifetime, sharedCacheTimeoutMs); } return result; - })().finally(() => inflight.delete(key)); - inflight.set(key, run); + })().finally(() => inflight.delete(flight)); + inflight.set(flight, run); return run; }; } From b91c1e85572d242c43825263d407a0b3f619bedb Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:19:42 +0200 Subject: [PATCH 4/5] fix: Keep Batch Lookups Apart by Cache and Sharing Policy --- packages/api/src/pulls/handler.spec.ts | 37 ++++++++++++++++++++++++++ packages/api/src/pulls/handler.ts | 5 ++++ 2 files changed, 42 insertions(+) diff --git a/packages/api/src/pulls/handler.spec.ts b/packages/api/src/pulls/handler.spec.ts index 6adc31a85d9..b4164ebe684 100644 --- a/packages/api/src/pulls/handler.spec.ts +++ b/packages/api/src/pulls/handler.spec.ts @@ -1117,6 +1117,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 = []; const lookup = jest.fn( diff --git a/packages/api/src/pulls/handler.ts b/packages/api/src/pulls/handler.ts index f0f18492f59..4c31fe4889d 100644 --- a/packages/api/src/pulls/handler.ts +++ b/packages/api/src/pulls/handler.ts @@ -478,6 +478,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); From 9e5951b1760334dd695bd577a48c73fa6c404cd9 Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:54:08 +0200 Subject: [PATCH 5/5] fix: Share Only Found Pull Requests and Rate Limits, Never Moving Them Backwards --- librechat.example.yaml | 7 +- .../pullRequests.cache_integration.spec.ts | 69 +++++ packages/api/src/cache/pullRequests.spec.ts | 20 +- packages/api/src/cache/pullRequests.ts | 34 ++- packages/api/src/pulls/handler.spec.ts | 1 + packages/api/src/pulls/handler.ts | 1 + packages/api/src/pulls/lookup.spec.ts | 287 ++++++++++++------ packages/api/src/pulls/lookup.ts | 69 +++-- packages/api/src/pulls/types.ts | 16 +- packages/data-provider/src/config.spec.ts | 6 + packages/data-provider/src/config.ts | 13 +- .../data-schemas/src/app/resolution.spec.ts | 5 + 12 files changed, 376 insertions(+), 152 deletions(-) create mode 100644 packages/api/src/cache/__tests__/pullRequests.cache_integration.spec.ts diff --git a/librechat.example.yaml b/librechat.example.yaml index d7eeb50b119..2a381add6a2 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -978,11 +978,14 @@ endpoints: # 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. Redis - # # holds at most cacheMaxEntries x cacheMaxCredentials results, one per slot. + # # 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. diff --git a/packages/api/src/cache/__tests__/pullRequests.cache_integration.spec.ts b/packages/api/src/cache/__tests__/pullRequests.cache_integration.spec.ts new file mode 100644 index 00000000000..1d34b6cd15c --- /dev/null +++ b/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; + 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(); + }); +}); diff --git a/packages/api/src/cache/pullRequests.spec.ts b/packages/api/src/cache/pullRequests.spec.ts index 302cdc3e378..9b64be25dda 100644 --- a/packages/api/src/cache/pullRequests.spec.ts +++ b/packages/api/src/cache/pullRequests.spec.ts @@ -16,7 +16,10 @@ function fakeClient() { const values = new Map(); const commands = { get: jest.fn(async (key: string) => values.get(key) ?? null), - set: jest.fn(async (key: string, value: string) => values.set(key, value)), + 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 }; @@ -51,11 +54,14 @@ describe('createPullRequestLookupStore', () => { const { client, commands } = fakeClient(); const config = { ...base, REDIS_KEY_PREFIX: 'tenant' }; const store = createPullRequestLookupStore({ client, config }); - await store?.set('k', { until: 2 }, 5_000.4, 40); + await store?.put('k', { key: 'owner', order: 2 }, 5_000.4, 40); const key = `tenant::${CacheKeys.PULL_REQUESTS}:k`; - expect(commands.set).toHaveBeenCalledWith(key, '{"until":2}', { PX: 5_001 }); + 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({ until: 2 }); + 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(); }); @@ -75,11 +81,13 @@ describe('createPullRequestLookupStore', () => { /** Never settles: the client keeps waiting to reconnect until it is destroyed. */ void redis.connect().catch(() => undefined); try { - const client = redis as unknown as PullRequestStoreClient; + /** 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); - await expect(store?.set('k', { until: 1 }, 1_000, 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(); diff --git a/packages/api/src/cache/pullRequests.ts b/packages/api/src/cache/pullRequests.ts index 9c00cd3f1e4..fc2512b1bbf 100644 --- a/packages/api/src/cache/pullRequests.ts +++ b/packages/api/src/cache/pullRequests.ts @@ -5,9 +5,27 @@ 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; - set(key: string, value: string, options: { PX: number }): Promise; + eval(script: string, options: { keys: string[]; arguments: string[] }): Promise; }; +/** + * 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; }; @@ -49,10 +67,16 @@ export function createPullRequestLookupStore({ const raw = await client.withCommandOptions({ timeout: timeoutMs }).get(prefix + key); return raw == null ? undefined : JSON.parse(raw); }, - set: async (key, value, ttlMs, timeoutMs) => { - await client - .withCommandOptions({ timeout: timeoutMs }) - .set(prefix + key, JSON.stringify(value), { PX: Math.max(1, Math.ceil(ttlMs)) }); + 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), + ], + }); }, }; } diff --git a/packages/api/src/pulls/handler.spec.ts b/packages/api/src/pulls/handler.spec.ts index b4164ebe684..a835a9f33c9 100644 --- a/packages/api/src/pulls/handler.spec.ts +++ b/packages/api/src/pulls/handler.spec.ts @@ -97,6 +97,7 @@ describe('createConversationPullRequestHandler', () => { cacheMaxCredentials: 256, sharedCache: false, sharedCacheTimeoutMs: 250, + sharedCacheSlots: 65_536, allowedRepositories: ['o/r'], limits: { requestTimeoutMs: 10_000, diff --git a/packages/api/src/pulls/handler.ts b/packages/api/src/pulls/handler.ts index 4c31fe4889d..3402a0bf58f 100644 --- a/packages/api/src/pulls/handler.ts +++ b/packages/api/src/pulls/handler.ts @@ -231,6 +231,7 @@ function lookupInput(settings: PullRequestSettings, token: string, lane: Eligibl 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, diff --git a/packages/api/src/pulls/lookup.spec.ts b/packages/api/src/pulls/lookup.spec.ts index 2b8d44f5614..d31bf7153fd 100644 --- a/packages/api/src/pulls/lookup.spec.ts +++ b/packages/api/src/pulls/lookup.spec.ts @@ -1,5 +1,5 @@ import type { TConversationPullRequest } from 'librechat-data-provider'; -import type { PullRequestLookupStore, PullRequestSource } from './types'; +import type { PullRequestSource, PullRequestLookupStore, PullRequestSharedValue } from './types'; import { createPullRequestLookup } from './lookup'; import { PullRequestSourceError } from './types'; @@ -475,17 +475,20 @@ describe('the repositories the credential may be used for', () => { expect(find).toHaveBeenCalledTimes(2); }); }); - const isPlainObject = (value: unknown): value is Record => typeof value === 'object' && value !== null; describe('the shared store', () => { const token = 'ghp_secret_token'; const shared = { ...input, token, sharedCache: true }; + const forever = Number.MAX_SAFE_INTEGER; - /** A Map with lifetimes, standing in for Redis between two replicas. */ + /** + * A Map with lifetimes, standing in for Redis between replicas, with the store's write rule: a + * value for the same owner with an order at least as large is kept. + */ function createStore(clock: () => number) { - const values = new Map(); + const values = new Map(); const store = { get: jest.fn(async (key: string, _timeoutMs: number) => { const held = values.get(key); @@ -493,9 +496,20 @@ describe('the shared store', () => { ? JSON.parse(JSON.stringify(held.value)) : undefined; }), - set: jest.fn(async (key: string, value: unknown, ttlMs: number, _timeoutMs: number) => { - values.set(key, { value, until: clock() + ttlMs }); - }), + put: jest.fn( + async (key: string, value: PullRequestSharedValue, ttlMs: number, _timeoutMs: number) => { + const held = values.get(key); + if ( + held != null && + held.until > clock() && + held.value.key === value.key && + held.value.order >= value.order + ) { + return; + } + values.set(key, { value, until: clock() + ttlMs }); + }, + ), }; return { store, values }; } @@ -503,7 +517,17 @@ describe('the shared store', () => { /** Background writes settle on the next turns of the event loop. */ const settle = () => new Promise((resolve) => setImmediate(resolve)); - it("serves one replica's answer to another without asking GitHub again", async () => { + /** The owner digest a lookup with `shared` writes, so a planted value differs only as labelled. */ + async function sharedKey(): Promise { + const { store } = createStore(Date.now); + await createPullRequestLookup({ source: { find: jest.fn().mockResolvedValue(value) }, store })( + shared, + ); + await settle(); + return store.put.mock.calls[0][1].key; + } + + it("serves one replica's answer to another without asking the provider again", async () => { const { store } = createStore(Date.now); const first = jest.fn().mockResolvedValue(value); const second = jest.fn().mockResolvedValue(value); @@ -519,6 +543,20 @@ describe('the shared store', () => { expect(store.get).not.toHaveBeenCalled(); }); + it('shares "no pull request" as an answer too', async () => { + const { store } = createStore(Date.now); + const second = jest.fn().mockResolvedValue(value); + const a = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(null) }, + store, + }); + const b = createPullRequestLookup({ source: { find: second }, store }); + await a(shared); + await settle(); + await expect(b(shared)).resolves.toEqual({ ok: true, value: null }); + expect(second).not.toHaveBeenCalled(); + }); + it('keeps a shared answer no longer than its lifetime', async () => { let clock = 0; const { store } = createStore(() => clock); @@ -536,6 +574,25 @@ describe('the shared store', () => { expect(find).toHaveBeenCalledTimes(1); }); + it.each([ + ['an upstream error', new Error('ECONNRESET')], + ['a source failure', new PullRequestSourceError('UPSTREAM_ERROR')], + ['an unconfigured source', new PullRequestSourceError('NOT_CONFIGURED')], + ])('never shares %s, so another replica tries the provider itself', async (_label, error) => { + const { store } = createStore(Date.now); + const find = jest.fn().mockResolvedValue(value); + const a = createPullRequestLookup({ + source: { find: jest.fn().mockRejectedValue(error) }, + store, + }); + const b = createPullRequestLookup({ source: { find }, store }); + await a(shared); + await settle(); + expect(store.put).not.toHaveBeenCalled(); + await expect(b(shared)).resolves.toEqual({ ok: true, value }); + expect(find).toHaveBeenCalledTimes(1); + }); + it('puts every replica into the cooldown one of them was rate limited into', async () => { const { store } = createStore(Date.now); const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); @@ -544,16 +601,54 @@ describe('the shared store', () => { const b = createPullRequestLookup({ source: { find }, store }); await a(shared); await settle(); - /** Another branch, so it is the cooldown and not a cached failure that answers. */ + /** Another branch, so it is the cooldown and not a cached answer that answers. */ await expect(b({ ...shared, branch: 'other' })).resolves.toEqual({ ok: false, error: { code: 'RATE_LIMITED' }, }); expect(find).not.toHaveBeenCalled(); - const cooldown = store.set.mock.calls.find(([key]) => key.includes('cooldown')); + const cooldown = store.put.mock.calls.find(([key]) => key.includes('cooldown')); expect(cooldown?.[2]).toBe(60_000); }); + it('writes a rate limit with its end, so a shorter one written later cannot end it', async () => { + let clock = 1_000; + const { store } = createStore(() => clock); + const long = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 3_600_000)); + const short = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED')); + const a = createPullRequestLookup({ source: { find: long }, store, now: () => clock }); + const b = createPullRequestLookup({ source: { find: short }, store, now: () => clock }); + /** At once, so neither has seen the other's limit; the long one is written first. */ + await Promise.all([a(shared), b({ ...shared, branch: 'other' })]); + await settle(); + const ends = store.put.mock.calls + .filter(([key]) => key.includes('cooldown')) + .map(([, stored]) => stored.order); + expect(ends).toEqual([1_000 + 3_600_000, 1_000 + 10_000]); + /** Past the short limit, a third replica still sees the long one. */ + clock += 60_000; + const find = jest.fn().mockResolvedValue(value); + const c = createPullRequestLookup({ source: { find }, store, now: () => clock }); + await expect(c({ ...shared, branch: 'third' })).resolves.toEqual({ + ok: false, + error: { code: 'RATE_LIMITED' }, + }); + expect(find).not.toHaveBeenCalled(); + }); + + it('writes an answer with its end, so a slower replica cannot replace a newer one', async () => { + const clock = 0; + const { store } = createStore(() => clock); + const lookup = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + now: () => clock, + }); + await lookup(shared); + await settle(); + expect(store.put.mock.calls[0][1].order).toBe(30_000); + }); + it('does not share a cooldown with lookups on another host', async () => { const { store } = createStore(Date.now); const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); @@ -565,6 +660,61 @@ describe('the shared store', () => { await expect(b({ ...shared, host: 'ghe.example.com' })).resolves.toEqual({ ok: true, value }); }); + it('finds a shared answer whatever cache capacity each caller has', async () => { + const { store } = createStore(Date.now); + const second = jest.fn().mockResolvedValue(value); + const a = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + }); + const b = createPullRequestLookup({ source: { find: second }, store }); + await a({ ...shared, cacheMaxEntries: 10, cacheMaxCredentials: 1 }); + await settle(); + await expect( + b({ ...shared, cacheMaxEntries: 5_000, cacheMaxCredentials: 64 }), + ).resolves.toEqual({ ok: true, value }); + expect(second).not.toHaveBeenCalled(); + }); + + it('holds no more results than its slot count, however many branches are looked up', async () => { + const { store, values } = createStore(Date.now); + const lookup = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + }); + for (let i = 0; i < 200; i++) { + await lookup({ ...shared, branch: `b${i}`, sharedCacheSlots: 20 }); + } + await settle(); + expect(store.put).toHaveBeenCalledTimes(200); + expect(values.size).toBeLessThanOrEqual(20); + }); + + it('holds no more rate limit cooldowns than its slot count', async () => { + const { store, values } = createStore(Date.now); + const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); + const lookup = createPullRequestLookup({ source: { find: limited }, store }); + for (let i = 0; i < 50; i++) { + await lookup({ ...shared, token: `t${i}`, sharedCacheSlots: 4 }); + } + await settle(); + expect([...values.keys()].filter((key) => key.includes('cooldown')).length).toBeLessThanOrEqual( + 4, + ); + }); + + it('gives the store its timeout as the deadline of every call', async () => { + const { store } = createStore(Date.now); + const lookup = createPullRequestLookup({ + source: { find: jest.fn().mockResolvedValue(value) }, + store, + }); + await lookup({ ...shared, sharedCacheTimeoutMs: 75 }); + await settle(); + expect(store.get.mock.calls.every((call) => call[1] === 75)).toBe(true); + expect(store.put.mock.calls[0][3]).toBe(75); + }); + it('never puts the token in a stored key or value', async () => { const { store } = createStore(Date.now); const find = jest @@ -575,19 +725,20 @@ describe('the shared store', () => { await lookup(shared); await lookup({ ...shared, branch: 'b'.repeat(5_000) }); await settle(); - expect(store.set).toHaveBeenCalledTimes(3); - for (const [key, stored] of store.set.mock.calls) { - expect(key).toMatch(/^pr:v1:/); + /** The answer and the cooldown; a failed lookup is never stored. */ + expect(store.put).toHaveBeenCalledTimes(2); + for (const [key, stored] of store.put.mock.calls) { + expect(key).toMatch(/^pr:v2:/); expect(key.length).toBeLessThan(100); expect(key).not.toContain(token); expect(JSON.stringify(stored)).not.toContain(token); } }); - it('answers from GitHub when the store fails', async () => { + it('answers from the provider when the store fails', async () => { const store: PullRequestLookupStore = { get: jest.fn().mockRejectedValue(new Error('ECONNREFUSED')), - set: jest.fn().mockRejectedValue(new Error('ECONNREFUSED')), + put: jest.fn().mockRejectedValue(new Error('ECONNREFUSED')), }; const find = jest.fn().mockResolvedValue(value); const lookup = createPullRequestLookup({ source: { find }, store }); @@ -599,7 +750,7 @@ describe('the shared store', () => { it('does not wait on a store that never answers for longer than its timeout', async () => { const hang = () => new Promise(() => undefined); - const store: PullRequestLookupStore = { get: jest.fn(hang), set: jest.fn(hang) }; + const store: PullRequestLookupStore = { get: jest.fn(hang), put: jest.fn(hang) }; const find = jest.fn().mockResolvedValue(value); const lookup = createPullRequestLookup({ source: { find }, store }); const started = Date.now(); @@ -611,62 +762,13 @@ describe('the shared store', () => { expect(find).toHaveBeenCalledTimes(1); }); - /** The key digest a lookup with `shared` writes, so a planted entry differs only as labelled. */ - async function sharedKey(): Promise { - const { store } = createStore(Date.now); - await createPullRequestLookup({ source: { find: jest.fn().mockResolvedValue(value) }, store })( - shared, - ); - await settle(); - return (store.set.mock.calls[0][1] as { key: string }).key; - } - - it('holds no more results than every partition together, however many branches are looked up', async () => { - const { store, values } = createStore(Date.now); - const find = jest.fn().mockResolvedValue(value); - const lookup = createPullRequestLookup({ source: { find }, store }); - for (let i = 0; i < 200; i++) { - await lookup({ ...shared, branch: `b${i}`, cacheMaxEntries: 10, cacheMaxCredentials: 2 }); - } - await settle(); - expect(store.set).toHaveBeenCalledTimes(200); - expect(values.size).toBeLessThanOrEqual(20); - }); - - it('holds one rate limit cooldown per credential slot', async () => { - const { store, values } = createStore(Date.now); - const limited = jest.fn().mockRejectedValue(new PullRequestSourceError('RATE_LIMITED', 60_000)); - const lookup = createPullRequestLookup({ source: { find: limited }, store }); - for (let i = 0; i < 50; i++) { - await lookup({ ...shared, token: `t${i}`, cacheMaxCredentials: 4 }); - } - await settle(); - expect([...values.keys()].filter((key) => key.includes('cooldown')).length).toBeLessThanOrEqual( - 4, - ); - }); - - it('gives the store its timeout as the deadline of every call', async () => { - const { store } = createStore(Date.now); - const lookup = createPullRequestLookup({ - source: { find: jest.fn().mockResolvedValue(value) }, - store, - }); - await lookup({ ...shared, sharedCacheTimeoutMs: 75 }); - await settle(); - expect(store.get.mock.calls.every((call) => call[1] === 75)).toBe(true); - expect(store.set.mock.calls[0][3]).toBe(75); - }); - - it('serves a well-formed entry stored under its own key', async () => { + it('serves a well-formed value stored under its own key', async () => { const key = await sharedKey(); const store: PullRequestLookupStore = { get: jest.fn(async (slot: string) => - slot.includes('cooldown') - ? undefined - : { result: { ok: true, value }, expiresAt: Number.MAX_SAFE_INTEGER, key }, + slot.includes('cooldown') ? undefined : { key, order: forever, value }, ), - set: jest.fn(async () => undefined), + put: jest.fn(async () => undefined), }; const find = jest.fn().mockResolvedValue(null); const lookup = createPullRequestLookup({ source: { find }, store }); @@ -677,51 +779,36 @@ describe('the shared store', () => { it.each([ ['nothing', undefined], ['a string', 'cached'], - [ - 'an answer for another key in the same slot', - { result: { ok: true, value }, expiresAt: Number.MAX_SAFE_INTEGER, key: 'other' }, - ], - ['an entry without an expiry', { result: { ok: true, value } }], - ['an expiry that is not a number', { result: { ok: true, value }, expiresAt: 'later' }], - ['a result without ok', { result: { value }, expiresAt: Number.MAX_SAFE_INTEGER }], - [ - 'an unknown error code', - { result: { ok: false, error: { code: 'TEAPOT' } }, expiresAt: Number.MAX_SAFE_INTEGER }, - ], - [ - 'a pull request missing fields', - { result: { ok: true, value: { number: 1 } }, expiresAt: Number.MAX_SAFE_INTEGER }, - ], + ['an answer for another key in the same slot', { key: 'other', order: forever, value }], + ['a value without an order', { value }], + ['an order that is not a number', { order: 'later', value }], + ['a value that has run out', { order: 0, value }], + ['a value missing entirely', { order: forever }], + ['a pull request missing fields', { order: forever, value: { number: 1 } }], [ 'a pull request with an unknown state', - { - result: { ok: true, value: { ...value, state: 'draft' } }, - expiresAt: Number.MAX_SAFE_INTEGER, - }, + { order: forever, value: { ...value, state: 'draft' } }, ], [ 'a pull request page that is not https', - { - result: { ok: true, value: { ...value, url: 'javascript:alert(1)' } }, - expiresAt: Number.MAX_SAFE_INTEGER, - }, + { order: forever, value: { ...value, url: 'javascript:alert(1)' } }, ], [ 'a pull request from an unknown provider', - { - result: { ok: true, value: { ...value, provider: 'gitea' } }, - expiresAt: Number.MAX_SAFE_INTEGER, - }, + { order: forever, value: { ...value, provider: 'gitea' } }, + ], + [ + 'a stored failure from an older replica', + { result: { ok: false, error: { code: 'UPSTREAM_ERROR' } }, expiresAt: forever }, ], - ['an entry that has expired', { result: { ok: true, value }, expiresAt: 0 }], ])('ignores %s in the store', async (_label, planted) => { const key = await sharedKey(); const stored = isPlainObject(planted) && !('key' in planted) ? { ...planted, key } : planted; const store: PullRequestLookupStore = { get: jest.fn(async (slot: string) => - slot.includes('cooldown') ? { until: 'soon' } : stored, + slot.includes('cooldown') ? { key, order: 'soon' } : stored, ), - set: jest.fn(async () => undefined), + put: jest.fn(async () => undefined), }; const find = jest.fn().mockResolvedValue(null); const lookup = createPullRequestLookup({ source: { find }, store }); @@ -740,7 +827,7 @@ describe('the shared store', () => { await lookup({ ...input, sharedCache: false, branch: 'b' }); await settle(); expect(store.get).not.toHaveBeenCalled(); - expect(store.set).not.toHaveBeenCalled(); + expect(store.put).not.toHaveBeenCalled(); }); it('serves a shared answer while this replica is cooling down', async () => { @@ -756,7 +843,7 @@ describe('the shared store', () => { await cooled({ ...shared, branch: 'limited' }); await settle(); await expect(cooled({ ...shared, branch: 'other' })).resolves.toEqual({ ok: true, value }); - /** Nothing stored for this branch, so the cooldown still answers without asking GitHub. */ + /** Nothing stored for this branch, so the cooldown still answers without asking the provider. */ await expect(cooled({ ...shared, branch: 'third' })).resolves.toEqual({ ok: false, error: { code: 'RATE_LIMITED' }, diff --git a/packages/api/src/pulls/lookup.ts b/packages/api/src/pulls/lookup.ts index 0360d2e2003..d0a12667827 100644 --- a/packages/api/src/pulls/lookup.ts +++ b/packages/api/src/pulls/lookup.ts @@ -6,13 +6,15 @@ import { PULL_REQUEST_MERGEABLE, PULL_REQUEST_PROVIDERS, } from 'librechat-data-provider'; +import type { TConversationPullRequest } from 'librechat-data-provider'; import type { PullRequestSource, + PullRequestSharedValue, PullRequestLookup, PullRequestLookupStore, PullRequestLookupResult, } from './types'; -import { PULL_REQUEST_ERROR_CODES, PullRequestSourceError } from './types'; +import { PullRequestSourceError } from './types'; import { getSafeErrorMetadata } from '~/utils'; /** A failed lookup is remembered briefly so a failing upstream is not hammered by every poll. */ @@ -26,13 +28,19 @@ const DEFAULT_MAX_ENTRIES = 500; const DEFAULT_MAX_PARTITIONS = 256; /** A shared store call slower than this is abandoned; Redis answers in single milliseconds. */ const DEFAULT_SHARED_TIMEOUT_MS = 250; +/** Shared store slots for results, and again for cooldowns, unless the deployment sets its own. */ +const DEFAULT_SHARED_SLOTS = 65_536; /** Versioned, so a change to the stored shape is a miss for replicas still on the old one. */ -const STORE_PREFIX = 'pr:v1:'; +const STORE_PREFIX = 'pr:v2:'; type Entry = { result: PullRequestLookupResult; expiresAt: number }; -/** `key` is the digest of the full key, since a slot is shared by every key that maps to it. */ -type Stored = Entry & { key: string }; -type Cooldown = { key: string; until: number }; +/** + * What the shared store holds. `key` is the digest of the full key, since a slot is shared by every + * key that maps to it, and `order` is when the value runs out: the store keeps the later of two + * values for one key. Only a found answer is shared, never a failure: an upstream error can be this + * replica's own network or proxy, and another replica may well reach the provider. + */ +type Stored = PullRequestSharedValue & { value: TConversationPullRequest | null }; /** Keys carry a digest, never the token, so no credential is held in a map key or a heap dump. */ const scopeOf = (token: string): string => @@ -43,7 +51,8 @@ const digestOf = (key: string): string => createHash('sha256').update(key).diges /** * The shared store is a fixed number of slots, so it holds no more than its slot count however * many repositories and branches are looked up: an answer goes in the slot its key's digest picks, - * replacing whatever was there, and a read whose slot holds another key is a miss. + * replacing whatever was there, and a read whose slot holds another key is a miss. The count is + * deployment-wide (read from the YAML only), so every caller addresses the same slot for a key. */ const slotOf = (kind: 'result' | 'cooldown', digest: string, slots: number): string => `${STORE_PREFIX}${kind}:${Number.parseInt(digest.slice(0, 12), 16) % slots}`; @@ -59,14 +68,8 @@ const isTime = (value: unknown): value is number => /** Another replica, or another version of one, wrote this; anything not shaped right is a miss. */ function isStored(value: unknown, key: string): value is Stored { - if (!isRecord(value) || value.key !== key) return false; - if (!isTime(value.expiresAt) || !isRecord(value.result)) return false; - const { result } = value; - if (result.ok === false) { - return isRecord(result.error) && isOneOf(PULL_REQUEST_ERROR_CODES, result.error.code); - } - if (result.ok !== true) return false; - const pr = result.value; + if (!isRecord(value) || value.key !== key || !isTime(value.order)) return false; + const pr = value.value; if (pr === null) return true; return ( isRecord(pr) && @@ -85,8 +88,9 @@ function isStored(value: unknown, key: string): value is Stored { ); } -const isCooldown = (value: unknown, key: string): value is Cooldown => - isRecord(value) && value.key === key && isTime(value.until); +/** A provider's rate limit on a credential, which holds for every replica until `order`. */ +const isCooldown = (value: unknown, key: string): value is PullRequestSharedValue => + isRecord(value) && value.key === key && isTime(value.order); /** * Runs one shared store call, bounded by `timeoutMs`. The store is an optimization, so a failure @@ -194,8 +198,9 @@ export function createPullRequestLookup({ } } + /** A cooldown only ever extends: a shorter hint arriving later never ends a longer one early. */ function startCooldown(scope: string, until: number): void { - cooldowns.set(scope, until); + cooldowns.set(scope, Math.max(cooldowns.get(scope) ?? 0, until)); for (const [other, end] of cooldowns) { if (end <= now()) cooldowns.delete(other); } @@ -205,12 +210,12 @@ export function createPullRequestLookup({ function share( target: PullRequestLookupStore, key: string, - value: Stored | Cooldown, + value: PullRequestSharedValue, ttlMs: number, timeoutMs: number, ): void { if (ttlMs <= 0) return; - void bounded(() => target.set(key, value, ttlMs, timeoutMs), timeoutMs); + void bounded(() => target.put(key, value, ttlMs, timeoutMs), timeoutMs); } return async ({ @@ -228,6 +233,7 @@ export function createPullRequestLookup({ cacheMaxCredentials, sharedCache, sharedCacheTimeoutMs = DEFAULT_SHARED_TIMEOUT_MS, + sharedCacheSlots = DEFAULT_SHARED_SLOTS, }) => { const scope = scopeOf(token); /** Everything that changes the answer or how long it may be reused is part of the key. */ @@ -268,11 +274,10 @@ export function createPullRequestLookup({ result.ok ? ttlMs : Math.min(ttlMs, FAILURE_TTL_MS); const capacity = cacheMaxEntries ?? maxEntries; const credentials = cacheMaxCredentials ?? DEFAULT_MAX_PARTITIONS; - /** As many results as every replica's partitions hold together, and a cooldown per credential. */ const resultDigest = digestOf(key); const cooldownDigest = digestOf(limitScope); - const storeKey = slotOf('result', resultDigest, capacity * credentials); - const cooldownKey = slotOf('cooldown', cooldownDigest, credentials); + const storeKey = slotOf('result', resultDigest, sharedCacheSlots); + const cooldownKey = slotOf('cooldown', cooldownDigest, sharedCacheSlots); const run = (async (): Promise => { if (shared != null) { @@ -280,15 +285,16 @@ export function createPullRequestLookup({ bounded(() => shared.get(storeKey, sharedCacheTimeoutMs), sharedCacheTimeoutMs), bounded(() => shared.get(cooldownKey, sharedCacheTimeoutMs), sharedCacheTimeoutMs), ]); - if (isStored(stored, resultDigest) && stored.expiresAt > now()) { + if (isStored(stored, resultDigest) && stored.order > now()) { + const result: PullRequestLookupResult = { ok: true, value: stored.value }; /** Never kept longer than this replica would have kept its own answer. */ - const expiresAt = Math.min(stored.expiresAt, now() + lifetimeOf(stored.result)); - remember(scope, key, { result: stored.result, expiresAt }, capacity, credentials); - return stored.result; + const expiresAt = Math.min(stored.order, now() + ttlMs); + remember(scope, key, { result, expiresAt }, capacity, credentials); + return result; } if (cooling) return rateLimited; - if (isCooldown(cooldown, cooldownDigest) && cooldown.until > now()) { - startCooldown(limitScope, Math.min(cooldown.until, now() + MAX_COOLDOWN_MS)); + if (isCooldown(cooldown, cooldownDigest) && cooldown.order > now()) { + startCooldown(limitScope, Math.min(cooldown.order, now() + MAX_COOLDOWN_MS)); return rateLimited; } } @@ -324,7 +330,7 @@ export function createPullRequestLookup({ share( shared, cooldownKey, - { key: cooldownDigest, until: end }, + { key: cooldownDigest, order: end }, wait, sharedCacheTimeoutMs, ); @@ -335,8 +341,9 @@ export function createPullRequestLookup({ const lifetime = lifetimeOf(result); const entry: Entry = { result, expiresAt: now() + lifetime }; remember(scope, key, entry, capacity, credentials); - if (shared != null) { - share(shared, storeKey, { ...entry, key: resultDigest }, lifetime, sharedCacheTimeoutMs); + if (shared != null && result.ok) { + const stored: Stored = { key: resultDigest, order: entry.expiresAt, value: result.value }; + share(shared, storeKey, stored, lifetime, sharedCacheTimeoutMs); } return result; })().finally(() => inflight.delete(flight)); diff --git a/packages/api/src/pulls/types.ts b/packages/api/src/pulls/types.ts index 1efd436ab37..5b48ec8de87 100644 --- a/packages/api/src/pulls/types.ts +++ b/packages/api/src/pulls/types.ts @@ -80,16 +80,24 @@ export type PullRequestLookupInput = PullRequestFindInput & { sharedCache?: boolean; /** Longest one shared store call may take before the lookup carries on without it. */ sharedCacheTimeoutMs?: number; + /** Slots the shared store spreads results over, and again cooldowns; deployment-wide. */ + sharedCacheSlots?: number; }; +/** A value in the shared store: `key` names whose it is, `order` how recent it is. */ +export type PullRequestSharedValue = { key: string; order: number }; + /** - * Storage the lookup shares with other replicas, so one answer or one rate limit from GitHub serves - * every replica. Values must round-trip as JSON; `ttlMs` is how long the store may keep one. A call - * still unsent at `timeoutMs` must be dropped, not left queued, and reject. + * Storage the lookup shares with other replicas, so one answer or one rate limit from the provider + * serves every replica. Values must round-trip as JSON; `ttlMs` is how long the store may keep one. + * A call still unsent at `timeoutMs` must be dropped, not left queued, and reject. */ export type PullRequestLookupStore = { get(key: string, timeoutMs: number): Promise; - set(key: string, value: unknown, ttlMs: number, timeoutMs: number): Promise; + /** Stores `value` unless the slot already holds one with the same `key` and an `order` at least + * as large, checked and written atomically, so a slower replica's write never replaces a newer + * answer and a shorter rate limit never ends a longer one. */ + put(key: string, value: PullRequestSharedValue, ttlMs: number, timeoutMs: number): Promise; }; export type PullRequestLookup = (input: PullRequestLookupInput) => Promise; diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index 79e26f7cee3..614d4285d17 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -1152,6 +1152,7 @@ describe('agent pull request config', () => { batchTimeoutSeconds: 20, sharedCache: false, sharedCacheTimeoutMs: 250, + sharedCacheSlots: 65_536, }); }); @@ -1269,6 +1270,9 @@ describe('agent pull request config', () => { ['sharedCacheTimeoutMs', 9], ['sharedCacheTimeoutMs', 5001], ['sharedCacheTimeoutMs', 1.5], + ['sharedCacheSlots', 63], + ['sharedCacheSlots', 1_000_001], + ['sharedCacheSlots', 1.5], ])('rejects %s of %s', (field, value) => { expect(parse({ [field]: value }).success).toBe(false); }); @@ -1290,6 +1294,8 @@ describe('agent pull request config', () => { ).toBe(true); expect(parse({ maxHeadComparisons: 0 }).success).toBe(true); expect(parse({ sharedCacheTimeoutMs: 10 }).success).toBe(true); + expect(parse({ sharedCacheSlots: 64 }).success).toBe(true); + expect(parse({ sharedCacheSlots: 1_000_000 }).success).toBe(true); }); it.each([4, 3601, 1.5])('rejects a cache lifetime of %s seconds', (cacheTtlSeconds) => { diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 2eb98a6675b..4e246a384fc 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -114,10 +114,12 @@ export const BASE_ONLY_CONFIG_SECTIONS = ['filters', 'mcpAppSandbox'] as const; export const BASE_PRINCIPAL_CONFIG_SECTIONS = ['langfuse'] as const; /** Fields inside overridable sections that only the operator's YAML sets: a role, group, user * or tenant config document can neither change nor remove them. The pull request provider and - * host decide where the deployment's token is sent. */ + * host decide where the deployment's token is sent, and the shared cache's slot count where + * every replica looks for a shared answer. */ export const BASE_ONLY_CONFIG_PATHS = [ 'endpoints.agents.pullRequests.provider', 'endpoints.agents.pullRequests.host', + 'endpoints.agents.pullRequests.sharedCacheSlots', ] as const; export const defaultRetrievalModels = [ @@ -1963,14 +1965,17 @@ export const agentsEndpointSchema = baseEndpointSchema /** Share looked-up pull requests and GitHub rate limit cooldowns between replicas * through Redis, so each replica does not spend the token on the same answer. Takes * effect only when the deployment runs with Redis (`USE_REDIS`); each replica keeps its - * own in-memory cache either way. Redis holds at most `cacheMaxEntries` times - * `cacheMaxCredentials` results and `cacheMaxCredentials` cooldowns: each goes in a - * slot picked from its key, replacing what was there. */ + * own in-memory cache either way. Only found pull requests and rate limits are shared; + * a failed lookup stays with the replica that saw it. */ sharedCache: z.boolean().optional().default(false), /** Longest one shared cache read or write may take, in milliseconds, before a lookup * carries on without it, so a slow Redis cannot hold the header request. A command * still queued at that point, as during a disconnect, is dropped rather than sent later. */ sharedCacheTimeoutMs: z.number().int().min(10).max(5000).optional().default(250), + /** Redis keys the shared cache uses at most, for results and again for rate limits: + * each goes in a slot picked from its key, replacing what was there. Read from the + * YAML only, so every replica and principal addresses the same slot for a lookup. */ + sharedCacheSlots: z.number().int().min(64).max(1_000_000).optional().default(65_536), }) .optional(), /** Conversational background-task delivery policy. Automatic completion wakeups are diff --git a/packages/data-schemas/src/app/resolution.spec.ts b/packages/data-schemas/src/app/resolution.spec.ts index 504bee0c095..c9d371fe5b7 100644 --- a/packages/data-schemas/src/app/resolution.spec.ts +++ b/packages/data-schemas/src/app/resolution.spec.ts @@ -168,6 +168,11 @@ describe('mergeConfigOverrides', () => { ]); expect(provider.provider).toBeUndefined(); + const slots = read([ + fakeConfig({ endpoints: { agents: { pullRequests: { sharedCacheSlots: 64 } } } }, 10), + ]); + expect(slots.sharedCacheSlots).toBeUndefined(); + const removed = read([fakeConfig({}, 10, ['endpoints.agents.pullRequests.host'])]); expect(removed.host).toBe('ghe.example.com');