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..4ec0007c4e1 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -23,6 +23,9 @@ const { createPullRequestSource, createProxyAwareFetch, createPullRequestLookup, + createPullRequestLookupStore, + keyvRedisClient, + cacheConfig, createConversationPullRequestsHandler, backgroundTaskRegistry, createSubagentThreadViewHandler, @@ -175,6 +178,7 @@ const backgroundTaskCancelHandler = createBackgroundTaskCancelHandler({ /** One lookup, so the header's single route and the sidebar's batch route share its cache. */ const pullRequestLookup = createPullRequestLookup({ source: createPullRequestSource({ fetchFn: createProxyAwareFetch() }), + store: createPullRequestLookupStore({ client: keyvRedisClient, config: cacheConfig }), }); const conversationPullRequestHandler = createConversationPullRequestHandler({ getConvoLaneGit: db.getConvoLaneGit, diff --git a/librechat.example.yaml b/librechat.example.yaml index 884fdcac373..2a381add6a2 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -977,6 +977,15 @@ endpoints: # maxCandidatePullRequests: 10 # pull requests listed per state when matching a reused branch (1 to 100) # maxCandidatePages: 1 # pages of those read per state, for branch names reused more often (1 to 10) # maxHeadComparisons: 3 # candidates compared with the chat's last commit (0 to 20) + # # Share looked-up pull requests and GitHub rate limit cooldowns between replicas, so they + # # do not each spend the token on the same answer. Only takes effect with USE_REDIS. Failed + # # lookups are never shared. + # sharedCache: false + # # Milliseconds one shared cache call may take before a lookup carries on without it (10 to 5000). + # sharedCacheTimeoutMs: 250 + # # Redis keys the shared cache uses at most, for results and again for rate limits (64 to + # # 1000000). Read from this file only. + # sharedCacheSlots: 65536 # # "run_in_background" makes Code Interpreter tools eligible by default and enables per-tool MCP opt-in. # # "tool_intents" enables live model-written labels for native tools and opted-in MCP tools. # # (optional) Require user approval before matching tool calls. Disabled by default. 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/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..9b64be25dda --- /dev/null +++ b/packages/api/src/cache/pullRequests.spec.ts @@ -0,0 +1,96 @@ +import { createServer } from 'node:net'; +import { CacheKeys } from 'librechat-data-provider'; +import { createClient, TimeoutError } from '@redis/client'; +import type { PullRequestStoreClient } from './pullRequests'; +import type { PullRequestStoreConfig } from './pullRequests'; +import { createPullRequestLookupStore } from './pullRequests'; + +const base: PullRequestStoreConfig = { + USE_REDIS: true, + FORCED_IN_MEMORY_CACHE_NAMESPACES: [], + REDIS_KEY_PREFIX: '', + GLOBAL_PREFIX_SEPARATOR: '::', +}; + +function fakeClient() { + const values = new Map(); + const commands = { + get: jest.fn(async (key: string) => values.get(key) ?? null), + eval: jest.fn(async (_script: string, options: { keys: string[]; arguments: string[] }) => { + values.set(options.keys[0], options.arguments[0]); + return 1; + }), + }; + const client = { withCommandOptions: jest.fn(() => commands) }; + return { client, commands, values }; +} + +/** A port nothing listens on, so a client there stays disconnected. */ +async function closedPort(): Promise { + 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', () => { + it('is absent without Redis, so nothing is shared', () => { + const config = { ...base, USE_REDIS: false }; + expect(createPullRequestLookupStore({ client: fakeClient().client, config })).toBeUndefined(); + }); + + it('is absent without a Redis client', () => { + expect(createPullRequestLookupStore({ client: null, config: base })).toBeUndefined(); + }); + + it('is absent when the namespace is forced in memory', () => { + const config = { ...base, FORCED_IN_MEMORY_CACHE_NAMESPACES: [CacheKeys.PULL_REQUESTS] }; + expect(createPullRequestLookupStore({ client: fakeClient().client, config })).toBeUndefined(); + }); + + it('round-trips JSON under the namespace with the lifetime and the deadline', async () => { + const { client, commands } = fakeClient(); + const config = { ...base, REDIS_KEY_PREFIX: 'tenant' }; + const store = createPullRequestLookupStore({ client, config }); + await store?.put('k', { key: 'owner', order: 2 }, 5_000.4, 40); + const key = `tenant::${CacheKeys.PULL_REQUESTS}:k`; + expect(commands.eval).toHaveBeenCalledWith(expect.any(String), { + keys: [key], + arguments: ['{"key":"owner","order":2}', '5001', 'owner', '2'], + }); + expect(client.withCommandOptions).toHaveBeenLastCalledWith({ timeout: 40 }); + await expect(store?.get('k', 30)).resolves.toEqual({ key: 'owner', order: 2 }); + expect(client.withCommandOptions).toHaveBeenLastCalledWith({ timeout: 30 }); + await expect(store?.get('missing', 30)).resolves.toBeUndefined(); + }); + + it('rejects a failed command', async () => { + const { client, commands } = fakeClient(); + commands.get.mockRejectedValueOnce(new Error('down')); + const store = createPullRequestLookupStore({ client, config: base }); + await expect(store?.get('k', 30)).rejects.toThrow('down'); + }); + + it('drops a command queued while Redis is unreachable at its deadline', async () => { + const redis = createClient({ + socket: { host: '127.0.0.1', port: await closedPort(), reconnectStrategy: () => 60_000 }, + }); + redis.on('error', () => undefined); + /** Never settles: the client keeps waiting to reconnect until it is destroyed. */ + void redis.connect().catch(() => undefined); + try { + /** The production client, typed as the store sees it, with no cast. */ + const client: PullRequestStoreClient = redis; + const store = createPullRequestLookupStore({ client, config: base }); + const started = Date.now(); + await expect(store?.get('k', 50)).rejects.toBeInstanceOf(TimeoutError); + const value = { key: 'owner', order: 1 }; + await expect(store?.put('k', value, 1_000, 50)).rejects.toBeInstanceOf(TimeoutError); + expect(Date.now() - started).toBeLessThan(2_000); + } finally { + redis.destroy(); + } + }); +}); diff --git a/packages/api/src/cache/pullRequests.ts b/packages/api/src/cache/pullRequests.ts new file mode 100644 index 00000000000..fc2512b1bbf --- /dev/null +++ b/packages/api/src/cache/pullRequests.ts @@ -0,0 +1,82 @@ +import { CacheKeys } from 'librechat-data-provider'; +import type { PullRequestLookupStore } from '~/pulls/types'; +import type { cacheConfig } from './cacheConfig'; + +/** The commands the store sends; the single-node and the cluster client both provide them. */ +type StoreCommands = { + get(key: string): Promise; + 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; +}; + +/** The deployment's cache settings the store follows: whether Redis is on, and its key prefix. */ +export type PullRequestStoreConfig = Pick< + typeof cacheConfig, + 'USE_REDIS' | 'FORCED_IN_MEMORY_CACHE_NAMESPACES' | 'REDIS_KEY_PREFIX' | 'GLOBAL_PREFIX_SEPARATOR' +>; + +/** + * The store pull request lookups share between replicas, or `undefined` when there is nothing to + * share it through: without Redis, or with the namespace forced in memory, each replica's own + * cache already is the whole cache. + * + * Every command carries the caller's timeout as its deadline. A command still unsent when it + * passes, such as one held in the offline queue while Redis is disconnected, is taken out of the + * queue and rejected, so a Redis outage does not pile up commands to replay on reconnect. + */ +export function createPullRequestLookupStore({ + client, + config, +}: { + client: PullRequestStoreClient | null; + config: PullRequestStoreConfig; +}): PullRequestLookupStore | undefined { + if ( + client == null || + !config.USE_REDIS || + config.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.PULL_REQUESTS) + ) { + return undefined; + } + const prefix = config.REDIS_KEY_PREFIX + ? `${config.REDIS_KEY_PREFIX}${config.GLOBAL_PREFIX_SEPARATOR}${CacheKeys.PULL_REQUESTS}:` + : `${CacheKeys.PULL_REQUESTS}:`; + return { + get: async (key, timeoutMs) => { + const raw = await client.withCommandOptions({ timeout: timeoutMs }).get(prefix + key); + return raw == null ? undefined : JSON.parse(raw); + }, + put: async (key, value, ttlMs, timeoutMs) => { + await client.withCommandOptions({ timeout: timeoutMs }).eval(PUT_SCRIPT, { + keys: [prefix + key], + arguments: [ + JSON.stringify(value), + String(Math.max(1, Math.ceil(ttlMs))), + value.key, + String(value.order), + ], + }); + }, + }; +} diff --git a/packages/api/src/pulls/handler.spec.ts b/packages/api/src/pulls/handler.spec.ts index 80c28de9534..a835a9f33c9 100644 --- a/packages/api/src/pulls/handler.spec.ts +++ b/packages/api/src/pulls/handler.spec.ts @@ -95,6 +95,9 @@ describe('createConversationPullRequestHandler', () => { ttlMs: 30_000, cacheMaxEntries: 500, cacheMaxCredentials: 256, + sharedCache: false, + sharedCacheTimeoutMs: 250, + sharedCacheSlots: 65_536, allowedRepositories: ['o/r'], limits: { requestTimeoutMs: 10_000, @@ -1115,6 +1118,43 @@ describe('createConversationPullRequestsHandler', () => { expect(hosts).toEqual(['github.com', 'ghe.example.com']); }); + it.each([ + ['sharing setting', { sharedCache: true }], + ['shared cache deadline', { sharedCache: true, sharedCacheTimeoutMs: 10 }], + ['cache capacity', { cacheMaxEntries: 50 }], + ])('does not join a lookup started with another %s', async (_label, changed) => { + const inputs: unknown[] = []; + const lookup = jest.fn( + (input: unknown) => + new Promise(() => { + inputs.push(input); + }), + ); + const settings = { ...enabled, batchTimeoutSeconds: 1 }; + const getAppConfig = jest + .fn() + .mockResolvedValueOnce(configWith(settings)) + .mockResolvedValueOnce(configWith(settings)) + .mockResolvedValue(configWith({ ...settings, ...changed })); + const handler = createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn(async () => [lane('a', 'o/r', 'feat/shared')]), + getAppConfig, + lookup: lookup as never, + env: { GH_TOKEN: 'ghp_secret' }, + }); + const call = () => + handler( + { user: { id: 'u1' }, body: { conversationIds: ['a'] } } as unknown as ServerRequest, + { status: jest.fn().mockReturnThis(), json: jest.fn() } as unknown as Response, + ); + void call(); + await flush(); + void call(); + await flush(); + expect(inputs).toHaveLength(2); + expect(inputs[1]).toMatchObject(changed); + }); + it('does not join a lookup started for another provider on the same host', async () => { const providers: Array = []; const lookup = jest.fn( diff --git a/packages/api/src/pulls/handler.ts b/packages/api/src/pulls/handler.ts index 4f7a190dc26..3402a0bf58f 100644 --- a/packages/api/src/pulls/handler.ts +++ b/packages/api/src/pulls/handler.ts @@ -229,6 +229,9 @@ function lookupInput(settings: PullRequestSettings, token: string, lane: Eligibl ttlMs: (settings.cacheTtlSeconds ?? 30) * 1000, cacheMaxEntries: settings.cacheMaxEntries ?? 500, cacheMaxCredentials: settings.cacheMaxCredentials ?? 256, + sharedCache: settings.sharedCache === true, + sharedCacheTimeoutMs: settings.sharedCacheTimeoutMs ?? 250, + sharedCacheSlots: settings.sharedCacheSlots ?? 65_536, allowedRepositories: repositoryScope(settings), limits: { requestTimeoutMs: (settings.requestTimeoutSeconds ?? 10) * 1000, @@ -476,6 +479,11 @@ export function createConversationPullRequestsHandler(deps: { input.ttlMs, input.limits, input.allowedRepositories, + /** Where the lookup caches and shares its answer, so each caller's own policy holds. */ + input.cacheMaxEntries, + input.cacheMaxCredentials, + input.sharedCache, + input.sharedCacheTimeoutMs, ])}`; const joined = sharedLookups.join(shareKey, scope, limit, input); const result = await withDeadline(joined.promise, remaining); diff --git a/packages/api/src/pulls/lookup.spec.ts b/packages/api/src/pulls/lookup.spec.ts index 76bd6e409c5..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 { PullRequestSource } from './types'; +import type { PullRequestSource, PullRequestLookupStore, PullRequestSharedValue } from './types'; import { createPullRequestLookup } from './lookup'; import { PullRequestSourceError } from './types'; @@ -475,3 +475,408 @@ 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 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 store = { + 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; + }), + 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 }; + } + + /** Background writes settle on the next turns of the event loop. */ + const settle = () => new Promise((resolve) => setImmediate(resolve)); + + /** 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); + 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('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); + 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.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)); + 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 answer that answers. */ + await expect(b({ ...shared, branch: 'other' })).resolves.toEqual({ + ok: false, + error: { code: 'RATE_LIMITED' }, + }); + expect(find).not.toHaveBeenCalled(); + 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)); + 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('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 + .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(); + /** 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 the provider when the store fails', async () => { + const store: PullRequestLookupStore = { + get: 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 }); + 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), put: 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('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 : { key, order: forever, value }, + ), + put: 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', { 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', + { order: forever, value: { ...value, state: 'draft' } }, + ], + [ + 'a pull request page that is not https', + { order: forever, value: { ...value, url: 'javascript:alert(1)' } }, + ], + [ + 'a pull request from an unknown provider', + { order: forever, value: { ...value, provider: 'gitea' } }, + ], + [ + 'a stored failure from an older replica', + { result: { ok: false, error: { code: 'UPSTREAM_ERROR' } }, expiresAt: forever }, + ], + ])('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') ? { key, order: 'soon' } : stored, + ), + put: 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.put).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 the provider. */ + 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({ + 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..d0a12667827 100644 --- a/packages/api/src/pulls/lookup.ts +++ b/packages/api/src/pulls/lookup.ts @@ -1,6 +1,19 @@ import { createHash } from 'node:crypto'; import { logger } from '@librechat/data-schemas'; -import type { PullRequestLookup, PullRequestLookupResult, PullRequestSource } from './types'; +import { + PULL_REQUEST_CHECKS, + PULL_REQUEST_STATES, + 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 { PullRequestSourceError } from './types'; import { getSafeErrorMetadata } from '~/utils'; @@ -13,13 +26,96 @@ 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; +/** 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:v2:'; type Entry = { result: PullRequestLookupResult; expiresAt: 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 => createHash('sha256').update(token).digest('hex').slice(0, 32); +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. 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}`; + +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 isStored(value: unknown, key: string): value is Stored { + if (!isRecord(value) || value.key !== key || !isTime(value.order)) return false; + const pr = value.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) + ); +} + +/** 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 + * 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; + 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 +124,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 +182,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 +198,26 @@ 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); + /** 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, Math.max(cooldowns.get(scope) ?? 0, 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: PullRequestSharedValue, + ttlMs: number, + timeoutMs: number, + ): void { + if (ttlMs <= 0) return; + void bounded(() => target.put(key, value, ttlMs, timeoutMs), timeoutMs); + } + return async ({ repo, branch, @@ -119,6 +231,9 @@ export function createPullRequestLookup({ allowedRepositories, cacheMaxEntries, 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. */ @@ -139,17 +254,50 @@ 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); } + /** 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; + const credentials = cacheMaxCredentials ?? DEFAULT_MAX_PARTITIONS; + const resultDigest = digestOf(key); + const cooldownDigest = digestOf(limitScope); + const storeKey = slotOf('result', resultDigest, sharedCacheSlots); + const cooldownKey = slotOf('cooldown', cooldownDigest, sharedCacheSlots); const run = (async (): Promise => { + if (shared != null) { + const [stored, cooldown] = await Promise.all([ + bounded(() => shared.get(storeKey, sharedCacheTimeoutMs), sharedCacheTimeoutMs), + bounded(() => shared.get(cooldownKey, sharedCacheTimeoutMs), sharedCacheTimeoutMs), + ]); + 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.order, now() + ttlMs); + remember(scope, key, { result, expiresAt }, capacity, credentials); + return result; + } + if (cooling) return rateLimited; + if (isCooldown(cooldown, cooldownDigest) && cooldown.order > now()) { + startCooldown(limitScope, Math.min(cooldown.order, now() + MAX_COOLDOWN_MS)); + return rateLimited; + } + } let result: PullRequestLookupResult; try { result = { @@ -172,21 +320,34 @@ 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), + ); + const end = now() + wait; + startCooldown(limitScope, end); + if (shared != null) { + share( + shared, + cooldownKey, + { key: cooldownDigest, order: end }, + 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 && result.ok) { + const stored: Stored = { key: resultDigest, order: entry.expiresAt, value: result.value }; + share(shared, storeKey, stored, lifetime, sharedCacheTimeoutMs); + } return result; - })().finally(() => inflight.delete(key)); - inflight.set(key, run); + })().finally(() => inflight.delete(flight)); + inflight.set(flight, run); return run; }; } diff --git a/packages/api/src/pulls/types.ts b/packages/api/src/pulls/types.ts index e0755adcec6..5b48ec8de87 100644 --- a/packages/api/src/pulls/types.ts +++ b/packages/api/src/pulls/types.ts @@ -75,6 +75,29 @@ 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; + /** 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 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; + /** 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 034d3620900..614d4285d17 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -1150,9 +1150,18 @@ describe('agent pull request config', () => { cacheMaxCredentials: 256, maxConcurrentLookups: 4, batchTimeoutSeconds: 20, + sharedCache: false, + sharedCacheTimeoutMs: 250, + sharedCacheSlots: 65_536, }); }); + 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 +1267,12 @@ describe('agent pull request config', () => { ['batchTimeoutSeconds', 0], ['batchTimeoutSeconds', 121], ['batchTimeoutSeconds', 1.5], + ['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); }); @@ -1274,9 +1289,13 @@ 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); + 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 ecb9f888d02..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 = [ @@ -1960,6 +1962,20 @@ 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. 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 @@ -4698,6 +4714,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. */ 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');