diff --git a/api/server/routes/__test-utils__/convos-route-mocks.js b/api/server/routes/__test-utils__/convos-route-mocks.js index ce6a630d933..402c0877dda 100644 --- a/api/server/routes/__test-utils__/convos-route-mocks.js +++ b/api/server/routes/__test-utils__/convos-route-mocks.js @@ -148,7 +148,9 @@ module.exports = { ), createBackgroundTaskPolicyMiddleware: jest.fn(() => (_req, _res, next) => next()), createGitHubPullRequestSource: jest.fn(() => ({ find: jest.fn() })), + createProxyAwareFetch: jest.fn(() => jest.fn()), createPullRequestLookup: jest.fn(() => jest.fn()), + createConversationPullRequestsHandler: jest.fn(() => jest.fn()), createConversationPullRequestHandler: jest.fn( () => (_req, res) => res.status(200).json({ pullRequest: null }), ), diff --git a/api/server/routes/__tests__/config.spec.js b/api/server/routes/__tests__/config.spec.js index 07f4de42348..c6386d6867d 100644 --- a/api/server/routes/__tests__/config.spec.js +++ b/api/server/routes/__tests__/config.spec.js @@ -181,6 +181,7 @@ describe('GET /api/config', () => { expect(response.body).not.toHaveProperty('conversationImportMaxFileSize'); expect(response.body).not.toHaveProperty('insightsEnabled'); expect(response.body).not.toHaveProperty('pullRequestsEnabled'); + expect(response.body).not.toHaveProperty('pullRequestsBatchVersion'); expect(response.body).not.toHaveProperty('mcpApps'); }); @@ -683,9 +684,26 @@ describe('GET /api/config', () => { }); const response = await request(createApp(mockUser)).get('/api/config'); expect(response.body.pullRequestsEnabled).toBe(expected); + /** The batch route's version rides with the flag, so a client never sees one without the other. */ + if (expected) { + expect(response.body.pullRequestsBatchVersion).toBe(1); + expect(response.body.pullRequestsMaxConcurrentLookups).toBe(4); + } else { + expect(response.body).not.toHaveProperty('pullRequestsBatchVersion'); + expect(response.body).not.toHaveProperty('pullRequestsMaxConcurrentLookups'); + } }, ); + it('advertises the configured lookup limit so a fallback client can keep to it', async () => { + mockGetAppConfig.mockResolvedValue({ + ...baseAppConfig, + endpoints: { agents: { pullRequests: { enabled: true, maxConcurrentLookups: 2 } } }, + }); + const response = await request(createApp(mockUser)).get('/api/config'); + expect(response.body.pullRequestsMaxConcurrentLookups).toBe(2); + }); + it('should not advertise pull requests for a config with no endpoints', async () => { mockGetAppConfig.mockResolvedValue(baseAppConfig); const response = await request(createApp(mockUser)).get('/api/config'); diff --git a/api/server/routes/config.js b/api/server/routes/config.js index 3336a94d4c7..62b0b7eb191 100644 --- a/api/server/routes/config.js +++ b/api/server/routes/config.js @@ -23,6 +23,7 @@ const { resolveCodeWorkspaceInheritanceCapability, resolveCodeEnvironmentTransitionVersion, loadConversationListLimits, + resolvePullRequestCapabilities, } = require('@librechat/api'); const { DEFAULT_MCP_APP_CSP_LIMITS, @@ -331,8 +332,7 @@ router.get('/', async function (req, res) { langfuseConnectionAccess, insightsEnabled: isEnabled(process.env.ENABLE_INSIGHTS), /** Lets the client skip the pull request lookup entirely when the feature is off. */ - pullRequestsEnabled: - appConfig?.endpoints?.[EModelEndpoint.agents]?.pullRequests?.enabled === true, + ...resolvePullRequestCapabilities(appConfig), compactionEnabled: appConfig?.summarization?.enabled !== false, ...(codeEnvironmentDecisionVersion != null ? { codeEnvironmentDecisionVersion } : {}), mcpApps: resolveMCPAppsPolicy( diff --git a/api/server/routes/convos.js b/api/server/routes/convos.js index 4c05abb57f6..ef6b220150c 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -21,7 +21,9 @@ const { createBackgroundTaskPolicyMiddleware, createConversationPullRequestHandler, createGitHubPullRequestSource, + createProxyAwareFetch, createPullRequestLookup, + createConversationPullRequestsHandler, backgroundTaskRegistry, createSubagentThreadViewHandler, createGeneratedTitleHandler, @@ -170,10 +172,20 @@ const backgroundTaskIndexHandler = createBackgroundTaskIndexHandler({ const backgroundTaskCancelHandler = createBackgroundTaskCancelHandler({ registry: backgroundTaskRegistry, }); +/** One lookup, so the header's single route and the sidebar's batch route share its cache. */ +const pullRequestLookup = createPullRequestLookup({ + source: createGitHubPullRequestSource({ fetchFn: createProxyAwareFetch() }), +}); const conversationPullRequestHandler = createConversationPullRequestHandler({ getConvoLaneGit: db.getConvoLaneGit, getAppConfig, - lookup: createPullRequestLookup({ source: createGitHubPullRequestSource({ fetchFn: fetch }) }), + lookup: pullRequestLookup, + env: process.env, +}); +const conversationPullRequestsHandler = createConversationPullRequestsHandler({ + getConvosLaneGit: db.getConvosLaneGit, + getAppConfig, + lookup: pullRequestLookup, env: process.env, }); router.use(requireJwtAuth); @@ -248,6 +260,7 @@ router.post( subagentControlHandler, ); router.get('/:parentConversationId/subagents', parentSubagentIndexHandler); +router.post('/pull-requests', conversationPullRequestsHandler); router.get('/:conversationId/pull-request', conversationPullRequestHandler); router.get('/:conversationId/background-tasks', backgroundTaskPolicy, backgroundTaskIndexHandler); router.post( diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 632d884eca7..446fe540048 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -2376,6 +2376,7 @@ async function loadToolsForExecution({ environmentId: codeExecutionContext.codeWorkspace.environmentId, workspaceId: codeExecutionContext.codeWorkspace.workspaceId, }, + admittedEpoch: req.resolvedConversation?.codeAttachmentEpoch, getConvoLaneContext, reserveConvoLaneGitSeq, setConvoLaneGit, diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index 82f98c1dc7f..213827e73f3 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -3236,15 +3236,9 @@ describe('ToolService - Action Capability Gating', () => { codeWorkspaces: [{ environmentId: 'personal-machine', workspaceId: 'project-a' }], }; mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); - mockResolveCodeExecutionContext.mockReturnValueOnce({ - baseUrl: 'http://attached-code.test/v1', - codeSessionKey: 'execute_code:stateful:attached', - executionProfile: 'stateful', - statefulSessions: true, - environmentType: 'attached', - environmentId: 'personal-machine', - bridgeWorkerId: 'worker-abc', - }); + req.config.endpoints[EModelEndpoint.agents].statefulCodeSessions = { + environments: [attachedEnvironment()], + }; mockCreateLaneGitRecorder.mockClear(); await loadToolsForExecution({ @@ -3255,6 +3249,7 @@ describe('ToolService - Action Capability Gating', () => { tools: [Tools.execute_code], stateful_code_sessions: true, stateful_code_environment: 'agent-user', + code_environment_id: 'personal-machine', }, conversationId: 'resolved-convo', toolNames: [AgentConstants.BASH_TOOL], @@ -3279,6 +3274,55 @@ describe('ToolService - Action Capability Gating', () => { expect(mockSetConvoLaneGit).toHaveBeenCalledWith(input); }); + it.each([ + ['a moved conversation', 3, 3], + ['a conversation never moved', 0, 0], + ['a conversation whose admitted read carried no epoch', undefined, undefined], + ])( + 'fences the lane report by the epoch of the admitted decision for %s', + async (_label, admitted, expected) => { + const capabilities = [ + AgentCapabilities.tools, + AgentCapabilities.execute_code, + AgentCapabilities.stateful_code_sessions, + ]; + const req = createMockReq(capabilities); + req.config.endpoints[EModelEndpoint.agents].pullRequests = { enabled: true }; + req.resolvedConversation = { + conversationId: 'resolved-convo', + codeWorkspaces: [{ environmentId: 'personal-machine', workspaceId: 'project-a' }], + ...(admitted === undefined ? {} : { codeAttachmentEpoch: admitted }), + }; + req.body = { + conversationId: 'resolved-convo', + codeWorkspaces: [{ environmentId: 'personal-machine', workspaceId: 'project-a' }], + }; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + req.config.endpoints[EModelEndpoint.agents].statefulCodeSessions = { + environments: [attachedEnvironment()], + }; + mockCreateLaneGitRecorder.mockClear(); + + await loadToolsForExecution({ + req, + res: {}, + agent: { + id: 'attached-agent', + tools: [Tools.execute_code], + stateful_code_sessions: true, + stateful_code_environment: 'agent-user', + code_environment_id: 'personal-machine', + }, + conversationId: 'resolved-convo', + toolNames: [AgentConstants.BASH_TOOL], + toolRegistry: new Map([[AgentConstants.BASH_TOOL, { name: AgentConstants.BASH_TOOL }]]), + actionsEnabled: false, + }); + + expect(mockCreateLaneGitRecorder.mock.calls[0][0].admittedEpoch).toBe(expected); + }, + ); + it.each([ ['unset', undefined], ['disabled', { enabled: false }], @@ -3294,15 +3338,9 @@ describe('ToolService - Action Capability Gating', () => { codeWorkspaces: [{ environmentId: 'personal-machine', workspaceId: 'project-a' }], }; mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); - mockResolveCodeExecutionContext.mockReturnValueOnce({ - baseUrl: 'http://attached-code.test/v1', - codeSessionKey: 'execute_code:stateful:attached', - executionProfile: 'stateful', - statefulSessions: true, - environmentType: 'attached', - environmentId: 'personal-machine', - bridgeWorkerId: 'worker-abc', - }); + req.config.endpoints[EModelEndpoint.agents].statefulCodeSessions = { + environments: [attachedEnvironment()], + }; mockCreateLaneGitRecorder.mockClear(); await loadToolsForExecution({ @@ -3313,6 +3351,7 @@ describe('ToolService - Action Capability Gating', () => { tools: [Tools.execute_code], stateful_code_sessions: true, stateful_code_environment: 'agent-user', + code_environment_id: 'personal-machine', }, conversationId: 'resolved-convo', toolNames: [AgentConstants.BASH_TOOL], diff --git a/client/src/components/Chat/PullRequest/Chip.tsx b/client/src/components/Chat/PullRequest/Chip.tsx index 18c3bde4c15..fbd5763748c 100644 --- a/client/src/components/Chat/PullRequest/Chip.tsx +++ b/client/src/components/Chat/PullRequest/Chip.tsx @@ -5,31 +5,10 @@ import { useAgentsMapContext, useChatContext } from '~/Providers'; import { TONE_DOT_CLASS, presentPullRequest } from './status'; import { URLIcon } from '~/components/Endpoints/URLIcon'; import { summarizePullRequest } from './summary'; +import PullRequestPanel from './Panel'; import { useLocalize } from '~/hooks'; import PullRequestIcon from './Icon'; -import PullRequestCard from './Card'; -import { cn } from '~/utils'; - -/** Slides in from the chip's side, left to right, and fades; reduced motion only fades. */ -const cardClass = cn( - 'border-border-light bg-surface-secondary text-text-primary z-[200] w-80 max-w-[calc(100vw-2rem)] rounded-xl border shadow-lg focus:outline-none', - 'origin-left -translate-x-3 opacity-0 transition duration-200 ease-out', - 'data-[enter]:translate-x-0 data-[enter]:opacity-100', - 'data-[leave]:-translate-x-3 data-[leave]:opacity-0', - 'motion-reduce:translate-x-0 motion-reduce:transition-opacity', -); - -function CiDot({ dotClass }: { dotClass: string }) { - return ( - - ); -} +import CiDot from './CiDot'; /** The agent's picture with the CI dot on its corner. */ function AgentAvatar({ @@ -72,7 +51,6 @@ function PullRequestChip({ conversationId }: { conversationId: string }) { hideTimeout: 150, }); const open = Ariakit.useStoreState(hovercard, 'open'); - const cardElement = Ariakit.useStoreState(hovercard, 'contentElement'); const { data, isError, refetch } = useConversationPullRequestQuery(conversationId); const pullRequest = data?.pullRequest; @@ -108,21 +86,12 @@ function PullRequestChip({ conversationId }: { conversationId: string }) { } /> - - void refetch()} - portalElement={cardElement} - /> - + void refetch()} + /> ); } diff --git a/client/src/components/Chat/PullRequest/CiDot.tsx b/client/src/components/Chat/PullRequest/CiDot.tsx new file mode 100644 index 00000000000..0e1e522c376 --- /dev/null +++ b/client/src/components/Chat/PullRequest/CiDot.tsx @@ -0,0 +1,25 @@ +import { cn } from '~/utils'; + +/** + * The CI status dot that sits on the corner of the picture or icon it describes. The ring is + * painted in the surface the dot sits on, so the caller names that surface. + */ +export default function CiDot({ + dotClass, + ringClassName = 'ring-presentation', +}: { + dotClass: string; + ringClassName?: string; +}) { + return ( + + ); +} diff --git a/client/src/components/Chat/PullRequest/Panel.tsx b/client/src/components/Chat/PullRequest/Panel.tsx new file mode 100644 index 00000000000..b9881e94071 --- /dev/null +++ b/client/src/components/Chat/PullRequest/Panel.tsx @@ -0,0 +1,58 @@ +import * as Ariakit from '@ariakit/react'; +import type { TConversationPullRequest } from 'librechat-data-provider'; +import { useLocalize } from '~/hooks'; +import PullRequestCard from './Card'; +import { cn } from '~/utils'; + +/** Slides in from the anchor's side, left to right, and fades; reduced motion only fades. */ +export const panelClass = cn( + 'border-border-light bg-surface-secondary text-text-primary z-[200] w-80 max-w-[calc(100vw-2rem)] rounded-xl border shadow-lg focus:outline-none', + 'origin-left -translate-x-3 opacity-0 transition duration-200 ease-out', + 'data-[enter]:translate-x-0 data-[enter]:opacity-100', + 'data-[leave]:-translate-x-3 data-[leave]:opacity-0', + 'motion-reduce:translate-x-0 motion-reduce:transition-opacity', +); + +/** + * The card as a hovercard beside its anchor. It is portaled, and a portal's events still bubble + * to the React tree that rendered it, so anything that sits inside a clickable row stops them + * here: a click on the card must not open the row it hangs from. + */ +export default function PullRequestPanel({ + store, + pullRequest, + refreshFailed, + onRetry, +}: { + store: Ariakit.HovercardStore; + pullRequest: TConversationPullRequest; + refreshFailed: boolean; + onRetry: () => void; +}) { + const localize = useLocalize(); + const panelElement = Ariakit.useStoreState(store, 'contentElement'); + const stop = (event: { stopPropagation: () => void }) => event.stopPropagation(); + + return ( + + + + ); +} diff --git a/client/src/components/Chat/PullRequest/RowMark.tsx b/client/src/components/Chat/PullRequest/RowMark.tsx new file mode 100644 index 00000000000..918b11746b6 --- /dev/null +++ b/client/src/components/Chat/PullRequest/RowMark.tsx @@ -0,0 +1,150 @@ +import { memo, useEffect } from 'react'; +import * as Ariakit from '@ariakit/react'; +import { Button } from '@librechat/client'; +import { TriangleAlert } from 'lucide-react'; +import type React from 'react'; +import { useRowPullRequestQuery } from '~/data-provider/PullRequest'; +import { TONE_DOT_CLASS, presentPullRequest } from './status'; +import PullRequestPanel, { panelClass } from './Panel'; +import { summarizePullRequest } from './summary'; +import { useLocalize } from '~/hooks'; +import PullRequestIcon from './Icon'; +import { cn } from '~/utils'; +import CiDot from './CiDot'; + +/** The row's own surface at each state, so the dot's ring blends into it. */ +const RING_SELECTED = 'ring-surface-nav-selected'; +const RING_IDLE = 'ring-surface-primary-alt group-hover:ring-surface-nav-hover'; + +/** + * The pull request a conversation opened, as a mark in its sidebar row: the state icon with the + * CI dot on its corner. It renders nothing until a pull request is known, so a row without one + * is exactly what it was. The row's title stays the conversation's own title; this only adds + * the icon. The mark is a button beside the row's own, so Tab reaches it: focusing it opens the + * card, which holds the GitHub link, and activating it must not also open the row. + */ +function PullRequestRowMark({ + conversationId, + labelId, + selected, + onDescribed, +}: { + conversationId: string; + /** The id the row lists in `aria-describedby`, so a screen reader hears what the colors say. */ + labelId: string; + selected: boolean; + /** Tells the row whether `labelId` is in the page, so it never points at text that is not. */ + onDescribed?: (described: boolean) => void; +}) { + const localize = useLocalize(); + const store = Ariakit.useHovercardStore({ + placement: 'right-start', + showTimeout: 100, + hideTimeout: 150, + }); + const open = Ariakit.useStoreState(store, 'open'); + const { data, isError, refetch } = useRowPullRequestQuery(conversationId); + const pullRequest = data?.pullRequest; + const described = pullRequest != null || isError; + + useEffect(() => { + onDescribed?.(described); + return () => onDescribed?.(false); + }, [described, onDescribed]); + + /* A first lookup that failed is not the same as a chat with no pull request: say so, and + offer a retry, instead of leaving the row looking like there is nothing to find. */ + if (pullRequest == null && isError) { + const failed = localize('com_ui_pr_load_failed'); + return ( + + } + aria-label={failed} + aria-expanded={open} + data-testid="convo-pull-request-failed" + onFocus={() => store.show()} + onClick={(event: React.MouseEvent) => { + event.stopPropagation(); + store.show(); + }} + onKeyDown={(event: React.KeyboardEvent) => { + if (event.key === 'Enter' || event.key === ' ') event.stopPropagation(); + }} + className="focus-visible:ring-text-primary text-status-warning relative mr-1 flex size-4 shrink-0 items-center justify-center rounded-sm outline-hidden focus-visible:ring-2" + > + + + + {failed} + + event.stopPropagation()} + onDoubleClick={(event) => event.stopPropagation()} + onContextMenu={(event) => event.stopPropagation()} + onKeyDown={(event) => event.stopPropagation()} + > + + {failed} + void refetch()}> + {localize('com_ui_retry')} + + + + + ); + } + + if (pullRequest == null) return null; + + const view = presentPullRequest(pullRequest); + const dotClass = view.dotTone == null ? null : TONE_DOT_CLASS[view.dotTone]; + const summary = summarizePullRequest(pullRequest, localize); + + return ( + + } + aria-label={summary} + aria-expanded={open} + data-testid="convo-pull-request" + onFocus={() => store.show()} + onClick={(event: React.MouseEvent) => { + /* The row itself opens the conversation on click; this opens the card instead. */ + event.stopPropagation(); + store.show(); + }} + onKeyDown={(event: React.KeyboardEvent) => { + /* Enter and Space belong to this button, not to the row's own shortcuts. */ + if (event.key === 'Enter' || event.key === ' ') event.stopPropagation(); + }} + className={cn( + 'focus-visible:ring-text-primary relative mr-1 flex size-4 shrink-0 items-center justify-center rounded-sm outline-hidden focus-visible:ring-2', + )} + > + + {dotClass != null && ( + + )} + + + {summary} + + void refetch()} + /> + + ); +} + +export default memo(PullRequestRowMark); diff --git a/client/src/components/Chat/PullRequest/__tests__/RowMark.test.tsx b/client/src/components/Chat/PullRequest/__tests__/RowMark.test.tsx new file mode 100644 index 00000000000..99d738c18e2 --- /dev/null +++ b/client/src/components/Chat/PullRequest/__tests__/RowMark.test.tsx @@ -0,0 +1,438 @@ +import React from 'react'; +import '@testing-library/jest-dom'; +import userEvent from '@testing-library/user-event'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { render, screen, waitFor, fireEvent } from '@testing-library/react'; +import type { TConversationPullRequest } from 'librechat-data-provider'; +import PullRequestRowMark from '../RowMark'; + +const mockGetMany = jest.fn(); +const mockGetOne = jest.fn(); +const mockStartup: { + current: { pullRequestsEnabled?: boolean; pullRequestsBatchVersion?: number } | undefined; +} = { + current: { pullRequestsEnabled: true, pullRequestsBatchVersion: 1 }, +}; + +const mockUser: { current: { id: string } | undefined } = { current: { id: 'user-a' } }; +jest.mock('recoil', () => ({ + ...jest.requireActual('recoil'), + useRecoilValue: () => mockUser.current, +})); +jest.mock('~/data-provider/Endpoints', () => ({ + useGetStartupConfig: () => ({ data: mockStartup.current }), +})); +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + getConversationPullRequests: (...args: unknown[]) => mockGetMany(...args), + getConversationPullRequest: (...args: unknown[]) => mockGetOne(...args), + }, + }; +}); +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string, values?: Record) => + values == null ? key : `${key}:${Object.values(values).join(',')}`, +})); + +const pr: TConversationPullRequest = { + number: 1234, + title: 'Simplify Single Tool Execution Path', + url: 'https://github.com/LibreChat-AI/LibreChat/pull/1234', + additions: 1234, + deletions: 56, + state: 'open', + isDraft: false, + mergeable: 'clean', + checks: 'passing', +}; + +const answer = (conversationId: string, value: TConversationPullRequest | null) => ({ + results: [{ conversationId, pullRequest: value }], +}); + +let pointerX = 100; +const movePointerOver = (element: HTMLElement) => { + pointerX += 7; + fireEvent.mouseMove(element, { + screenX: pointerX, + screenY: pointerX, + movementX: 7, + movementY: 7, + }); +}; + +const renderMark = (props: Partial> = {}) => { + const rowClick = jest.fn(); + const rowKey = jest.fn(); + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const view = render( + + + + + , + ); + return { ...view, rowClick, rowKey, client }; +}; + +describe('PullRequestRowMark', () => { + beforeEach(() => { + mockGetMany.mockReset(); + mockGetOne.mockReset(); + mockStartup.current = { pullRequestsEnabled: true, pullRequestsBatchVersion: 1 }; + mockUser.current = { id: 'user-a' }; + }); + + it.each([ + [ + 'while loading', + () => new Promise((resolve) => setTimeout(() => resolve(answer('convo-1', null)), 40)), + ], + ['without a pull request', () => Promise.resolve(answer('convo-1', null))], + ])('renders nothing %s, so the row is unchanged', async (_label, respond) => { + mockGetMany.mockImplementation(respond); + const { container } = renderMark(); + await waitFor(() => expect(mockGetMany).toHaveBeenCalled()); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(container.querySelector('[data-testid="convo-pull-request"]')).toBeNull(); + expect(container.querySelector('[data-testid="convo-pull-request-failed"]')).toBeNull(); + }); + + it.each([ + ['the request fails', () => Promise.reject(new Error('503'))], + [ + 'that conversation failed on the server', + () => + Promise.resolve({ + results: [{ conversationId: 'convo-1', error: { code: 'RATE_LIMITED' } }], + }), + ], + ])('says the lookup failed, with a retry, when %s', async (_label, respond) => { + mockGetMany.mockImplementationOnce(respond); + renderMark(); + const failed = await screen.findByTestId('convo-pull-request-failed'); + expect(failed).toHaveAccessibleName('com_ui_pr_load_failed'); + expect(screen.queryByTestId('convo-pull-request')).toBeNull(); + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + fireEvent.focus(failed); + const retry = await screen.findByRole('button', { name: 'com_ui_retry' }); + fireEvent.click(retry); + expect(await screen.findByTestId('convo-pull-request')).toBeInTheDocument(); + expect(mockGetMany).toHaveBeenCalledTimes(2); + }); + + it.each([ + ['an older server that has the flag but not the batch route', { pullRequestsEnabled: true }], + [ + 'a batch version this client does not know', + { pullRequestsEnabled: true, pullRequestsBatchVersion: 2 }, + ], + ['the version without the flag', { pullRequestsBatchVersion: 1 }], + ])('does not ask %s, so a rolling upgrade shows no failure marks', async (_label, startup) => { + mockStartup.current = startup; + const { container } = renderMark(); + await new Promise((resolve) => setTimeout(resolve, 80)); + expect(mockGetMany).not.toHaveBeenCalled(); + expect(container.querySelector('[data-testid="convo-pull-request"]')).toBeNull(); + expect(container.querySelector('[data-testid="convo-pull-request-failed"]')).toBeNull(); + }); + + it('shows the failure when a refresh fails after an empty answer was cached', async () => { + jest.useFakeTimers(); + try { + mockGetMany + .mockResolvedValueOnce(answer('convo-1', null)) + .mockRejectedValue(new Error('503')); + renderMark(); + await jest.advanceTimersByTimeAsync(100); + expect(screen.queryByTestId('convo-pull-request-failed')).toBeNull(); + await jest.advanceTimersByTimeAsync(61_000); + window.dispatchEvent(new Event('focus')); + await jest.advanceTimersByTimeAsync(100); + expect(screen.getByTestId('convo-pull-request-failed')).toBeInTheDocument(); + } finally { + jest.useRealTimers(); + } + }); + + it('keeps its spacing on the rendered mark, so an absent mark takes no room', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + const first = renderMark(); + expect(await screen.findByTestId('convo-pull-request')).toHaveClass('mr-1'); + first.unmount(); + mockGetMany.mockRejectedValueOnce(new Error('503')); + renderMark(); + expect(await screen.findByTestId('convo-pull-request-failed')).toHaveClass('mr-1'); + }); + + it('falls back to the single route when the replica that answered has no batch route', async () => { + mockGetMany.mockRejectedValue( + Object.assign(new Error('Not Found'), { response: { status: 404 } }), + ); + mockGetOne.mockResolvedValue({ pullRequest: pr }); + renderMark(); + expect(await screen.findByTestId('convo-pull-request')).toBeInTheDocument(); + expect(mockGetOne).toHaveBeenCalledWith('convo-1', { signal: expect.any(AbortSignal) }); + expect(screen.queryByTestId('convo-pull-request-failed')).toBeNull(); + }); + + it('shows no failure for a conversation without a pull request when only the batch route is missing', async () => { + mockGetMany.mockRejectedValue( + Object.assign(new Error('Not Found'), { response: { status: 404 } }), + ); + mockGetOne.mockResolvedValue({ pullRequest: null }); + const { container } = renderMark(); + await waitFor(() => expect(mockGetOne).toHaveBeenCalled()); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(container.querySelector('[data-testid="convo-pull-request-failed"]')).toBeNull(); + expect(container.querySelector('[data-testid="convo-pull-request"]')).toBeNull(); + }); + + it('does not ask for a user it cannot name, so a batch is never sent for no one', async () => { + mockUser.current = undefined; + const { container } = renderMark(); + await new Promise((resolve) => setTimeout(resolve, 80)); + expect(mockGetMany).not.toHaveBeenCalled(); + expect(container.querySelector('[data-testid="convo-pull-request"]')).toBeNull(); + }); + + it('asks again for the next account, in a queue of its own, once the user changes', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + const first = renderMark({ conversationId: 'convo-1' }); + await screen.findByTestId('convo-pull-request'); + first.unmount(); + + /** The next account's row is not held up by anything the previous account left queued. */ + mockUser.current = { id: 'user-b' }; + mockGetMany.mockResolvedValue(answer('convo-2', { ...pr, number: 2 })); + renderMark({ conversationId: 'convo-2' }); + await screen.findByTestId('convo-pull-request'); + expect(mockGetMany).toHaveBeenLastCalledWith(['convo-2'], { signal: expect.any(AbortSignal) }); + }); + + it('does not hold the next account behind work the previous account left in flight', async () => { + /** The first account's request never answers, as a stalled upstream would leave it. */ + mockGetMany.mockImplementationOnce(() => new Promise(() => undefined)); + const stalled = renderMark({ conversationId: 'convo-1' }); + await waitFor(() => expect(mockGetMany).toHaveBeenCalledTimes(1)); + stalled.unmount(); + + mockUser.current = { id: 'user-b' }; + mockGetMany.mockResolvedValue(answer('convo-2', { ...pr, number: 2 })); + renderMark({ conversationId: 'convo-2' }); + /** Without a queue of its own it would wait out the other account's request, which never ends. */ + expect(await screen.findByTestId('convo-pull-request')).toBeInTheDocument(); + expect(mockGetMany).toHaveBeenLastCalledWith(['convo-2'], { signal: expect.any(AbortSignal) }); + }); + + it('does not ask when the deployment does not advertise the feature', async () => { + mockStartup.current = { pullRequestsEnabled: false }; + const { container } = renderMark(); + await new Promise((resolve) => setTimeout(resolve, 80)); + expect(mockGetMany).not.toHaveBeenCalled(); + expect(container.querySelector('[data-testid="convo-pull-request"]')).toBeNull(); + }); + + it('asks through the batch endpoint with the conversation id', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark(); + await screen.findByTestId('convo-pull-request'); + expect(mockGetMany).toHaveBeenCalledWith(['convo-1'], { signal: expect.any(AbortSignal) }); + }); + + it('shows the state icon with the CI dot, in the row surface colors', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark({ selected: false }); + const dot = await screen.findByTestId('pull-request-ci-dot'); + expect(dot).toHaveClass('bg-status-success', 'ring-surface-primary-alt'); + expect(screen.getByTestId('convo-pull-request').querySelector('svg')).toHaveClass( + 'text-status-success', + ); + }); + + it('paints the dot ring in the selected surface for the open conversation', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark({ selected: true }); + expect(await screen.findByTestId('pull-request-ci-dot')).toHaveClass( + 'ring-surface-nav-selected', + ); + }); + + it.each([ + ['failing', { checks: 'failing' as const }, 'bg-status-error', 'text-status-success'], + ['running', { checks: 'running' as const }, 'bg-status-warning', 'text-status-success'], + ['conflicts', { mergeable: 'conflicting' as const }, 'bg-status-success', 'text-status-error'], + ])('colors the dot and icon for %s', async (_label, patch, dotClass, iconClass) => { + mockGetMany.mockResolvedValue(answer('convo-1', { ...pr, ...patch })); + renderMark(); + expect(await screen.findByTestId('pull-request-ci-dot')).toHaveClass(dotClass); + expect(screen.getByTestId('convo-pull-request').querySelector('svg')).toHaveClass(iconClass); + }); + + it('shows no dot when there are no checks, or the pull request is finished', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', { ...pr, checks: 'none' })); + const first = renderMark(); + await screen.findByTestId('convo-pull-request'); + expect(screen.queryByTestId('pull-request-ci-dot')).not.toBeInTheDocument(); + first.unmount(); + mockGetMany.mockResolvedValue(answer('convo-1', { ...pr, state: 'merged' })); + renderMark(); + await screen.findByTestId('convo-pull-request'); + expect(screen.queryByTestId('pull-request-ci-dot')).not.toBeInTheDocument(); + }); + + it('tells the row when its description exists, and when it stops existing', async () => { + const onDescribed = jest.fn(); + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + const { unmount } = renderMark({ onDescribed }); + await screen.findByTestId('convo-pull-request'); + expect(onDescribed).toHaveBeenLastCalledWith(true); + unmount(); + expect(onDescribed).toHaveBeenLastCalledWith(false); + }); + + it('reports no description while loading or without a pull request', async () => { + const onDescribed = jest.fn(); + mockGetMany.mockResolvedValue(answer('convo-1', null)); + renderMark({ onDescribed }); + await waitFor(() => expect(mockGetMany).toHaveBeenCalled()); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(onDescribed).not.toHaveBeenCalledWith(true); + }); + + it('reports a description for a failed lookup too, since it renders the failure text', async () => { + const onDescribed = jest.fn(); + mockGetMany.mockRejectedValueOnce(new Error('503')); + renderMark({ onDescribed }); + await screen.findByTestId('convo-pull-request-failed'); + expect(onDescribed).toHaveBeenLastCalledWith(true); + }); + + it('writes what the colors say into text the row can announce', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark({ labelId: 'row-pr' }); + await screen.findByTestId('convo-pull-request'); + const text = document.getElementById('row-pr'); + expect(text).toHaveTextContent(pr.title); + expect(text).toHaveTextContent('com_ui_pr_state_open'); + expect(text).toHaveTextContent('com_ui_pr_checks_passing'); + }); + + it('is a button with an accessible name that says what the colors say', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark(); + const mark = await screen.findByTestId('convo-pull-request'); + expect(mark.tagName).toBe('BUTTON'); + expect(mark).toHaveAccessibleName(expect.stringContaining(pr.title)); + expect(mark).toHaveAccessibleName(expect.stringContaining('com_ui_pr_checks_passing')); + expect(mark).toHaveAttribute('aria-expanded', 'false'); + }); + + it('opens the card for keyboard focus, so its GitHub link is reachable', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark(); + const mark = await screen.findByTestId('convo-pull-request'); + await userEvent.tab(); + expect(mark).toHaveFocus(); + expect(await screen.findByTestId('pull-request-card')).toBeInTheDocument(); + expect(mark).toHaveAttribute('aria-expanded', 'true'); + }); + + it('opens the card on click without opening the row, and keeps Enter and Space off the row', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + const { rowClick, rowKey } = renderMark(); + const mark = await screen.findByTestId('convo-pull-request'); + fireEvent.click(mark); + expect(await screen.findByTestId('pull-request-card')).toBeInTheDocument(); + fireEvent.keyDown(mark, { key: 'Enter' }); + fireEvent.keyDown(mark, { key: ' ' }); + expect(rowClick).not.toHaveBeenCalled(); + expect(rowKey).not.toHaveBeenCalled(); + }); + + it('opens the card to the side on hover, and closes it when the pointer leaves', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark(); + const mark = await screen.findByTestId('convo-pull-request'); + expect(screen.queryByTestId('pull-request-card')).not.toBeInTheDocument(); + movePointerOver(mark); + const card = await screen.findByTestId('pull-request-card'); + expect(card).toHaveTextContent(pr.title); + expect(screen.getByTestId('pull-request-github-link')).toHaveAttribute('href', pr.url); + fireEvent.mouseLeave(mark); + fireEvent.mouseMove(document.body, { screenX: 900, screenY: 900, movementX: 7, movementY: 7 }); + await waitFor(() => expect(screen.queryByTestId('pull-request-card')).not.toBeInTheDocument()); + }); + + it('does not open the row when the card is clicked', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + const { rowClick } = renderMark(); + movePointerOver(await screen.findByTestId('convo-pull-request')); + const card = await screen.findByTestId('pull-request-card'); + fireEvent.click(card); + fireEvent.click(screen.getByTestId('pull-request-github-link')); + expect(rowClick).not.toHaveBeenCalled(); + }); + + it('refreshes every row in one request when the window regains focus, once its answer is stale', async () => { + jest.useFakeTimers(); + try { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark(); + await jest.advanceTimersByTimeAsync(100); + expect(mockGetMany).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(61_000); + window.dispatchEvent(new Event('focus')); + await jest.advanceTimersByTimeAsync(100); + expect(mockGetMany).toHaveBeenCalledTimes(2); + } finally { + jest.useRealTimers(); + } + }); + + it('does not refresh on focus while its answer is still fresh', async () => { + jest.useFakeTimers(); + try { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + renderMark(); + await jest.advanceTimersByTimeAsync(100); + window.dispatchEvent(new Event('focus')); + await jest.advanceTimersByTimeAsync(100); + expect(mockGetMany).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } + }); + + it('never polls, so a list of rows does not poll GitHub once per row', async () => { + jest.useFakeTimers(); + try { + mockGetMany.mockResolvedValue(answer('convo-1', { ...pr, checks: 'running' })); + renderMark(); + await jest.advanceTimersByTimeAsync(100); + expect(mockGetMany).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(10 * 60_000); + expect(mockGetMany).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } + }); + + it('shares its cache entry with the header, so one conversation is fetched once', async () => { + mockGetMany.mockResolvedValue(answer('convo-1', pr)); + const { client } = renderMark(); + await screen.findByTestId('convo-pull-request'); + expect(client.getQueryData(['conversationPullRequest', 'convo-1'])).toEqual({ + pullRequest: pr, + }); + }); +}); diff --git a/client/src/components/Chat/PullRequest/index.ts b/client/src/components/Chat/PullRequest/index.ts index 505e79db72d..b87d27e1063 100644 --- a/client/src/components/Chat/PullRequest/index.ts +++ b/client/src/components/Chat/PullRequest/index.ts @@ -1,2 +1,3 @@ export { default as PullRequestChip } from './Chip'; export { default as usePullRequestMenu } from './menu'; +export { default as PullRequestRowMark } from './RowMark'; diff --git a/client/src/components/Conversations/Convo.tsx b/client/src/components/Conversations/Convo.tsx index 197075123eb..02767489190 100644 --- a/client/src/components/Conversations/Convo.tsx +++ b/client/src/components/Conversations/Convo.tsx @@ -4,7 +4,11 @@ import { Link2 } from 'lucide-react'; import { useRecoilValue } from 'recoil'; import { useParams } from 'react-router-dom'; import { Spinner, useToastContext, useMediaQuery } from '@librechat/client'; -import { Constants, supportsConversationTitleOwnership } from 'librechat-data-provider'; +import { + Constants, + PULL_REQUEST_BATCH_VERSION, + supportsConversationTitleOwnership, +} from 'librechat-data-provider'; import type { TConversation } from 'librechat-data-provider'; import type { ConversationDragItem } from './dnd'; import { @@ -13,6 +17,7 @@ import { useUpdateConversationMutation, } from '~/data-provider'; import { cn, logger, setDocumentTitle, isConversationUnseen, hasRealTitle } from '~/utils'; +import PullRequestRowMark from '~/components/Chat/PullRequest/RowMark'; import { useNavigateToConvo, useLocalize, useShiftKey } from '~/hooks'; import ConversationEndpointIcon from './ConversationEndpointIcon'; import useDrawerViewport from '~/hooks/Nav/useDrawerViewport'; @@ -70,6 +75,13 @@ function Conversation({ const sharedLinksEnabled = startupConfig?.sharedLinksEnabled === true; const isSharedBadgeVisible = conversation.isShared === true && sharedLinksEnabled; const projectLabelId = useId(); + const pullRequestLabelId = useId(); + /** The mark owns the description text; the row points at it only while it is in the page. */ + const [pullRequestDescribed, setPullRequestDescribed] = useState(false); + const showPullRequest = + startupConfig?.pullRequestsEnabled === true && + startupConfig.pullRequestsBatchVersion === PULL_REQUEST_BATCH_VERSION && + !isGenerating; const projectBadgeProjectId = showProjectBadge ? conversation.chatProjectId : undefined; const isUnseen = isConversationUnseen(conversation); const isShiftHeld = useShiftKey(); @@ -387,7 +399,14 @@ function Conversation({ isSmallScreen={isSmallScreen} localize={localize} keyShortcuts={keyShortcuts} - describedBy={projectBadgeProjectId ? projectLabelId : undefined} + describedBy={ + [ + projectBadgeProjectId ? projectLabelId : null, + showPullRequest && pullRequestDescribed ? pullRequestLabelId : null, + ] + .filter(Boolean) + .join(' ') || undefined + } > {/* Status sits on the avatar so the row's trailing edge stays free for its badges and menu. The ring is 2.125rem around the 1.25rem icon: offset by half the @@ -447,6 +466,19 @@ function Conversation({ {/* Only render ConvoOptions when user interacts (hover/focus) or for active conversation */} {actionContent} + {/* After the action slot on purpose: that slot grows from zero on hover, and anything before + it slides left under the pointer, so the mark would move out from under the cursor that + is trying to reach it. The title is the only thing that gives way. */} + {showPullRequest && ( + + + + )} ); } diff --git a/client/src/components/Conversations/__tests__/Convo.pullRequest.spec.tsx b/client/src/components/Conversations/__tests__/Convo.pullRequest.spec.tsx new file mode 100644 index 00000000000..fde5eb708d7 --- /dev/null +++ b/client/src/components/Conversations/__tests__/Convo.pullRequest.spec.tsx @@ -0,0 +1,176 @@ +import React, { useEffect } from 'react'; +import { DndProvider } from 'react-dnd'; +import { HTML5Backend } from 'react-dnd-html5-backend'; +import { render, screen } from '@testing-library/react'; +import type { TConversation } from 'librechat-data-provider'; + +const mockStartup: { + current: { pullRequestsEnabled?: boolean; pullRequestsBatchVersion?: number }; +} = { + current: { pullRequestsEnabled: true, pullRequestsBatchVersion: 1 }, +}; +const mockMarkProps: Array> = []; +/** What the mock mark reports to the row, as the real one does once it has text to show. */ +const mockMarkState = { described: true }; + +jest.mock('@librechat/client', () => ({ + useMediaQuery: () => false, + useRemScale: () => 1, + useToastContext: () => ({ showToast: jest.fn() }), + Spinner: () => , + Button: ({ children, ...props }: React.ComponentProps<'button'>) => ( + {children} + ), +})); +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string) => key, + useNavigateToConvo: () => ({ navigateToConvo: jest.fn() }), + useShiftKey: () => false, +})); +jest.mock('~/data-provider', () => ({ + useGetStartupConfig: () => ({ data: mockStartup.current }), + useUpdateConversationMutation: () => ({ mutateAsync: jest.fn() }), + usePinConversationMutation: () => ({ mutate: jest.fn() }), + useProjectName: () => 'Scheduling', +})); +jest.mock('react-router-dom', () => ({ useParams: () => ({ conversationId: 'convo-1' }) })); +jest.mock('recoil', () => ({ useRecoilValue: () => [] })); +jest.mock('~/store', () => ({ + __esModule: true, + default: { + allConversationsSelector: 'allConversationsSelector', + conversationIdByIndex: () => 'conversationIdByIndex', + }, +})); +jest.mock('~/utils', () => ({ + cn: (...classes: unknown[]) => classes.filter(Boolean).join(' '), + logger: { error: jest.fn() }, + setDocumentTitle: jest.fn(), + isConversationUnseen: () => false, + hasRealTitle: (title: string) => !!title && title !== 'New Chat', +})); +jest.mock('../ConvoOptions', () => ({ ConvoOptions: () => })); +jest.mock('../ConversationEndpointIcon', () => ({ + __esModule: true, + default: () => , +})); +jest.mock('../RenameForm', () => ({ __esModule: true, default: () => })); +/** Named with the `mock` prefix so the factory below may reference it, and a real component so + * its hook is legal. It reports its description to the row the way the real mark does. */ +function MockRowMark({ + labelId, + onDescribed, + ...rest +}: Record & { labelId: string; onDescribed?: (described: boolean) => void }) { + mockMarkProps.push({ labelId, ...rest }); + useEffect(() => { + onDescribed?.(mockMarkState.described); + return () => onDescribed?.(false); + }, [onDescribed]); + if (!mockMarkState.described) return null; + return ( + + + + ); +} +const mockRowMark = MockRowMark; + +jest.mock('~/components/Chat/PullRequest/RowMark', () => ({ + __esModule: true, + default: (props: Record) => mockRowMark(props as never), +})); + +import Conversation from '../Convo'; + +const conversation = { + conversationId: 'convo-1', + title: 'Tool Approval UI', +} as TConversation; + +const renderRow = (props: { isGenerating?: boolean } = {}) => + render( + + + , + ); + +const rowButton = () => screen.getByRole('button', { name: /^com_ui_conversation_label/ }); + +describe('Conversation row pull request', () => { + beforeEach(() => { + mockMarkProps.length = 0; + mockMarkState.described = true; + mockStartup.current = { pullRequestsEnabled: true, pullRequestsBatchVersion: 1 }; + }); + + it("keeps the title the conversation's own and puts the pull request beside it, outside the row button", () => { + renderRow(); + expect(screen.getByText('Tool Approval UI')).toBeInTheDocument(); + const mark = screen.getByTestId('convo-pull-request'); + expect(rowButton()).not.toContainElement(mark); + expect(screen.getByTestId('convo-item')).toContainElement(mark); + expect(mark.parentElement?.parentElement).toBe(screen.getByTestId('convo-item')); + }); + + it('sits after the row menu, which grows on hover, so it does not slide out from under the pointer', () => { + renderRow(); + const row = screen.getByTestId('convo-item'); + const mark = screen.getByTestId('convo-pull-request'); + const menuSlot = screen.getByTestId('convo-options').parentElement as HTMLElement; + expect(row).toContainElement(menuSlot); + expect(menuSlot.compareDocumentPosition(mark) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(row.lastElementChild).toBe(mark.parentElement); + }); + + it("names the pull request in the row's description so a screen reader hears what the colors say", () => { + renderRow(); + const labelId = String(mockMarkProps[0].labelId); + expect(rowButton().getAttribute('aria-describedby')).toContain(labelId); + expect(document.getElementById(labelId)).toBeInTheDocument(); + }); + + it('tells the mark which conversation it is for and that this row is the open one', () => { + renderRow(); + expect(mockMarkProps[0]).toMatchObject({ conversationId: 'convo-1', selected: true }); + }); + + it('reserves no space around the mark, since an absent mark must leave the row unchanged', () => { + renderRow(); + const wrapper = screen.getByTestId('convo-pull-request').parentElement as HTMLElement; + expect(wrapper.className).not.toMatch(/\bm[rlxy]?-\d/); + expect(wrapper.className).not.toMatch(/\bp[rlxy]?-\d/); + }); + + it('does not point the row at a description that is not in the page', () => { + mockMarkState.described = false; + renderRow(); + expect(screen.queryByTestId('convo-pull-request')).not.toBeInTheDocument(); + expect(rowButton().getAttribute('aria-describedby')).toBeNull(); + const ids = (rowButton().getAttribute('aria-describedby') ?? '').split(' ').filter(Boolean); + expect(ids.every((id) => document.getElementById(id) != null)).toBe(true); + }); + + it.each([ + ['not advertised', {}], + ['advertised off', { pullRequestsEnabled: false }], + ['advertised without the batch route', { pullRequestsEnabled: true }], + ])('draws nothing and describes nothing when the feature is %s', (_label, startup) => { + mockStartup.current = startup; + renderRow(); + expect(screen.queryByTestId('convo-pull-request')).not.toBeInTheDocument(); + expect(mockMarkProps).toHaveLength(0); + expect(rowButton().getAttribute('aria-describedby')).toBeNull(); + }); + + it('leaves the mark out while the chat runs, so the ring is the only status', () => { + renderRow({ isGenerating: true }); + expect(screen.queryByTestId('convo-pull-request')).not.toBeInTheDocument(); + expect(screen.getByTestId('status-ring')).toBeInTheDocument(); + }); +}); diff --git a/client/src/data-provider/PullRequest/batch.test.ts b/client/src/data-provider/PullRequest/batch.test.ts new file mode 100644 index 00000000000..8ff6dd9afd8 --- /dev/null +++ b/client/src/data-provider/PullRequest/batch.test.ts @@ -0,0 +1,758 @@ +import type { TConversationPullRequest } from 'librechat-data-provider'; +import { PullRequestBatchError, createBatchFetcher, createPullRequestBatcher } from './batch'; + +const pr: TConversationPullRequest = { + number: 1, + title: 't', + url: 'https://github.com/o/r/pull/1', + additions: 1, + deletions: 0, + state: 'open', + isDraft: false, + mergeable: 'clean', + checks: 'passing', +}; + +describe('createPullRequestBatcher', () => { + beforeEach(() => jest.useFakeTimers()); + afterEach(() => jest.useRealTimers()); + + it('asks once for every conversation requested in the same moment', async () => { + const fetchMany = jest.fn().mockResolvedValue({ + results: [ + { conversationId: 'a', pullRequest: pr }, + { conversationId: 'b', pullRequest: null }, + ], + }); + const { load } = createPullRequestBatcher({ fetchMany }); + const first = load('a'); + const second = load('b'); + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledTimes(1); + expect(fetchMany).toHaveBeenCalledWith(['a', 'b'], expect.any(AbortSignal)); + await expect(first).resolves.toEqual({ pullRequest: pr }); + await expect(second).resolves.toEqual({ pullRequest: null }); + }); + + it('shares one entry between callers asking for the same conversation', async () => { + const fetchMany = jest + .fn() + .mockResolvedValue({ results: [{ conversationId: 'a', pullRequest: pr }] }); + const { load } = createPullRequestBatcher({ fetchMany }); + const calls = [load('a'), load('a')]; + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledWith(['a'], expect.any(AbortSignal)); + await expect(Promise.all(calls)).resolves.toEqual([{ pullRequest: pr }, { pullRequest: pr }]); + }); + + it('answers no pull request for a conversation the server left out', async () => { + const fetchMany = jest.fn().mockResolvedValue({ results: [] }); + const { load } = createPullRequestBatcher({ fetchMany }); + const answer = load('a'); + await jest.advanceTimersByTimeAsync(50); + await expect(answer).resolves.toEqual({ pullRequest: null }); + }); + + it('splits a crowd into requests the server accepts', async () => { + const fetchMany = jest.fn(async (ids: string[]) => ({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: null })), + })); + const { load } = createPullRequestBatcher({ fetchMany, maxBatch: 2 }); + const all = ['a', 'b', 'c', 'd', 'e'].map((id) => load(id)); + await jest.advanceTimersByTimeAsync(50); + await Promise.all(all); + expect(fetchMany.mock.calls.map(([ids]) => ids)).toEqual([['a', 'b'], ['c', 'd'], ['e']]); + }); + + describe('requests in flight', () => { + /** A server that answers only when told to, so overlap is visible. */ + const slowServer = () => { + let active = 0; + let peak = 0; + const releases: Array<() => void> = []; + const fetchMany = jest.fn( + (ids: string[]) => + new Promise<{ results: Array<{ conversationId: string; pullRequest: null }> }>( + (resolve) => { + active += 1; + peak = Math.max(peak, active); + releases.push(() => { + active -= 1; + resolve({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: null })), + }); + }); + }, + ), + ); + return { fetchMany, releases, peak: () => peak }; + }; + + it('sends the chunks of a crowd one at a time, so the server limit is not multiplied', async () => { + const { fetchMany, releases, peak } = slowServer(); + const { load } = createPullRequestBatcher({ fetchMany, maxBatch: 2 }); + const all = ['a', 'b', 'c', 'd', 'e'].map((id) => load(id)); + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledTimes(1); + releases.shift()?.(); + await jest.advanceTimersByTimeAsync(0); + expect(fetchMany).toHaveBeenCalledTimes(2); + releases.shift()?.(); + await jest.advanceTimersByTimeAsync(0); + expect(fetchMany).toHaveBeenCalledTimes(3); + releases.shift()?.(); + await Promise.all(all); + expect(peak()).toBe(1); + expect(fetchMany.mock.calls.map(([ids]) => ids)).toEqual([['a', 'b'], ['c', 'd'], ['e']]); + }); + + it('does not start a later flush while an earlier request is still out', async () => { + const { fetchMany, releases, peak } = slowServer(); + const { load } = createPullRequestBatcher({ fetchMany }); + const first = load('a'); + await jest.advanceTimersByTimeAsync(50); + const second = load('b'); + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledTimes(1); + releases.shift()?.(); + await jest.advanceTimersByTimeAsync(0); + expect(fetchMany).toHaveBeenCalledTimes(2); + releases.shift()?.(); + await Promise.all([first, second]); + expect(peak()).toBe(1); + }); + + it('moves on when a request never answers, failing only the callers it carried', async () => { + const fetchMany = jest + .fn() + .mockImplementationOnce(() => new Promise(() => undefined)) + .mockResolvedValue({ results: [{ conversationId: 'b', pullRequest: null }] }); + const { load } = createPullRequestBatcher({ fetchMany, requestTimeoutMs: 1000 }); + const hung = load('a').catch((error) => error); + await jest.advanceTimersByTimeAsync(50); + const later = load('b'); + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1000); + const error = await hung; + expect(error).toBeInstanceOf(PullRequestBatchError); + expect(error.code).toBe('TIMEOUT'); + await jest.advanceTimersByTimeAsync(50); + await expect(later).resolves.toEqual({ pullRequest: null }); + }); + + it('waits longer than the longest server deadline a deployment can configure by default', async () => { + const fetchMany = jest + .fn() + .mockImplementationOnce( + () => + new Promise((resolve) => + setTimeout( + () => resolve({ results: [{ conversationId: 'a', pullRequest: null }] }), + 125_000, + ), + ), + ) + .mockResolvedValue({ results: [{ conversationId: 'b', pullRequest: null }] }); + const { load } = createPullRequestBatcher({ fetchMany }); + const slow = load('a'); + await jest.advanceTimersByTimeAsync(50); + const queued = load('b'); + await jest.advanceTimersByTimeAsync(124_000); + /** A 120 s server batch has not finished at 124 s of waiting, and the next must not have started. */ + expect(fetchMany).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(2_000); + await expect(slow).resolves.toEqual({ pullRequest: null }); + await jest.advanceTimersByTimeAsync(50); + await expect(queued).resolves.toEqual({ pullRequest: null }); + expect(fetchMany).toHaveBeenCalledTimes(2); + }); + + it('does not time out a request that answers in time', async () => { + const fetchMany = jest + .fn() + .mockResolvedValue({ results: [{ conversationId: 'a', pullRequest: null }] }); + const { load } = createPullRequestBatcher({ fetchMany, requestTimeoutMs: 1000 }); + const answer = load('a'); + await jest.advanceTimersByTimeAsync(50); + await expect(answer).resolves.toEqual({ pullRequest: null }); + await jest.advanceTimersByTimeAsync(5000); + }); + + it.each([ + ['no body', undefined], + ['a body without results', {}], + ['results that are not a list', { results: 'nope' }], + ['a result that is not an entry', { results: [null] }], + ])( + 'fails its callers, and not the queue, when the server answers with %s', + async (_label, body) => { + const fetchMany = jest + .fn() + .mockResolvedValueOnce(body) + .mockResolvedValue({ results: [{ conversationId: 'b', pullRequest: null }] }); + const { load } = createPullRequestBatcher({ fetchMany, maxBatch: 1 }); + const bad = load('a').then( + () => 'resolved', + () => 'rejected', + ); + const good = load('b'); + await jest.advanceTimersByTimeAsync(50); + await expect(bad).resolves.toBe('rejected'); + await expect(good).resolves.toEqual({ pullRequest: null }); + expect(fetchMany).toHaveBeenCalledTimes(2); + }, + ); + + it('keeps sending after a request fails', async () => { + const fetchMany = jest + .fn() + .mockRejectedValueOnce(new Error('503')) + .mockResolvedValue({ results: [{ conversationId: 'c', pullRequest: null }] }); + const { load } = createPullRequestBatcher({ fetchMany, maxBatch: 1 }); + const failed = load('a').catch((error) => error.message); + const fine = load('c'); + await jest.advanceTimersByTimeAsync(50); + await expect(failed).resolves.toBe('503'); + await expect(fine).resolves.toEqual({ pullRequest: null }); + }); + }); + + it('rejects only the conversation whose entry failed, carrying its code', async () => { + const fetchMany = jest.fn().mockResolvedValue({ + results: [ + { conversationId: 'a', error: { code: 'RATE_LIMITED' } }, + { conversationId: 'b', pullRequest: pr }, + ], + }); + const { load } = createPullRequestBatcher({ fetchMany }); + const failed = load('a'); + const fine = load('b'); + const outcome = failed.catch((error) => error); + await jest.advanceTimersByTimeAsync(50); + const error = await outcome; + expect(error).toBeInstanceOf(PullRequestBatchError); + expect(error.code).toBe('RATE_LIMITED'); + await expect(fine).resolves.toEqual({ pullRequest: pr }); + }); + + it('rejects every caller of a request that fails outright', async () => { + const fetchMany = jest.fn().mockRejectedValue(new Error('503')); + const { load } = createPullRequestBatcher({ fetchMany }); + const outcomes = [load('a'), load('b')].map((promise) => + promise.catch((error) => error.message), + ); + await jest.advanceTimersByTimeAsync(50); + await expect(Promise.all(outcomes)).resolves.toEqual(['503', '503']); + }); + + it('starts a new request for a conversation asked for after the last went out', async () => { + const fetchMany = jest.fn(async (ids: string[]) => ({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: null })), + })); + const { load } = createPullRequestBatcher({ fetchMany }); + const first = load('a'); + await jest.advanceTimersByTimeAsync(50); + await first; + const second = load('b'); + await jest.advanceTimersByTimeAsync(50); + await second; + expect(fetchMany.mock.calls.map(([ids]) => ids)).toEqual([['a'], ['b']]); + }); +}); + +describe('createPullRequestBatcher dispose', () => { + beforeEach(() => jest.useFakeTimers()); + afterEach(() => jest.useRealTimers()); + + it('rejects ids that have not gone out, and sends nothing for them', async () => { + const fetchMany = jest.fn(); + const batcher = createPullRequestBatcher({ fetchMany }); + const waiting = batcher.load('a').catch((error) => error); + batcher.dispose(); + await jest.advanceTimersByTimeAsync(100); + const error = await waiting; + expect(error).toBeInstanceOf(PullRequestBatchError); + expect(error.code).toBe('DISPOSED'); + expect(fetchMany).not.toHaveBeenCalled(); + }); + + it('does not send chunks that were still queued behind a request in flight', async () => { + const releases: Array<() => void> = []; + const fetchMany = jest.fn( + (ids: string[]) => + new Promise<{ results: Array<{ conversationId: string; pullRequest: null }> }>((resolve) => + releases.push(() => + resolve({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: null })), + }), + ), + ), + ); + const batcher = createPullRequestBatcher({ fetchMany, maxBatch: 1 }); + const outcomes = ['a', 'b', 'c'].map((id) => + batcher.load(id).then( + () => 'sent', + (error) => error.code, + ), + ); + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledTimes(1); + batcher.dispose(); + releases.shift()?.(); + await jest.advanceTimersByTimeAsync(50); + /** The request on the wire is aborted with the batcher, and the queued chunks never go out. */ + expect(await Promise.all(outcomes)).toEqual(['ABORTED', 'DISPOSED', 'DISPOSED']); + expect(fetchMany).toHaveBeenCalledTimes(1); + }); + + it('refuses new work once disposed', async () => { + const batcher = createPullRequestBatcher({ fetchMany: jest.fn() }); + batcher.dispose(); + await expect(batcher.load('a')).rejects.toMatchObject({ code: 'DISPOSED' }); + }); +}); + +describe('createBatchFetcher', () => { + const entry = (conversationId: string) => ({ conversationId, pullRequest: null }); + const notFound = () => Object.assign(new Error('Not Found'), { response: { status: 404 } }); + + it('uses the batch route when the server has it', async () => { + const fetchMany = jest.fn().mockResolvedValue({ results: [entry('a')] }); + const fetchOne = jest.fn(); + await expect( + createBatchFetcher({ fetchMany, fetchOne })(['a'], new AbortController().signal), + ).resolves.toEqual({ + results: [entry('a')], + }); + expect(fetchOne).not.toHaveBeenCalled(); + }); + + it('asks the single route for each conversation when the replica has no batch route', async () => { + const fetchMany = jest.fn().mockRejectedValue(notFound()); + const fetchOne = jest.fn(async (id: string) => ({ pullRequest: id === 'a' ? pr : null })); + const results = await createBatchFetcher({ fetchMany, fetchOne })( + ['a', 'b'], + new AbortController().signal, + ); + expect(results).toEqual({ + results: [ + { conversationId: 'a', pullRequest: pr }, + { conversationId: 'b', pullRequest: null }, + ], + }); + expect(fetchOne.mock.calls.map(([id]) => id)).toEqual(['a', 'b']); + }); + + it("keeps one conversation's failure to itself and reports its code", async () => { + const fetchMany = jest.fn().mockRejectedValue(notFound()); + const fetchOne = jest.fn(async (id: string) => { + if (id === 'a') { + throw Object.assign(new Error('x'), { + response: { status: 503, data: { code: 'RATE_LIMITED' } }, + }); + } + return { pullRequest: pr }; + }); + const { results } = await createBatchFetcher({ fetchMany, fetchOne })( + ['a', 'b'], + new AbortController().signal, + ); + expect(results).toEqual([ + { conversationId: 'a', error: { code: 'RATE_LIMITED' } }, + { conversationId: 'b', pullRequest: pr }, + ]); + }); + + it('answers a failure with no code as an upstream error, never with its text', async () => { + const fetchMany = jest.fn().mockRejectedValue(notFound()); + const fetchOne = jest.fn().mockRejectedValue(new Error('secret host name')); + const { results } = await createBatchFetcher({ fetchMany, fetchOne })( + ['a'], + new AbortController().signal, + ); + expect(results).toEqual([{ conversationId: 'a', error: { code: 'UPSTREAM_ERROR' } }]); + expect(JSON.stringify(results)).not.toContain('secret'); + }); + + it.each([ + ['a server error', { response: { status: 503 } }], + ['a rate limit', { response: { status: 429 } }], + ['an expired session', { response: { status: 401 } }], + ['a network failure', new Error('Network Error')], + ])( + 'does not fall back on %s, since the route exists and really failed', + async (_label, failure) => { + const error = failure instanceof Error ? failure : Object.assign(new Error('x'), failure); + const fetchMany = jest.fn().mockRejectedValue(error); + const fetchOne = jest.fn(); + await expect( + createBatchFetcher({ fetchMany, fetchOne })(['a'], new AbortController().signal), + ).rejects.toBe(error); + expect(fetchOne).not.toHaveBeenCalled(); + }, + ); + + const measuredFallback = async (concurrency?: () => number | undefined) => { + let active = 0; + let peak = 0; + const fetchMany = jest.fn().mockRejectedValue(notFound()); + const fetchOne = jest.fn(async () => { + active += 1; + peak = Math.max(peak, active); + await new Promise((resolve) => setTimeout(resolve, 5)); + active -= 1; + return { pullRequest: null }; + }); + jest.useRealTimers(); + const ids = Array.from({ length: 12 }, (_, i) => `c${i}`); + const { results } = await createBatchFetcher({ fetchMany, fetchOne, concurrency })( + ids, + new AbortController().signal, + ); + expect(results.map((r) => r.conversationId)).toEqual(ids); + return peak; + }; + + it('runs the single-route calls one at a time when the server never said how many it allows', async () => { + expect(await measuredFallback()).toBe(1); + expect(await measuredFallback(() => undefined)).toBe(1); + }); + + it('runs the single-route calls at the limit the server advertised, no more', async () => { + expect(await measuredFallback(() => 3)).toBe(3); + expect(await measuredFallback(() => 1)).toBe(1); + }); + + it.each([0, -2, 1.5, Number.NaN])( + 'treats an unusable advertised limit of %s as one', + async (limit) => { + expect(await measuredFallback(() => limit)).toBe(1); + }, + ); + + it('reads the advertised limit when the fallback runs, not when the fetcher was made', async () => { + let limit = 1; + const fetchMany = jest.fn().mockRejectedValue(notFound()); + let active = 0; + let peak = 0; + const fetchOne = jest.fn(async () => { + active += 1; + peak = Math.max(peak, active); + await new Promise((resolve) => setTimeout(resolve, 5)); + active -= 1; + return { pullRequest: null }; + }); + const fetcher = createBatchFetcher({ fetchMany, fetchOne, concurrency: () => limit }); + /** The startup config arrives after the fetcher exists, as it does for the module-level one. */ + limit = 3; + jest.useRealTimers(); + const ids = Array.from({ length: 9 }, (_, i) => `c${i}`); + await fetcher(ids, new AbortController().signal); + expect(peak).toBe(3); + }); +}); + +describe('the single-route fallback and the queue', () => { + const notFound = () => Object.assign(new Error('Not Found'), { response: { status: 404 } }); + + /** A single-route call that ends only when released or aborted, like a stalled GET. */ + const stalledSingleRoute = () => { + const calls: string[] = []; + const aborted: string[] = []; + const fetchOne = jest.fn( + (id: string, signal: AbortSignal) => + new Promise<{ pullRequest: null }>((_, reject) => { + calls.push(id); + signal.addEventListener('abort', () => { + aborted.push(id); + reject(new Error('aborted')); + }); + }), + ); + return { fetchOne, calls, aborted }; + }; + + it('stops its calls when the signal aborts, and starts no more', async () => { + const { fetchOne, calls, aborted } = stalledSingleRoute(); + const controller = new AbortController(); + const ids = Array.from({ length: 12 }, (_, i) => `c${i}`); + const fetcher = createBatchFetcher({ + fetchMany: jest.fn().mockRejectedValue(notFound()), + fetchOne, + concurrency: () => 4, + }); + const settled = fetcher(ids, controller.signal).catch((error) => error); + await Promise.resolve(); + await new Promise((resolve) => setTimeout(resolve, 5)); + expect(calls).toHaveLength(4); + controller.abort(); + const error = await settled; + expect(error).toBeInstanceOf(PullRequestBatchError); + expect(error.code).toBe('ABORTED'); + expect(aborted).toHaveLength(4); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(calls).toHaveLength(4); + }); + + it('starts nothing when the signal is already aborted', async () => { + const controller = new AbortController(); + controller.abort(); + const fetchOne = jest.fn(); + const fetcher = createBatchFetcher({ + fetchMany: jest.fn().mockRejectedValue(notFound()), + fetchOne, + }); + await expect(fetcher(['a', 'b'], controller.signal)).rejects.toMatchObject({ code: 'ABORTED' }); + expect(fetchOne).not.toHaveBeenCalled(); + }); + + it("does not report an aborted call as one conversation's failure", async () => { + const controller = new AbortController(); + const fetchOne = jest.fn( + (_id: string, signal: AbortSignal) => + new Promise((_, reject) => + signal.addEventListener('abort', () => reject(new Error('x'))), + ), + ); + const fetcher = createBatchFetcher({ + fetchMany: jest.fn().mockRejectedValue(notFound()), + fetchOne, + }); + const settled = fetcher(['a'], controller.signal).catch((error) => error); + await new Promise((resolve) => setTimeout(resolve, 5)); + controller.abort(); + expect(await settled).toMatchObject({ code: 'ABORTED' }); + }); + + it('aborts a request that runs past its timeout, so its fallback stops before the next request starts', async () => { + jest.useFakeTimers(); + try { + const { fetchOne, calls, aborted } = stalledSingleRoute(); + const fetchMany = jest + .fn() + .mockRejectedValueOnce(notFound()) + .mockResolvedValue({ results: [{ conversationId: 'z', pullRequest: null }] }); + const fetcher = createBatchFetcher({ fetchMany, fetchOne }); + const batcher = createPullRequestBatcher({ + fetchMany: fetcher, + maxBatch: 1, + requestTimeoutMs: 1000, + }); + const first = batcher.load('a').catch((error) => error.code); + const second = batcher.load('z'); + await jest.advanceTimersByTimeAsync(50); + expect(calls).toEqual(['a']); + expect(fetchMany).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1000); + expect(await first).toBe('TIMEOUT'); + /** The stalled call was aborted before the queue handed the next request its turn. */ + expect(aborted).toEqual(['a']); + await jest.advanceTimersByTimeAsync(50); + await expect(second).resolves.toEqual({ pullRequest: null }); + expect(fetchMany).toHaveBeenCalledTimes(2); + } finally { + jest.useRealTimers(); + } + }); + + it('aborts the request on the wire when the batcher is disposed', async () => { + const { fetchOne, calls, aborted } = stalledSingleRoute(); + const fetcher = createBatchFetcher({ + fetchMany: jest.fn().mockRejectedValue(notFound()), + fetchOne, + }); + const batcher = createPullRequestBatcher({ fetchMany: fetcher, delayMs: 1 }); + const outcome = batcher.load('a').catch((error) => error.code); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(calls).toEqual(['a']); + batcher.dispose(); + expect(aborted).toEqual(['a']); + expect(await outcome).toBe('ABORTED'); + }); + + it("hands the batcher's own signal to the fetcher, live while the request runs", async () => { + let seen: AbortSignal | undefined; + const fetchMany = jest.fn(async (_ids: string[], signal: AbortSignal) => { + seen = signal; + return { results: [{ conversationId: 'a', pullRequest: null }] }; + }); + const batcher = createPullRequestBatcher({ fetchMany, delayMs: 1 }); + await batcher.load('a'); + expect(seen).toBeDefined(); + expect(seen?.aborted).toBe(false); + }); +}); + +describe('the batch POST receives the request signal', () => { + it("hands the batcher's signal to the primary batch call, and aborts it on timeout", async () => { + jest.useFakeTimers(); + try { + const seen: AbortSignal[] = []; + const fetchMany = jest.fn( + (_ids: string[], signal: AbortSignal) => + new Promise(() => { + seen.push(signal); + }), + ); + const fetcher = createBatchFetcher({ fetchMany, fetchOne: jest.fn() }); + const batcher = createPullRequestBatcher({ fetchMany: fetcher, requestTimeoutMs: 1000 }); + const outcome = batcher.load('a').catch((error) => error.code); + await jest.advanceTimersByTimeAsync(50); + expect(seen).toHaveLength(1); + expect(seen[0].aborted).toBe(false); + await jest.advanceTimersByTimeAsync(1000); + expect(await outcome).toBe('TIMEOUT'); + expect(seen[0].aborted).toBe(true); + } finally { + jest.useRealTimers(); + } + }); + + it('aborts the batch call on the wire when the batcher is disposed', async () => { + jest.useFakeTimers(); + try { + let seen: AbortSignal | undefined; + const fetchMany = jest.fn( + (_ids: string[], signal: AbortSignal) => + new Promise(() => { + seen = signal; + }), + ); + const batcher = createPullRequestBatcher({ + fetchMany: createBatchFetcher({ fetchMany, fetchOne: jest.fn() }), + }); + const outcome = batcher.load('a').catch(() => 'rejected'); + await jest.advanceTimersByTimeAsync(50); + batcher.dispose(); + expect(seen?.aborted).toBe(true); + expect(await outcome).toBe('rejected'); + } finally { + jest.useRealTimers(); + } + }); +}); + +describe('withdrawing a caller that stopped waiting', () => { + beforeEach(() => jest.useFakeTimers()); + afterEach(() => jest.useRealTimers()); + + const answer = (ids: string[]) => ({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: null })), + }); + + it('rejects the caller at once', async () => { + const fetchMany = jest.fn(async (ids: string[]) => answer(ids)); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + const waiting = load('a', controller.signal).catch((error) => error); + controller.abort(); + const error = await waiting; + expect(error).toBeInstanceOf(PullRequestBatchError); + expect(error.code).toBe('ABORTED'); + }); + + it('never asks for an id whose only caller left before the request went out', async () => { + const fetchMany = jest.fn(async (ids: string[]) => answer(ids)); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + void load('gone', controller.signal).catch(() => undefined); + const kept = load('kept'); + controller.abort(); + await jest.advanceTimersByTimeAsync(50); + await kept; + expect(fetchMany).toHaveBeenCalledTimes(1); + expect(fetchMany.mock.calls[0][0]).toEqual(['kept']); + }); + + it('sends nothing at all when every caller left', async () => { + const fetchMany = jest.fn(async (ids: string[]) => answer(ids)); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + void load('a', controller.signal).catch(() => undefined); + void load('b', controller.signal).catch(() => undefined); + controller.abort(); + await jest.advanceTimersByTimeAsync(100); + expect(fetchMany).not.toHaveBeenCalled(); + }); + + it('keeps an id that another caller still waits for', async () => { + const fetchMany = jest.fn(async (ids: string[]) => answer(ids)); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + void load('shared', controller.signal).catch(() => undefined); + const stayed = load('shared'); + controller.abort(); + await jest.advanceTimersByTimeAsync(50); + await expect(stayed).resolves.toEqual({ pullRequest: null }); + expect(fetchMany.mock.calls[0][0]).toEqual(['shared']); + }); + + it('skips a queued request nobody waits for any more, so the next one is not held behind it', async () => { + const releases: Array<() => void> = []; + const fetchMany = jest.fn( + (ids: string[]) => + new Promise>((resolve) => { + releases.push(() => resolve(answer(ids))); + }), + ); + const { load } = createPullRequestBatcher({ fetchMany, maxBatch: 1 }); + const first = load('first'); + const controller = new AbortController(); + void load('queued', controller.signal).catch(() => undefined); + const last = load('last'); + await jest.advanceTimersByTimeAsync(50); + expect(fetchMany).toHaveBeenCalledTimes(1); + controller.abort(); + releases.shift()?.(); + await first; + await jest.advanceTimersByTimeAsync(0); + expect(fetchMany.mock.calls.map(([ids]) => ids)).toEqual([['first'], ['last']]); + releases.shift()?.(); + await last; + }); + + it('aborts the request on the wire once nobody waits for it', async () => { + let seen: AbortSignal | undefined; + const fetchMany = jest.fn( + (_ids: string[], signal: AbortSignal) => + new Promise(() => { + seen = signal; + }), + ); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + void load('a', controller.signal).catch(() => undefined); + await jest.advanceTimersByTimeAsync(50); + expect(seen?.aborted).toBe(false); + controller.abort(); + expect(seen?.aborted).toBe(true); + }); + + it('leaves the request on the wire running while another caller still waits for it', async () => { + let seen: AbortSignal | undefined; + const fetchMany = jest.fn( + (_ids: string[], signal: AbortSignal) => + new Promise(() => { + seen = signal; + }), + ); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + void load('a', controller.signal).catch(() => undefined); + void load('b').catch(() => undefined); + await jest.advanceTimersByTimeAsync(50); + controller.abort(); + expect(seen?.aborted).toBe(false); + }); + + it('rejects a caller whose signal had already aborted, without queueing it', async () => { + const fetchMany = jest.fn(async (ids: string[]) => answer(ids)); + const { load } = createPullRequestBatcher({ fetchMany }); + const controller = new AbortController(); + controller.abort(); + await expect(load('a', controller.signal)).rejects.toMatchObject({ code: 'ABORTED' }); + await jest.advanceTimersByTimeAsync(100); + expect(fetchMany).not.toHaveBeenCalled(); + }); +}); diff --git a/client/src/data-provider/PullRequest/batch.ts b/client/src/data-provider/PullRequest/batch.ts new file mode 100644 index 00000000000..9a914e0da25 --- /dev/null +++ b/client/src/data-provider/PullRequest/batch.ts @@ -0,0 +1,300 @@ +import { + PULL_REQUEST_BATCH_MAX, + PULL_REQUEST_BATCH_TIMEOUT_MAX_SECONDS, +} from 'librechat-data-provider'; +import type { + TConversationPullRequestResponse, + TConversationPullRequestsResponse, +} from 'librechat-data-provider'; + +/** A failed entry keeps its stable server code, so a caller can tell a rate limit from an outage. */ +export class PullRequestBatchError extends Error { + constructor(public readonly code: string) { + super(`Pull request lookup failed: ${code}`); + this.name = 'PullRequestBatchError'; + } +} + +type Waiter = { + resolve: (value: TConversationPullRequestResponse) => void; + reject: (reason: unknown) => void; + /** Set when its caller stopped waiting; nothing is sent, or kept running, only for such waiters. */ + cancelled?: boolean; +}; + +/** A request that has been queued or sent and not yet settled. */ +type Dispatch = { waiters: Map; controller?: AbortController }; + +const hasLiveWaiter = (list: Waiter[] | undefined): boolean => + list?.some((waiter) => !waiter.cancelled) === true; + +type BatcherOptions = { + /** + * `signal` aborts when the request has run out of time or the batcher was disposed. A fetcher + * that starts work of its own must stop it on abort: the queue moves on the moment the request + * is settled, so work left running would overlap the next request. + */ + fetchMany: ( + conversationIds: string[], + signal: AbortSignal, + ) => Promise; + /** How long to gather more ids before asking, so rows mounted together share one request. */ + delayMs?: number; + /** Ids per request; the server refuses more than its own bound. */ + maxBatch?: number; + /** + * Longest a request may hold the queue. Requests go out one at a time, so one that never + * answers would otherwise stop every later one. It defaults to just past the longest deadline a + * deployment may configure for the server's own batch, so the server always answers first and + * the next request never starts while the previous one's lookups are still running. + */ + requestTimeoutMs?: number; +}; + +const DEFAULT_DELAY_MS = 30; +const DEFAULT_REQUEST_TIMEOUT_MS = (PULL_REQUEST_BATCH_TIMEOUT_MAX_SECONDS + 10) * 1000; + +/** + * Collects the conversation ids asked for within a short window and fetches them in one request. + * Each caller gets its own conversation's answer, and callers asking for the same conversation + * share one entry. A request that fails rejects only the callers it carried. + */ +export function createPullRequestBatcher({ + fetchMany, + delayMs = DEFAULT_DELAY_MS, + maxBatch = PULL_REQUEST_BATCH_MAX, + requestTimeoutMs = DEFAULT_REQUEST_TIMEOUT_MS, +}: BatcherOptions) { + const pending = new Map(); + let timer: ReturnType | undefined; + let disposed = false; + + /** + * Requests go out one at a time. Each request lets the server run its own number of lookups at + * once, so requests in flight together would multiply the limit the operator configured: a long + * pinned list, or a second flush while one is still out, would otherwise put several times that + * many lookups on GitHub at once. The tail never rejects, so one failed request cannot stall + * the ones queued behind it. + */ + let tail: Promise = Promise.resolve(); + /** Aborts whatever request is on the wire, so disposal does not leave its work running. */ + let active: AbortController | undefined; + /** Requests queued or on the wire, so a cancel can tell when one has nobody left waiting. */ + const dispatches = new Set(); + + const send = (dispatch: Dispatch): Promise => { + const { waiters } = dispatch; + /** Ids nobody waits for any more (their rows unmounted while queued) are never asked for. */ + const ids = [...waiters.keys()].filter((id) => hasLiveWaiter(waiters.get(id))); + if (!disposed && ids.length === 0) { + dispatches.delete(dispatch); + return Promise.resolve(); + } + /** Work that was still waiting its turn when the batcher was disposed never goes out. */ + if (disposed) { + for (const list of waiters.values()) { + list.forEach((waiter) => waiter.reject(new PullRequestBatchError('DISPOSED'))); + } + return Promise.resolve(); + } + const settle = ({ results }: TConversationPullRequestsResponse) => { + const byId = new Map(results.map((entry) => [entry.conversationId, entry])); + for (const [id, list] of waiters) { + const entry = byId.get(id); + for (const waiter of list) { + if (entry != null && 'error' in entry) { + waiter.reject(new PullRequestBatchError(entry.error.code)); + } else { + waiter.resolve({ pullRequest: entry?.pullRequest ?? null }); + } + } + } + }; + const fail = (error: unknown) => { + for (const list of waiters.values()) list.forEach((waiter) => waiter.reject(error)); + }; + const controller = new AbortController(); + dispatch.controller = controller; + active = controller; + let timer: ReturnType | undefined; + /** + * Settles the request on its timeout, and on an abort from disposal, so callers never wait on + * a fetcher that ignores its signal. The timeout aborts first, so the code it reports stays + * `TIMEOUT` and the abort listener only ever fires for disposal. + */ + const timedOut = new Promise((_, reject) => { + timer = setTimeout(() => { + reject(new PullRequestBatchError('TIMEOUT')); + controller.abort(); + }, requestTimeoutMs); + controller.signal.addEventListener('abort', () => + reject(new PullRequestBatchError('ABORTED')), + ); + }); + /** + * `settle` is inside the guarded chain: a response that is not the shape it should be would + * otherwise throw out of it, leave these callers waiting forever and, because the queue chains + * on this promise, stop every later request. Anything that goes wrong here fails the callers + * this request carried and nothing else. + */ + return Promise.race([fetchMany(ids, controller.signal), timedOut]) + .then(settle) + .catch(fail) + .finally(() => { + clearTimeout(timer); + dispatches.delete(dispatch); + if (active === controller) active = undefined; + }); + }; + + const flush = () => { + timer = undefined; + const ids = [...pending.keys()]; + for (let start = 0; start < ids.length; start += maxBatch) { + const chunk = ids.slice(start, start + maxBatch); + const waiters = new Map(); + for (const id of chunk) { + waiters.set(id, pending.get(id) ?? []); + pending.delete(id); + } + const dispatch: Dispatch = { waiters }; + dispatches.add(dispatch); + tail = tail.then(() => send(dispatch)); + } + }; + + return { + /** + * Ends this batcher's life, for the identity it was made for. What has not gone out yet is + * rejected instead of being sent under whoever is signed in next, and nothing queued behind a + * request still in flight waits on it for the next session. A request already on the wire + * finishes on its own. + */ + dispose(): void { + disposed = true; + active?.abort(); + clearTimeout(timer); + timer = undefined; + for (const list of pending.values()) { + list.forEach((waiter) => waiter.reject(new PullRequestBatchError('DISPOSED'))); + } + pending.clear(); + }, + /** + * `signal` withdraws this caller: it is rejected at once, its id leaves the queue if it has not + * gone out, and a request nobody is waiting for any more is not started, or is aborted if it + * is already on the wire, so rows that scrolled away never hold up the ones now visible. + */ + load(conversationId: string, signal?: AbortSignal): Promise { + if (disposed) return Promise.reject(new PullRequestBatchError('DISPOSED')); + if (signal?.aborted) return Promise.reject(new PullRequestBatchError('ABORTED')); + return new Promise((resolve, reject) => { + const waiter: Waiter = { resolve, reject }; + const waiters = pending.get(conversationId) ?? []; + waiters.push(waiter); + pending.set(conversationId, waiters); + timer ??= setTimeout(flush, delayMs); + signal?.addEventListener( + 'abort', + () => { + waiter.cancelled = true; + reject(new PullRequestBatchError('ABORTED')); + const queued = pending.get(conversationId); + if (queued != null) { + const rest = queued.filter((candidate) => candidate !== waiter); + if (rest.length > 0) pending.set(conversationId, rest); + else pending.delete(conversationId); + if (pending.size === 0) { + clearTimeout(timer); + timer = undefined; + } + } + for (const dispatch of dispatches) { + const ids = [...dispatch.waiters.values()]; + if (ids.every((list) => !hasLiveWaiter(list))) dispatch.controller?.abort(); + } + }, + { once: true }, + ); + }); + }, + }; +} + +/** An unknown or unusable limit is one call at a time, never more than the operator allowed. */ +const usableConcurrency = (value: number | undefined): number => + value != null && Number.isInteger(value) && value >= 1 ? value : 1; + +const statusOf = (error: unknown): number | undefined => { + if (error == null || typeof error !== 'object') return undefined; + const candidate = error as { status?: number; response?: { status?: number } }; + return candidate.response?.status ?? candidate.status; +}; + +const codeOf = (error: unknown): string => { + const code = (error as { response?: { data?: { code?: unknown } } } | null)?.response?.data?.code; + return typeof code === 'string' ? code : 'UPSTREAM_ERROR'; +}; + +/** + * Asks the batch route, and when the server that answered does not have it (a replica that has + * not been upgraded yet answers 404, whatever the startup config of another replica advertised), + * asks the single route for each conversation instead, at the limit the server advertised (one at + * a time when it did not say), since the single route has no limiter. Anything but a 404 is a + * real failure and is passed on untouched. + */ +export function createBatchFetcher({ + fetchMany, + fetchOne, + concurrency, +}: { + /** The limit the server advertised, read when the fallback runs; one call at a time without it. */ + concurrency?: () => number | undefined; + fetchMany: ( + conversationIds: string[], + signal: AbortSignal, + ) => Promise; + fetchOne: ( + conversationId: string, + signal: AbortSignal, + ) => Promise; +}) { + return async ( + conversationIds: string[], + signal: AbortSignal, + ): Promise => { + try { + return await fetchMany(conversationIds, signal); + } catch (error) { + if (statusOf(error) !== 404) throw error; + } + const results: TConversationPullRequestsResponse['results'] = new Array(conversationIds.length); + let next = 0; + /** + * The fallback is many requests, and the batcher moves on the moment this one is settled, so + * it stops itself: once the signal aborts no worker starts another call, and the call each + * has in flight is aborted with it, so nothing carries on beside the next request. + */ + const worker = async (): Promise => { + for (let index = next++; index < conversationIds.length; index = next++) { + if (signal.aborted) return; + const conversationId = conversationIds[index]; + try { + const { pullRequest } = await fetchOne(conversationId, signal); + results[index] = { conversationId, pullRequest }; + } catch (error) { + /** After an abort the whole fetch throws below, so a result written here is never read. */ + results[index] = { conversationId, error: { code: codeOf(error) } }; + } + } + }; + await Promise.all( + Array.from( + { length: Math.min(usableConcurrency(concurrency?.()), conversationIds.length) }, + worker, + ), + ); + if (signal.aborted) throw new PullRequestBatchError('ABORTED'); + return { results }; + }; +} diff --git a/client/src/data-provider/PullRequest/index.ts b/client/src/data-provider/PullRequest/index.ts index 3cf1ef310b4..a4878becc89 100644 --- a/client/src/data-provider/PullRequest/index.ts +++ b/client/src/data-provider/PullRequest/index.ts @@ -1 +1,2 @@ export * from './queries'; +export * from './batch'; diff --git a/client/src/data-provider/PullRequest/queries.test.tsx b/client/src/data-provider/PullRequest/queries.test.tsx index fe7448d7c9a..444ed220fb3 100644 --- a/client/src/data-provider/PullRequest/queries.test.tsx +++ b/client/src/data-provider/PullRequest/queries.test.tsx @@ -1,11 +1,20 @@ import React from 'react'; -import { renderHook, waitFor } from '@testing-library/react'; +import { RecoilRoot } from 'recoil'; +import { act, renderHook, waitFor } from '@testing-library/react'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import type { TConversationPullRequest } from 'librechat-data-provider'; -import { pullRequestRefetchInterval, useConversationPullRequestQuery } from './queries'; +import type { TStartupConfig } from 'librechat-data-provider'; +import { + endPullRequestSession, + pullRequestRefetchInterval, + useRowPullRequestQuery, + useConversationPullRequestQuery, +} from './queries'; +import store from '~/store'; const mockGet = jest.fn(); -const mockStartup: { current: { pullRequestsEnabled?: boolean } | undefined } = { +const mockGetMany = jest.fn(); +const mockStartup: { current: Partial | undefined } = { current: { pullRequestsEnabled: true }, }; @@ -20,6 +29,7 @@ jest.mock('librechat-data-provider', () => { dataService: { ...actual.dataService, getConversationPullRequest: (...args: unknown[]) => mockGet(...args), + getConversationPullRequests: (...args: unknown[]) => mockGetMany(...args), }, }; }); @@ -116,3 +126,134 @@ describe('useConversationPullRequestQuery', () => { expect(mockGet).toHaveBeenCalledTimes(1); }); }); + +describe('useRowPullRequestQuery', () => { + const rowStartup = { + pullRequestsEnabled: true, + pullRequestsBatchVersion: 1, + pullRequestsMaxConcurrentLookups: 2, + } as const; + const answerFor = (ids: string[]) => ({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: pr })), + }); + + const wrapRows = (userId = 'user-1') => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const Wrapper = ({ children }: { children: React.ReactNode }) => ( + + set(store.user, { id: userId } as never)}> + {children} + + + ); + return { Wrapper, client }; + }; + + beforeEach(() => { + endPullRequestSession(); + mockGet.mockReset().mockResolvedValue({ pullRequest: null }); + mockGetMany.mockReset().mockImplementation(async (ids: string[]) => answerFor(ids)); + mockStartup.current = { ...rowStartup }; + }); + afterEach(() => endPullRequestSession()); + + it('asks the batch route once for rows mounted together', async () => { + const { Wrapper } = wrapRows(); + const a = renderHook(() => useRowPullRequestQuery('a'), { wrapper: Wrapper }); + const b = renderHook(() => useRowPullRequestQuery('b'), { wrapper: Wrapper }); + await waitFor(() => expect(a.result.current.data).toEqual({ pullRequest: pr })); + await waitFor(() => expect(b.result.current.data).toEqual({ pullRequest: pr })); + expect(mockGetMany).toHaveBeenCalledTimes(1); + expect(mockGetMany.mock.calls[0][0]).toEqual(['a', 'b']); + }); + + it('does not ask for a row that unmounted before its request went out', async () => { + const { Wrapper } = wrapRows(); + const gone = renderHook(() => useRowPullRequestQuery('gone'), { wrapper: Wrapper }); + const kept = renderHook(() => useRowPullRequestQuery('kept'), { wrapper: Wrapper }); + gone.unmount(); + await waitFor(() => expect(kept.result.current.data).toEqual({ pullRequest: pr })); + expect(mockGetMany).toHaveBeenCalledTimes(1); + expect(mockGetMany.mock.calls[0][0]).toEqual(['kept']); + }); + + it('aborts a request on the wire once the last row waiting for it unmounts', async () => { + let seen: AbortSignal | undefined; + mockGetMany.mockReset().mockImplementation( + (_ids: string[], options?: { signal?: AbortSignal }) => + new Promise(() => { + seen = options?.signal; + }), + ); + const { Wrapper } = wrapRows(); + const row = renderHook(() => useRowPullRequestQuery('a'), { wrapper: Wrapper }); + await waitFor(() => expect(seen).toBeDefined()); + expect(seen?.aborted).toBe(false); + row.unmount(); + await waitFor(() => expect(seen?.aborted).toBe(true)); + }); + + it('starts a fresh queue after the session ends, even for the same account', async () => { + const stalled: Array = []; + mockGetMany + .mockReset() + .mockImplementationOnce( + (_ids: string[], options?: { signal?: AbortSignal }) => + new Promise(() => { + stalled.push(options?.signal); + }), + ) + .mockImplementation(async (ids: string[]) => answerFor(ids)); + const first = wrapRows(); + renderHook(() => useRowPullRequestQuery('a'), { wrapper: first.Wrapper }); + await waitFor(() => expect(stalled).toHaveLength(1)); + act(() => endPullRequestSession()); + expect(stalled[0]?.aborted).toBe(true); + const second = wrapRows(); + const row = renderHook(() => useRowPullRequestQuery('b'), { wrapper: second.Wrapper }); + await waitFor(() => expect(row.result.current.data).toEqual({ pullRequest: pr })); + expect(mockGetMany.mock.calls.at(-1)?.[0]).toEqual(['b']); + }); + + it('falls back to the single route no faster than the limit the server advertised', async () => { + const notFound = Object.assign(new Error('Not Found'), { response: { status: 404 } }); + mockGetMany.mockReset().mockRejectedValue(notFound); + let active = 0; + let peak = 0; + mockGet.mockReset().mockImplementation(async () => { + active += 1; + peak = Math.max(peak, active); + await new Promise((resolve) => setTimeout(resolve, 10)); + active -= 1; + return { pullRequest: null }; + }); + const { Wrapper } = wrapRows(); + const rows = ['a', 'b', 'c', 'd', 'e', 'f'].map((id) => + renderHook(() => useRowPullRequestQuery(id), { wrapper: Wrapper }), + ); + await waitFor(() => rows.forEach((row) => expect(row.result.current.isSuccess).toBe(true))); + expect(mockGet).toHaveBeenCalledTimes(6); + expect(peak).toBe(2); + }); + + it('falls back one call at a time when the server advertised no limit', async () => { + mockStartup.current = { pullRequestsEnabled: true, pullRequestsBatchVersion: 1 }; + const notFound = Object.assign(new Error('Not Found'), { response: { status: 404 } }); + mockGetMany.mockReset().mockRejectedValue(notFound); + let active = 0; + let peak = 0; + mockGet.mockReset().mockImplementation(async () => { + active += 1; + peak = Math.max(peak, active); + await new Promise((resolve) => setTimeout(resolve, 5)); + active -= 1; + return { pullRequest: null }; + }); + const { Wrapper } = wrapRows(); + const rows = ['a', 'b', 'c'].map((id) => + renderHook(() => useRowPullRequestQuery(id), { wrapper: Wrapper }), + ); + await waitFor(() => rows.forEach((row) => expect(row.result.current.isSuccess).toBe(true))); + expect(peak).toBe(1); + }); +}); diff --git a/client/src/data-provider/PullRequest/queries.ts b/client/src/data-provider/PullRequest/queries.ts index 21294add9e4..16e135194c7 100644 --- a/client/src/data-provider/PullRequest/queries.ts +++ b/client/src/data-provider/PullRequest/queries.ts @@ -1,8 +1,16 @@ +import { useRecoilValue } from 'recoil'; import { useQuery } from '@tanstack/react-query'; -import { Constants, QueryKeys, dataService } from 'librechat-data-provider'; +import { + Constants, + QueryKeys, + dataService, + PULL_REQUEST_BATCH_VERSION, +} from 'librechat-data-provider'; import type { TConversationPullRequestResponse } from 'librechat-data-provider'; import type { UseQueryOptions } from '@tanstack/react-query'; +import { createBatchFetcher, createPullRequestBatcher } from './batch'; import { useGetStartupConfig } from '../Endpoints'; +import store from '~/store'; const SETTLED_REFRESH_MS = 60_000; /** Checks and conflict state change while a pull request is being worked on. */ @@ -23,6 +31,11 @@ export const pullRequestRefetchInterval = ( return pr.mergeable === 'unknown' ? ACTIVE_REFRESH_MS : SETTLED_REFRESH_MS; }; +const isSavedConversation = (conversationId: string) => + conversationId !== '' && + conversationId !== Constants.NEW_CONVO && + conversationId !== Constants.PENDING_CONVO; + /** * Asks only when the deployment advertises the feature and the conversation is a saved one, so a * default installation never calls the endpoint. The server answers no pull request for a chat @@ -37,11 +50,7 @@ export const useConversationPullRequestQuery = ( [QueryKeys.conversationPullRequest, conversationId], () => dataService.getConversationPullRequest(conversationId), { - enabled: - startupConfig?.pullRequestsEnabled === true && - conversationId !== '' && - conversationId !== Constants.NEW_CONVO && - conversationId !== Constants.PENDING_CONVO, + enabled: startupConfig?.pullRequestsEnabled === true && isSavedConversation(conversationId), staleTime: 15_000, retry: false, refetchOnWindowFocus: true, @@ -51,3 +60,71 @@ export const useConversationPullRequestQuery = ( }, ); }; + +/** What the server last advertised, for a fallback that has no batch limiter to lean on. */ +let advertisedConcurrency: number | undefined; + +const fetchRows = createBatchFetcher({ + fetchMany: (conversationIds, signal) => + dataService.getConversationPullRequests(conversationIds, { signal }), + fetchOne: (conversationId, signal) => + dataService.getConversationPullRequest(conversationId, { signal }), + concurrency: () => advertisedConcurrency, +}); + +type RowBatcher = ReturnType; +let rowBatcher: { userId: string; batcher: RowBatcher } | undefined; + +/** + * The batcher for the signed-in user: every sidebar row mounted in the same moment shares one + * request. It belongs to one identity, so when the user changes the old one is disposed, and ids + * it still held are neither sent under the next account's authorization nor made to wait ahead + * of that account's own rows. + */ +const batcherFor = (userId: string): RowBatcher => { + if (rowBatcher?.userId === userId) return rowBatcher.batcher; + rowBatcher?.batcher.dispose(); + const batcher = createPullRequestBatcher({ fetchMany: fetchRows }); + rowBatcher = { userId, batcher }; + return batcher; +}; + +/** + * Ends the row queue with the session. Called wherever the session's client state is cleared, so + * signing back in as the same account starts a fresh queue and neither waits behind, nor sends, + * anything the previous session left. + */ +export const endPullRequestSession = (): void => { + rowBatcher?.batcher.dispose(); + rowBatcher = undefined; +}; + +/** + * The pull request of a sidebar row. It shares its cache entry with the header's query, so the + * open conversation is never fetched twice, but it never polls: a list of rows would otherwise + * poll GitHub once per row. It refreshes when the window regains focus and a row's answer is + * older than a minute; every mounted row asks in the same moment, so the batcher turns that + * into one request for the whole list rather than one per row. + */ +export const useRowPullRequestQuery = (conversationId: string) => { + const { data: startupConfig } = useGetStartupConfig(); + const userId = useRecoilValue(store.user)?.id ?? ''; + return useQuery( + [QueryKeys.conversationPullRequest, conversationId], + ({ signal }) => { + advertisedConcurrency = startupConfig?.pullRequestsMaxConcurrentLookups; + return batcherFor(userId).load(conversationId, signal); + }, + { + /** The batch route is newer than the flag: ask only a server that says it has it. */ + enabled: + startupConfig?.pullRequestsEnabled === true && + startupConfig.pullRequestsBatchVersion === PULL_REQUEST_BATCH_VERSION && + userId !== '' && + isSavedConversation(conversationId), + staleTime: SETTLED_REFRESH_MS, + retry: false, + refetchOnWindowFocus: true, + }, + ); +}; diff --git a/client/src/hooks/AuthContext.tsx b/client/src/hooks/AuthContext.tsx index 69264b70972..e4752f80556 100644 --- a/client/src/hooks/AuthContext.tsx +++ b/client/src/hooks/AuthContext.tsx @@ -37,6 +37,7 @@ import { import { useGetRole, useGetUserQuery, + endPullRequestSession, useLoginUserMutation, useLogoutUserMutation, useRefreshTokenMutation, @@ -64,6 +65,7 @@ const endSessionClientState = (): void => { getDefaultStore().set(resetFacetsAtom); clearRetainedFileDeletions(); clearComposerDraftStorage(); + endPullRequestSession(); }; /** * Only recognized codes override the HTTP status used by the login error translation. diff --git a/client/src/hooks/__tests__/AuthContext.spec.tsx b/client/src/hooks/__tests__/AuthContext.spec.tsx index 135fc80ac49..f7a667bfc77 100644 --- a/client/src/hooks/__tests__/AuthContext.spec.tsx +++ b/client/src/hooks/__tests__/AuthContext.spec.tsx @@ -52,6 +52,7 @@ let mockCapturedLogoutOptions: { onError: (...args: unknown[]) => void; }; +const mockEndPullRequestSession = jest.fn(); const mockRefreshMutate = jest.fn(); const mockLoginMutate = jest.fn(); const mockLogoutMutate = jest.fn(); @@ -83,6 +84,7 @@ jest.mock('~/data-provider', () => ({ error: null, })), useGetRole: jest.fn(() => ({ data: null })), + endPullRequestSession: jest.fn(() => mockEndPullRequestSession()), useListRoles: jest.fn(() => ({ data: undefined })), })); @@ -508,6 +510,28 @@ describe('AuthContextProvider — logout onSuccess/onError handling', () => { jotaiStore.set(resetChatFilterSessionAtom); }); + it('ends the sidebar pull request queue at the logout session boundary', () => { + mockEndPullRequestSession.mockClear(); + renderProvider(); + + act(() => { + mockCapturedLogoutOptions.onSuccess({ message: 'Logout successful' }); + }); + + expect(mockEndPullRequestSession).toHaveBeenCalled(); + }); + + it('does not end the sidebar pull request queue while a session is being established', () => { + mockEndPullRequestSession.mockClear(); + renderProvider(); + + act(() => { + capturedAuthContext?.completeAuthentication('auth-token', { id: 'user-1' } as never); + }); + + expect(mockEndPullRequestSession).not.toHaveBeenCalled(); + }); + it('does not call window.location.replace when redirect is absent', async () => { const replaceSpy = jest.spyOn(window.location, 'replace').mockImplementation(() => {}); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 8217be56e96..181479917de 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -7,6 +7,7 @@ "com_ui_pr_checks_running": "Running", "com_ui_pr_deletions": "{{0}} lines removed", "com_ui_pr_label": "PR #{{0}}", + "com_ui_pr_load_failed": "Couldn't load the pull request", "com_ui_pr_open_in_github": "Open in GitHub", "com_ui_pr_refresh_failed": "Couldn't refresh the pull request", "com_ui_pr_state": "State: {{0}}", diff --git a/librechat.example.yaml b/librechat.example.yaml index 28ff6059494..655c5153051 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -944,11 +944,17 @@ endpoints: # # Credentials cached at once (1 to 10000), deployment-wide; the one idle longest is dropped # # past it. When principals configure different values, the largest applies. # cacheMaxCredentials: 256 + # # Lookups one sidebar request runs at once (1 to 16); each is a few GitHub requests. + # maxConcurrentLookups: 4 + # # Seconds one sidebar request may stay open in total (1 to 120); entries still waiting then + # # answer with an upstream error and the row shows nothing until the next refresh. + # batchTimeoutSeconds: 20 # # Bounds on one lookup. A repository with many check runs reads one page of 100 per request. # requestTimeoutSeconds: 10 # longest single GitHub request (1 to 60) # lookupTimeoutSeconds: 30 # longest the whole lookup may take (1 to 120) # maxCheckRunPages: 10 # pages read before checks are reported as running (1 to 50) # 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) # # "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. diff --git a/packages/api/src/code/decision.ts b/packages/api/src/code/decision.ts index 32d16e2f8ca..ee6fd69c162 100644 --- a/packages/api/src/code/decision.ts +++ b/packages/api/src/code/decision.ts @@ -20,7 +20,12 @@ export interface ConversationCodeEnvironmentDecision { export type StoredConversationDecision = Pick< TConversation, 'conversationId' | 'codeEnvironmentMode' | 'codeWorkspaces' -> & { codeEnvironmentRevision?: number }; +> & { + codeEnvironmentRevision?: number; + /** Read in the same snapshot as the workspaces, so a lane report can be fenced by exactly the + * attachment the run was admitted on. Absent when the read did not carry it. */ + codeAttachmentEpoch?: number; +}; function canonicalSelections(selections: CodeWorkspaceSelection[]): CodeWorkspaceSelection[] { return canonicalizeCodeWorkspaceSelections(selections); diff --git a/packages/api/src/code/lane.spec.ts b/packages/api/src/code/lane.spec.ts index 922003e5731..e7038e9f99d 100644 --- a/packages/api/src/code/lane.spec.ts +++ b/packages/api/src/code/lane.spec.ts @@ -466,6 +466,41 @@ describe('createLaneGitRecorder target conversation', () => { }); }); + it('fences an ordinary conversation by the epoch of the snapshot the run was admitted on', async () => { + const getConvoLaneContext = jest.fn().mockResolvedValue({ codeAttachmentEpoch: 2 }); + const { record, setConvoLaneGit } = build({ getConvoLaneContext, admittedEpoch: 0 }); + await record(laneGit); + expect(written(setConvoLaneGit)[0].workspace.epoch).toBe(0); + }); + + it('keeps an admitted epoch of zero, rather than reading a newer one in its place', async () => { + const getConvoLaneContext = jest.fn().mockResolvedValue({ codeAttachmentEpoch: 5 }); + const { record, setConvoLaneGit } = build({ getConvoLaneContext, admittedEpoch: 0 }); + await record(laneGit); + expect(written(setConvoLaneGit)[0].workspace.epoch).toBe(0); + }); + + it('falls back to the route read when the admitted decision carried no epoch', async () => { + const getConvoLaneContext = jest.fn().mockResolvedValue({ codeAttachmentEpoch: 3 }); + const { record, setConvoLaneGit } = build({ getConvoLaneContext }); + await record(laneGit); + expect(written(setConvoLaneGit)[0].workspace.epoch).toBe(3); + }); + + it("fences a subagent thread by its root's epoch, which the admitted decision cannot know", async () => { + const getConvoLaneContext = jest.fn().mockResolvedValue({ + subagentThread: { rootConversationId: 'visible-root' }, + codeAttachmentEpoch: 6, + }); + const { record, setConvoLaneGit } = build({ + conversationId: 'child-thread', + getConvoLaneContext, + admittedEpoch: 1, + }); + await record(laneGit); + expect(written(setConvoLaneGit)[0].workspace.epoch).toBe(6); + }); + it('records an ordinary conversation on itself', async () => { const getConvoLaneContext = jest.fn().mockResolvedValue({ codeAttachmentEpoch: 0 }); const { record, setConvoLaneGit } = build({ conversationId: 'plain', getConvoLaneContext }); diff --git a/packages/api/src/code/lane.ts b/packages/api/src/code/lane.ts index 61368ecb27b..77322e9722c 100644 --- a/packages/api/src/code/lane.ts +++ b/packages/api/src/code/lane.ts @@ -79,6 +79,7 @@ export async function createLaneGitRecorder({ conversationId, repo, workspace, + admittedEpoch, getConvoLaneContext, reserveConvoLaneGitSeq, setConvoLaneGit, @@ -89,6 +90,13 @@ export async function createLaneGitRecorder({ conversationId: string | undefined; repo?: string; workspace?: { environmentId: string; workspaceId: string }; + /** + * The attachment epoch read in the same snapshot as the workspace this run was admitted on (the + * admitted decision carries both). For a conversation that is its own lane it is the epoch the + * fence uses, so an A to B to A move between two separate reads cannot satisfy both. A subagent + * thread writes to its root, whose epoch only the route read can supply, so it still reads it. + */ + admittedEpoch?: number; getConvoLaneContext?: LaneContextReader; reserveConvoLaneGitSeq: LaneSeqReserver; setConvoLaneGit: LaneGitWriter; @@ -103,7 +111,7 @@ export async function createLaneGitRecorder({ placed = { conversationId: root || conversationId, required: Boolean(root), - epoch: context?.codeAttachmentEpoch ?? 0, + epoch: (root ? undefined : admittedEpoch) ?? context?.codeAttachmentEpoch ?? 0, }; } catch (error) { logger.warn('[LaneGit] Failed to place a lane recorder', getSafeErrorMetadata(error)); diff --git a/packages/api/src/pulls/capabilities.spec.ts b/packages/api/src/pulls/capabilities.spec.ts new file mode 100644 index 00000000000..dd4288ad5ed --- /dev/null +++ b/packages/api/src/pulls/capabilities.spec.ts @@ -0,0 +1,36 @@ +import { PULL_REQUEST_BATCH_VERSION } from 'librechat-data-provider'; +import type { AppConfig } from '@librechat/data-schemas'; +import { resolvePullRequestCapabilities } from './capabilities'; + +const configWith = (pullRequests?: Record) => + ({ endpoints: { agents: pullRequests ? { pullRequests } : {} } }) as unknown as AppConfig; + +describe('resolvePullRequestCapabilities', () => { + it('advertises the flag, the batch version and the lookup limit together when the feature is on', () => { + expect(resolvePullRequestCapabilities(configWith({ enabled: true }))).toEqual({ + pullRequestsEnabled: true, + pullRequestsBatchVersion: PULL_REQUEST_BATCH_VERSION, + pullRequestsMaxConcurrentLookups: 4, + }); + }); + + it('advertises the configured lookup limit', () => { + expect( + resolvePullRequestCapabilities(configWith({ enabled: true, maxConcurrentLookups: 2 })), + ).toMatchObject({ pullRequestsMaxConcurrentLookups: 2 }); + }); + + it.each([ + ['no config', undefined], + ['a null config', null], + ['a config without endpoints', {} as AppConfig], + ['no pull request settings', configWith()], + ['the feature switched off', configWith({ enabled: false })], + ['a non-true value', configWith({ enabled: 'yes' })], + ])('advertises neither the flag nor a version with %s', (_label, appConfig) => { + const capabilities = resolvePullRequestCapabilities(appConfig); + expect(capabilities).toEqual({ pullRequestsEnabled: false }); + expect(capabilities).not.toHaveProperty('pullRequestsBatchVersion'); + expect(capabilities).not.toHaveProperty('pullRequestsMaxConcurrentLookups'); + }); +}); diff --git a/packages/api/src/pulls/capabilities.ts b/packages/api/src/pulls/capabilities.ts new file mode 100644 index 00000000000..846d876d6e2 --- /dev/null +++ b/packages/api/src/pulls/capabilities.ts @@ -0,0 +1,33 @@ +import { EModelEndpoint, PULL_REQUEST_BATCH_VERSION } from 'librechat-data-provider'; +import type { TAgentsEndpoint } from 'librechat-data-provider'; +import type { AppConfig } from '@librechat/data-schemas'; + +export type PullRequestCapabilities = { + pullRequestsEnabled: boolean; + /** Present only with the feature on, so a client never sees a version without the flag. */ + pullRequestsBatchVersion?: typeof PULL_REQUEST_BATCH_VERSION; + /** The lookups one batch runs at once, so a client that falls back to the single route can + * keep to the same limit while a replica without the batch route is still serving it. */ + pullRequestsMaxConcurrentLookups?: number; +}; + +/** + * What the startup config tells a client about pull requests: whether the deployment turned the + * feature on, and, once it has, which version of the batch route this server serves, so the + * sidebar asks only a server that has it. + */ +export function resolvePullRequestCapabilities( + appConfig: Pick | null | undefined, +): PullRequestCapabilities { + const agents = appConfig?.endpoints?.[EModelEndpoint.agents] as TAgentsEndpoint | undefined; + const enabled = agents?.pullRequests?.enabled === true; + return { + pullRequestsEnabled: enabled, + ...(enabled + ? { + pullRequestsBatchVersion: PULL_REQUEST_BATCH_VERSION, + pullRequestsMaxConcurrentLookups: agents?.pullRequests?.maxConcurrentLookups ?? 4, + } + : {}), + }; +} diff --git a/packages/api/src/pulls/github.spec.ts b/packages/api/src/pulls/github.spec.ts index 2351923b342..93e9b2f2e6e 100644 --- a/packages/api/src/pulls/github.spec.ts +++ b/packages/api/src/pulls/github.spec.ts @@ -486,9 +486,9 @@ describe('matching the recorded head', () => { expect(compares.length).toBeLessThanOrEqual(3); }); - it.each([null, undefined])('does not compare when no head was recorded (%s)', async (head) => { + it('does not compare when the caller never recorded a head', async () => { const { source, fetchFn } = sourceFor(routes('diverged')); - await expect(findWith(source, head)).resolves.toMatchObject({ number: 7 }); + await expect(findWith(source, undefined)).resolves.toMatchObject({ number: 7 }); expect(fetchFn.mock.calls.some(([url]) => String(url).includes('/compare/'))).toBe(false); }); @@ -677,3 +677,260 @@ describe('a secondary rate limit without limiting headers', () => { expect(String(error.message)).not.toContain('secondary'); }); }); + +describe('a 429 that carries no wait', () => { + const body = (message: string) => () => + new Response(JSON.stringify({ message }), { status: 429 }); + + it('is read for the secondary limit message and waits at least a minute', async () => { + const { source } = sourceFor({ + '/pulls?state=open': body('You have exceeded a secondary rate limit.'), + }); + const error = await find(source).catch((caught) => caught); + expect(error.code).toBe('RATE_LIMITED'); + expect(error.retryAfterMs).toBeGreaterThanOrEqual(60_000); + }); + + it('keeps the wait GitHub named, whatever the body says', async () => { + const { source } = sourceFor({ + '/pulls?state=open': () => + new Response(JSON.stringify({ message: 'secondary rate limit' }), { + status: 429, + headers: { 'retry-after': '7' }, + }), + }); + const error = await find(source).catch((caught) => caught); + expect(error.retryAfterMs).toBe(7_000); + }); + + it('leaves the wait unnamed for an ordinary 429', async () => { + const { source } = sourceFor({ '/pulls?state=open': body('Too many requests') }); + const error = await find(source).catch((caught) => caught); + expect(error.code).toBe('RATE_LIMITED'); + expect(error.retryAfterMs).toBeUndefined(); + }); +}); + +describe('check runs of a head that vanished', () => { + it('is an upstream error, not a pull request with no checks', async () => { + const { source } = sourceFor({ + '/pulls?state=open': () => json(listed()), + '/pulls/7': () => json(pull()), + '/check-runs': () => new Response('', { status: 404 }), + }); + await expect(find(source)).rejects.toMatchObject({ code: 'UPSTREAM_ERROR' }); + }); + + it('is an upstream error when a later page of check runs disappears', async () => { + let calls = 0; + const { source } = sourceFor({ + '/pulls?state=open': () => json(listed()), + '/pulls/7': () => json(pull()), + '/check-runs': () => + ++calls === 1 + ? json({ + total_count: 150, + check_runs: Array.from({ length: 100 }, () => ({ + status: 'completed', + conclusion: 'success', + })), + }) + : new Response('', { status: 404 }), + }); + await expect(find(source)).rejects.toMatchObject({ code: 'UPSTREAM_ERROR' }); + }); +}); + +describe('a lane with a branch and no commit yet', () => { + it('matches nothing, so an earlier incarnation of the branch name is never shown', async () => { + const { source, fetchFn } = sourceFor({ + '/pulls?state=open': () => json([]), + '/pulls?state=closed': () => json([{ number: 9, state: 'closed', head: { sha } }]), + '/pulls/9': () => json(pull({ number: 9, state: 'closed', merged: true })), + '/check-runs': () => json({ check_runs: [] }), + }); + await expect( + source.find({ repo: 'o/r', branch: 'feat/x', head: null, token: 't' }), + ).resolves.toBeNull(); + expect(fetchFn).not.toHaveBeenCalled(); + }); + + it('still matches by branch when the caller never recorded a head', async () => { + const { source } = sourceFor({ + '/pulls?state=open': () => json(listed()), + '/pulls/7': () => json(pull()), + '/check-runs': () => json({ check_runs: [] }), + }); + await expect(find(source)).resolves.toMatchObject({ number: 7 }); + }); +}); + +describe('open pull requests before closed ones, from any repository', () => { + const recorded = 'c'.repeat(40); + it('prefers an open pull request from a fork over a closed one in the base repository', async () => { + const { source } = sourceFor({ + '/pulls?state=open': () => json([]), + '/pulls?state=closed': () => json([{ number: 5, state: 'closed', head: { sha: recorded } }]), + '/pulls?per_page': () => + json([{ number: 12, state: 'open', head: { sha: recorded, ref: 'feat/x' } }]), + '/pulls/12': () => json(pull({ number: 12, head: { sha: recorded } })), + '/pulls/5': () => json(pull({ number: 5, state: 'closed', head: { sha: recorded } })), + '/check-runs': () => json({ total_count: 0, check_runs: [] }), + }); + await expect( + source.find({ repo: 'o/r', branch: 'feat/x', head: recorded, token: 't' }), + ).resolves.toMatchObject({ number: 12, state: 'open' }); + }); + + it('still takes a closed pull request when no open one matches anywhere', async () => { + const { source } = sourceFor({ + '/pulls?state=open': () => json([]), + '/pulls?state=closed': () => json([{ number: 5, state: 'closed', head: { sha: recorded } }]), + '/pulls?per_page': () => json([]), + '/pulls/5': () => + json(pull({ number: 5, state: 'closed', merged: true, head: { sha: recorded } })), + '/check-runs': () => json({ total_count: 0, check_runs: [] }), + }); + await expect( + source.find({ repo: 'o/r', branch: 'feat/x', head: recorded, token: 't' }), + ).resolves.toMatchObject({ number: 5, state: 'merged' }); + }); +}); + +describe('the checks of a fork pull request', () => { + const recorded = 'c'.repeat(40); + const routes = () => ({ + '/pulls?state=open': () => json([]), + '/pulls?state=closed': () => json([]), + '/pulls?per_page': () => + json([{ number: 12, state: 'open', head: { sha: recorded, ref: 'feat/x' } }]), + '/pulls/12': () => + json( + pull({ + number: 12, + head: { sha: recorded, repo: { full_name: 'fork-owner/r' } }, + }), + ), + '/check-runs': () => + json({ total_count: 1, check_runs: [{ status: 'completed', conclusion: 'failure' }] }), + }); + const run = ( + source: ReturnType, + allowedRepositories?: string[], + ) => + source.find({ + repo: 'o/r', + branch: 'feat/x', + head: recorded, + token: 't', + allowedRepositories, + }); + const checkUrls = (fetchFn: jest.Mock) => + fetchFn.mock.calls.map(([url]) => String(url)).filter((url) => url.includes('/check-runs')); + + it('reads them from the fork when an administrator allowed that repository', async () => { + const { source, fetchFn } = sourceFor(routes()); + await expect(run(source, ['o/r', 'fork-owner/*'])).resolves.toMatchObject({ + checks: 'failing', + }); + expect(checkUrls(fetchFn)[0]).toContain('/repos/fork-owner/r/commits/'); + }); + + it('never sends the token to a fork nobody allowed, and reads the base repository', async () => { + const { source, fetchFn } = sourceFor(routes()); + await run(source, ['o/r']); + expect(checkUrls(fetchFn)[0]).toContain('/repos/o/r/commits/'); + expect(fetchFn.mock.calls.some(([url]) => String(url).includes('fork-owner'))).toBe(false); + }); + + it('reads the base repository when the head lives there', async () => { + const { source, fetchFn } = sourceFor({ + ...routes(), + '/pulls/12': () => + json(pull({ number: 12, head: { sha: recorded, repo: { full_name: 'O/R' } } })), + }); + await run(source, ['o/r', 'fork-owner/*']); + expect(checkUrls(fetchFn)[0]).toContain('/repos/o/r/commits/'); + }); +}); + +describe('reading more than one page of candidates', () => { + const recorded = 'c'.repeat(40); + const item = (number: number, sha_: string) => ({ number, state: 'open', head: { sha: sha_ } }); + const pageOf = (url: string) => Number(new URL(url).searchParams.get('page') ?? '1'); + + function paged(pages: unknown[][]) { + const fetchFn = jest.fn(async (input: string) => { + const url = String(input); + if (url.includes('/pulls?state=open')) return json(pages[pageOf(url) - 1] ?? []); + if (url.includes('/pulls?state=closed')) return json([]); + if (url.includes('/compare/')) { + return json({ status: url.includes(`...${'d'.repeat(40)}`) ? 'ahead' : 'diverged' }); + } + if (url.includes('/pulls/')) { + return json( + pull({ number: Number(url.split('/pulls/')[1]), head: { sha: 'd'.repeat(40) } }), + ); + } + if (url.includes('/check-runs')) return json({ total_count: 0, check_runs: [] }); + if (url.includes('/commits/')) return json([]); + throw new Error(`unexpected ${url}`); + }); + return { fetchFn, source: createGitHubPullRequestSource({ fetchFn }) }; + } + const first = Array.from({ length: 2 }, (_, i) => item(i + 1, 'e'.repeat(40))); + const second = [item(30, 'd'.repeat(40))]; + const limits = { maxCandidatePullRequests: 2, maxHeadComparisons: 5 }; + + it('reaches an older matching pull request on a later page within the budget', async () => { + const { source } = paged([first, second]); + await expect( + source.find({ + repo: 'o/r', + branch: 'feat/x', + head: recorded, + token: 't', + limits: { ...limits, maxCandidatePages: 2 }, + }), + ).resolves.toMatchObject({ number: 30 }); + }); + + it('stops at one page by default, as before', async () => { + const { source, fetchFn } = paged([first, second]); + await expect( + source.find({ repo: 'o/r', branch: 'feat/x', head: recorded, token: 't', limits }), + ).resolves.toBeNull(); + expect( + fetchFn.mock.calls.filter(([url]) => String(url).includes('/pulls?state=open')), + ).toHaveLength(1); + }); + + it('does not ask for another page after a short one', async () => { + const { source, fetchFn } = paged([[item(1, 'e'.repeat(40))], second]); + await source.find({ + repo: 'o/r', + branch: 'feat/x', + head: recorded, + token: 't', + limits: { ...limits, maxCandidatePages: 5 }, + }); + expect( + fetchFn.mock.calls.filter(([url]) => String(url).includes('/pulls?state=open')), + ).toHaveLength(1); + }); + + it('never reads past the configured page budget', async () => { + const full = [first, first, first, first, first]; + const { source, fetchFn } = paged(full); + await source.find({ + repo: 'o/r', + branch: 'feat/x', + head: recorded, + token: 't', + limits: { ...limits, maxCandidatePages: 3 }, + }); + expect( + fetchFn.mock.calls.filter(([url]) => String(url).includes('/pulls?state=open')), + ).toHaveLength(3); + }); +}); diff --git a/packages/api/src/pulls/github.ts b/packages/api/src/pulls/github.ts index 2b273c0b9af..80d3cff1117 100644 --- a/packages/api/src/pulls/github.ts +++ b/packages/api/src/pulls/github.ts @@ -5,6 +5,7 @@ import type { TConversationPullRequest, } from 'librechat-data-provider'; import type { PullRequestSource } from './types'; +import { isAllowedRepository } from './repository'; import { PullRequestSourceError } from './types'; const GITHUB_API_BASE = 'https://api.github.com'; @@ -15,6 +16,7 @@ const DEFAULT_MAX_CHECK_RUN_PAGES = 10; const CHECK_RUN_PAGE_SIZE = 100; const DEFAULT_MAX_CANDIDATE_PULL_REQUESTS = 10; const DEFAULT_MAX_HEAD_COMPARISONS = 3; +const DEFAULT_MAX_CANDIDATE_PAGES = 1; /** * Comparing `recorded...candidate`: `ahead` means the candidate builds on the recorded commit and * `identical` that it is the recorded commit. `behind` means the candidate does not contain it @@ -48,6 +50,8 @@ type GitHubPull = { mergeable: boolean | null; /** The head this detail describes; the check runs are read for this same commit. */ headSha: string; + /** The repository that head lives in; a fork's checks run there, not in the base repository. */ + headRepo?: string; }; type GitHubCheckRun = { status: string; conclusion: string | null }; @@ -87,6 +91,10 @@ function parsePull(value: unknown): GitHubPull { ) { throw new PullRequestSourceError('UPSTREAM_ERROR'); } + const headRepo = + isRecord(value.head.repo) && typeof value.head.repo.full_name === 'string' + ? value.head.repo.full_name + : undefined; return { number: value.number, title: value.title, @@ -98,6 +106,7 @@ function parsePull(value: unknown): GitHubPull { deletions: value.deletions, mergeable: value.mergeable, headSha: value.head.sha, + ...(headRepo ? { headRepo } : {}), }; } @@ -198,7 +207,7 @@ function isRateLimited(response: Response): boolean { * message. The body is read bounded and only matched, never stored or returned. */ async function isSecondaryLimit(response: Response): Promise { - if (response.status !== 403) return false; + if (response.status !== 403 && response.status !== 429) return false; try { const text = (await response.text()).slice(0, MAX_ERROR_BODY_CHARS); return SECONDARY_LIMIT_MESSAGE.test(text); @@ -245,6 +254,7 @@ export function createGitHubPullRequestSource({ maxCheckRunPages: number; maxCandidatePullRequests: number; maxHeadComparisons: number; + maxCandidatePages: number; }; async function getJson( @@ -252,6 +262,8 @@ export function createGitHubPullRequestSource({ lookup: Lookup, /** Statuses that mean "nothing there" for this request, besides 404. */ absent: readonly number[] = [], + /** Whether a 404 means "nothing there". Check runs for a commit that was just read are not. */ + notFoundIsAbsent = true, ): Promise { let response: Response; try { @@ -275,13 +287,19 @@ export function createGitHubPullRequestSource({ } } if (isRateLimited(response)) { - throw new PullRequestSourceError('RATE_LIMITED', retryAfterMs(response)); + const hinted = retryAfterMs(response); + /** A 429 that names no wait may still be a secondary limit, told apart by its message. */ + const wait = + hinted == null && (await isSecondaryLimit(response)) ? SECONDARY_LIMIT_WAIT_MS : hinted; + throw new PullRequestSourceError('RATE_LIMITED', wait); } if (await isSecondaryLimit(response)) { throw new PullRequestSourceError('RATE_LIMITED', SECONDARY_LIMIT_WAIT_MS); } /** A repository the token cannot see is indistinguishable from one without pull requests. */ - if (response.status === 404 || absent.includes(response.status)) return null; + if ((notFoundIsAbsent && response.status === 404) || absent.includes(response.status)) { + return null; + } throw new PullRequestSourceError('UPSTREAM_ERROR'); } @@ -292,11 +310,14 @@ export function createGitHubPullRequestSource({ ): Promise<{ runs: GitHubCheckRun[]; incomplete: boolean }> { const runs: GitHubCheckRun[] = []; for (let page = 1; page <= lookup.maxCheckRunPages; page++) { + /** The head was read a moment ago, so a 404 here is a head that vanished, not no checks. */ const body = await getJson( `${pathname}?per_page=${CHECK_RUN_PAGE_SIZE}&page=${page}`, lookup, + [], + false, ); - if (body == null) return { runs, incomplete: false }; + if (body == null) throw new PullRequestSourceError('UPSTREAM_ERROR'); const parsed = parseCheckRuns(body); runs.push(...parsed.runs); const done = @@ -309,7 +330,13 @@ export function createGitHubPullRequestSource({ } return { - async find({ repo, branch, head: recordedHead, token, limits }) { + async find({ repo, branch, head: recordedHead, token, limits, allowedRepositories }) { + /** + * A lane with a branch and no commit yet has nothing to match a pull request against, and + * every candidate would be accepted, including an earlier incarnation of a reused branch + * name. `undefined` is a caller that never recorded a head, which keeps matching by branch. + */ + if (recordedHead === null) return null; const lookup: Lookup = { token, requestTimeoutMs: limits?.requestTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS, @@ -318,6 +345,7 @@ export function createGitHubPullRequestSource({ maxCandidatePullRequests: limits?.maxCandidatePullRequests ?? DEFAULT_MAX_CANDIDATE_PULL_REQUESTS, maxHeadComparisons: limits?.maxHeadComparisons ?? DEFAULT_MAX_HEAD_COMPARISONS, + maxCandidatePages: limits?.maxCandidatePages ?? DEFAULT_MAX_CANDIDATE_PAGES, }; const match = REPO_PATTERN.exec(repo); if (match == null || branch.length === 0) return null; @@ -326,9 +354,9 @@ export function createGitHubPullRequestSource({ const base = `/repos/${encodeURIComponent(owner)}/${encodeURIComponent(name)}`; const headFilter = encodeURIComponent(`${owner}:${branch}`); /** `undefined`: the repository is not visible. An empty list: it is, and nothing matched. */ - const listPulls = async (state: 'open' | 'closed', perPage: number) => { + const listPulls = async (state: 'open' | 'closed', perPage: number, page = 1) => { const listed = await getJson( - `${base}/pulls?state=${state}&sort=updated&direction=desc&per_page=${perPage}&head=${headFilter}`, + `${base}/pulls?state=${state}&sort=updated&direction=desc&per_page=${perPage}&page=${page}&head=${headFilter}`, lookup, ); if (listed == null) return undefined; @@ -363,15 +391,35 @@ export function createGitHubPullRequestSource({ } return null; }; + /** + * Reads pages of one state until a short page, the configured page budget, or (without a + * recorded commit, where only the newest matters) the first page. `undefined` is a + * repository the token cannot see. + */ + const listAll = async (state: 'open' | 'closed'): Promise => { + const perPage = + state === 'closed' && recorded == null ? 1 : lookup.maxCandidatePullRequests; + const pages = state === 'closed' && recorded == null ? 1 : lookup.maxCandidatePages; + const all: ListItem[] = []; + for (let page = 1; page <= pages; page++) { + const items = await listPulls(state, perPage, page); + if (items === undefined) return page === 1 ? undefined : all; + all.push(...items); + if (items.length < perPage) break; + } + return all; + }; /** Open first on its own, so closed history on a reused branch name cannot hide it. */ - const open = await listPulls('open', lookup.maxCandidatePullRequests); + const open = await listAll('open'); if (open === undefined) return null; /** * The head filter names the repository's own owner, so a pull request from a fork is not in * these lists. The recorded commit finds it instead: GitHub lists the pull requests that * carry a commit whatever fork they come from, and the branch name narrows them. */ - const forkCandidates = async (): Promise => { + let forks: Promise | undefined; + const forkCandidates = (): Promise => (forks ??= findForkCandidates()); + const findForkCandidates = async (): Promise => { if (recorded == null) return []; const found = await getJson( `${base}/commits/${recorded}/pulls?per_page=${lookup.maxCandidatePullRequests}`, @@ -382,15 +430,17 @@ export function createGitHubPullRequestSource({ if (!Array.isArray(found)) throw new PullRequestSourceError('UPSTREAM_ERROR'); return found .filter((item) => isRecord(item) && isRecord(item.head) && item.head.ref === branch) - .map(parseListItem) - .sort((a, b) => Number(a.state !== 'open') - Number(b.state !== 'open')); + .map(parseListItem); }; + /** + * Open pull requests win over closed ones whichever repository they come from, so a closed + * pull request in the base repository cannot hide an open one from a fork. + */ const chosen = (await choose(open)) ?? - (await choose( - (await listPulls('closed', recorded == null ? 1 : lookup.maxCandidatePullRequests)) ?? [], - )) ?? - (await choose(await forkCandidates())); + (await choose((await forkCandidates()).filter((item) => item.state === 'open'))) ?? + (await choose((await listAll('closed')) ?? [])) ?? + (await choose((await forkCandidates()).filter((item) => item.state !== 'open'))); if (chosen == null) return null; /** @@ -403,7 +453,24 @@ export function createGitHubPullRequestSource({ const pull = parsePull(detail); /** A push between the list and the detail can move the head off the recorded commit. */ if (pull.headSha !== chosen.sha && !(await contains(pull.headSha))) return null; - const checks = await readCheckRuns(`${base}/commits/${pull.headSha}/check-runs`, lookup); + /** + * A fork's checks run in the fork. Its repository is read only when an administrator named + * it, so the server's token is never used against a repository nobody authorized. + */ + const checksBase = (() => { + const headRepo = pull.headRepo; + if (headRepo == null || headRepo.toLowerCase() === `${owner}/${name}`.toLowerCase()) { + return base; + } + const parts = REPO_PATTERN.exec(headRepo); + if (parts == null || !isPathSegment(parts[1]) || !isPathSegment(parts[2])) return base; + if (!isAllowedRepository(headRepo, allowedRepositories)) return base; + return `/repos/${encodeURIComponent(parts[1])}/${encodeURIComponent(parts[2])}`; + })(); + const checks = await readCheckRuns( + `${checksBase}/commits/${pull.headSha}/check-runs`, + lookup, + ); return toConversationPullRequest(pull, checks.runs, checks.incomplete); }, }; diff --git a/packages/api/src/pulls/handler.spec.ts b/packages/api/src/pulls/handler.spec.ts index ec862de7b40..dfbeaddedef 100644 --- a/packages/api/src/pulls/handler.spec.ts +++ b/packages/api/src/pulls/handler.spec.ts @@ -1,7 +1,11 @@ import type { TConversationPullRequest } from 'librechat-data-provider'; import type { Response } from 'express'; import type { ServerRequest } from '~/types'; -import { createConversationPullRequestHandler, resolveTokenReference } from './handler'; +import { + createConversationPullRequestHandler, + createConversationPullRequestsHandler, + resolveTokenReference, +} from './handler'; jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), @@ -84,12 +88,14 @@ describe('createConversationPullRequestHandler', () => { ttlMs: 30_000, cacheMaxEntries: 500, cacheMaxCredentials: 256, + allowedRepositories: ['o/r'], limits: { requestTimeoutMs: 10_000, lookupTimeoutMs: 30_000, maxCheckRunPages: 10, maxCandidatePullRequests: 10, maxHeadComparisons: 3, + maxCandidatePages: 1, }, }); expect(res.status).toHaveBeenCalledWith(200); @@ -104,6 +110,7 @@ describe('createConversationPullRequestHandler', () => { lookupTimeoutSeconds: 8, maxCheckRunPages: 4, maxCandidatePullRequests: 25, + maxCandidatePages: 4, maxHeadComparisons: 6, cacheMaxEntries: 77, cacheMaxCredentials: 12, @@ -120,6 +127,7 @@ describe('createConversationPullRequestHandler', () => { maxCheckRunPages: 4, maxCandidatePullRequests: 25, maxHeadComparisons: 6, + maxCandidatePages: 4, }, }), ); @@ -281,3 +289,645 @@ describe('createConversationPullRequestHandler', () => { expect(JSON.stringify(res.json.mock.calls)).not.toContain('ghp_secret'); }); }); + +describe('createConversationPullRequestsHandler', () => { + type Row = { + conversationId: string; + laneGit: { branch: string | null; head: string | null; repo?: string }; + }; + const lane = (conversationId: string, repo = 'o/r', branch = `feat/${conversationId}`): Row => ({ + conversationId, + laneGit: { branch, head: null, repo }, + }); + + function batch( + options: { + settings?: Record | null; + lanes?: Row[]; + lookup?: jest.Mock; + env?: Record; + } = {}, + ) { + const settings = options.settings === undefined ? enabled : (options.settings ?? undefined); + const lanes = options.lanes ?? [lane('a'), lane('b')]; + const lookup = options.lookup ?? jest.fn().mockResolvedValue({ ok: true, value: pr }); + const getConvosLaneGit = jest.fn().mockResolvedValue(lanes); + const handler = createConversationPullRequestsHandler({ + getConvosLaneGit, + getAppConfig: jest.fn().mockResolvedValue(configWith(settings)), + lookup, + env: options.env ?? { GH_TOKEN: 'ghp_secret' }, + }); + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const run = (...args: [body?: unknown, user?: string | null]) => { + /** Read by count, so an explicit `undefined` is a real "none" and not the default. */ + const body = args.length > 0 ? args[0] : { conversationIds: ['a', 'b'] }; + const user = args.length > 1 ? args[1] : 'u1'; + return handler( + { user: user ? { id: user } : undefined, body } as unknown as ServerRequest, + res as unknown as Response, + ); + }; + return { run, res, lookup, getConvosLaneGit }; + } + + it('answers each conversation in the order asked, from one owner-scoped read', async () => { + const { run, res, getConvosLaneGit, lookup } = batch(); + await run({ conversationIds: ['b', 'a'] }); + expect(getConvosLaneGit).toHaveBeenCalledTimes(1); + expect(getConvosLaneGit).toHaveBeenCalledWith('u1', ['b', 'a']); + expect(lookup).toHaveBeenCalledTimes(2); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'b', pullRequest: pr }, + { conversationId: 'a', pullRequest: pr }, + ], + }); + }); + + it('answers no pull request for a conversation without a stored lane, without a lookup', async () => { + const { run, res, lookup } = batch({ lanes: [lane('a')] }); + await run({ conversationIds: ['a', 'ghost'] }); + expect(lookup).toHaveBeenCalledTimes(1); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', pullRequest: pr }, + { conversationId: 'ghost', pullRequest: null }, + ], + }); + }); + + it('never looks up a repository the administrator did not allow', async () => { + const { run, res, lookup } = batch({ lanes: [lane('a', 'o/r'), lane('b', 'evil/secret')] }); + await run(); + expect(lookup).toHaveBeenCalledTimes(1); + expect(lookup).toHaveBeenCalledWith(expect.objectContaining({ repo: 'o/r' })); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', pullRequest: pr }, + { conversationId: 'b', pullRequest: null }, + ], + }); + }); + + it('keeps one failure from hiding the others and exposes only its code', async () => { + const lookup = jest + .fn() + .mockResolvedValueOnce({ ok: false, error: { code: 'RATE_LIMITED' } }) + .mockResolvedValueOnce({ ok: true, value: pr }); + const { run, res } = batch({ lookup }); + await run(); + expect(res.status).toHaveBeenCalledWith(200); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', error: { code: 'RATE_LIMITED' } }, + { conversationId: 'b', pullRequest: pr }, + ], + }); + }); + + it('answers no pull request for everything, and reads no token, when the feature is off', async () => { + const { run, res, lookup } = batch({ settings: { ...enabled, enabled: false } }); + await run(); + expect(lookup).not.toHaveBeenCalled(); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', pullRequest: null }, + { conversationId: 'b', pullRequest: null }, + ], + }); + }); + + it('says not configured only for the conversations that needed the token, and keeps the rest', async () => { + const { run, res, lookup } = batch({ env: {}, lanes: [lane('a')] }); + await run({ conversationIds: ['a', 'ghost'] }); + expect(res.status).toHaveBeenCalledWith(200); + expect(lookup).not.toHaveBeenCalled(); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', error: { code: 'NOT_CONFIGURED' } }, + { conversationId: 'ghost', pullRequest: null }, + ], + }); + }); + + it('answers a plain no pull request when nothing would have needed the token', async () => { + const { run, res } = batch({ env: {}, lanes: [] }); + await run(); + expect(res.status).toHaveBeenCalledWith(200); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', pullRequest: null }, + { conversationId: 'b', pullRequest: null }, + ], + }); + }); + + it('collapses duplicate ids so one conversation is looked up once', async () => { + const { run, lookup, getConvosLaneGit } = batch({ lanes: [lane('a')] }); + await run({ conversationIds: ['a', 'a', 'a'] }); + expect(getConvosLaneGit).toHaveBeenCalledWith('u1', ['a']); + expect(lookup).toHaveBeenCalledTimes(1); + }); + + it('runs no more lookups at once than configured', async () => { + let active = 0; + let peak = 0; + const lookup = jest.fn(async () => { + active += 1; + peak = Math.max(peak, active); + await new Promise((resolve) => setTimeout(resolve, 5)); + active -= 1; + return { ok: true, value: pr }; + }); + const ids = ['a', 'b', 'c', 'd', 'e', 'f']; + const { run } = batch({ + lookup, + settings: { ...enabled, maxConcurrentLookups: 2 }, + lanes: ids.map((id) => lane(id)), + }); + await run({ conversationIds: ids }); + expect(lookup).toHaveBeenCalledTimes(6); + expect(peak).toBe(2); + }); + + it.each([ + ['no body', undefined], + ['no list', {}], + ['an empty list', { conversationIds: [] }], + ['a list that is not an array', { conversationIds: 'a' }], + ['a non-string id', { conversationIds: ['a', 7] }], + ['a blank id', { conversationIds: ['a', ' '] }], + ['an overlong id', { conversationIds: ['x'.repeat(257)] }], + ['more ids than the bound', { conversationIds: Array.from({ length: 51 }, (_, i) => `c${i}`) }], + ])('rejects %s before reading anything', async (_label, body) => { + const { run, res, getConvosLaneGit, lookup } = batch(); + await run(body); + expect(res.status).toHaveBeenCalledWith(400); + expect(getConvosLaneGit).not.toHaveBeenCalled(); + expect(lookup).not.toHaveBeenCalled(); + }); + + it('accepts exactly the bound', async () => { + const ids = Array.from({ length: 50 }, (_, i) => `c${i}`); + const { run, res } = batch({ lanes: [] }); + await run({ conversationIds: ids }); + expect(res.status).toHaveBeenCalledWith(200); + }); + + it('rejects an unauthenticated request', async () => { + const { run, res, getConvosLaneGit } = batch(); + await run({ conversationIds: ['a'] }, null); + expect(res.status).toHaveBeenCalledWith(400); + expect(getConvosLaneGit).not.toHaveBeenCalled(); + }); + + it('answers 500 without the error text when the read fails', async () => { + const { run, res, getConvosLaneGit } = batch(); + getConvosLaneGit.mockRejectedValue(new Error('mongodb://u:secret@h')); + await run(); + expect(res.status).toHaveBeenCalledWith(500); + expect(JSON.stringify((res.json as jest.Mock).mock.calls)).not.toContain('secret'); + }); + + describe('overall deadline', () => { + const never = () => new Promise(() => undefined); + + it('answers within the deadline when a lookup stalls, keeping what finished', async () => { + const lookup = jest + .fn() + .mockResolvedValueOnce({ ok: true, value: pr }) + .mockImplementationOnce(never); + const { run, res } = batch({ lookup, settings: { ...enabled, batchTimeoutSeconds: 0.08 } }); + const started = Date.now(); + await run(); + expect(Date.now() - started).toBeLessThan(1000); + expect(res.status).toHaveBeenCalledWith(200); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', pullRequest: pr }, + { conversationId: 'b', error: { code: 'UPSTREAM_ERROR' } }, + ], + }); + }); + + it('starts nothing new once the deadline has passed', async () => { + const lookup = jest.fn().mockImplementation(never); + const ids = ['a', 'b', 'c']; + const { run, res } = batch({ + lookup, + lanes: ids.map((id) => lane(id)), + settings: { ...enabled, maxConcurrentLookups: 1, batchTimeoutSeconds: 0.08 }, + }); + await run({ conversationIds: ids }); + expect(lookup).toHaveBeenCalledTimes(1); + expect(res.json).toHaveBeenCalledWith({ + results: ids.map((conversationId) => ({ + conversationId, + error: { code: 'UPSTREAM_ERROR' }, + })), + }); + }); + + it('counts the time spent reading the config and lanes against the deadline', async () => { + /** The reads take most of the 0.2 s deadline, so the lookups get what is left of it, not a fresh 0.2 s. */ + const lookup = jest.fn(() => new Promise(() => undefined)); + const slowRead = (value: T, ms: number) => + new Promise((resolve) => setTimeout(() => resolve(value), ms)); + const handler = createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn(() => slowRead([lane('a')], 150)), + getAppConfig: jest.fn(() => + slowRead(configWith({ ...enabled, batchTimeoutSeconds: 0.2 }), 10), + ) as never, + lookup: lookup as never, + env: { GH_TOKEN: 'ghp_secret' }, + }); + const out = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const started = Date.now(); + await handler( + { user: { id: 'u1' }, body: { conversationIds: ['a'] } } as unknown as ServerRequest, + out as unknown as Response, + ); + const took = Date.now() - started; + expect(took).toBeLessThan(330); + expect(out.json).toHaveBeenCalledWith({ + results: [{ conversationId: 'a', error: { code: 'UPSTREAM_ERROR' } }], + }); + }); + + const respond503 = { + error: 'Pull request lookup is unavailable', + code: 'UPSTREAM_ERROR', + }; + const call = (handler: ReturnType) => { + const out = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const pending = handler( + { user: { id: 'u1' }, body: { conversationIds: ['a'] } } as unknown as ServerRequest, + out as unknown as Response, + ); + return { out, pending }; + }; + + it('ends a request whose lane read never answers at the configured deadline, not the maximum', async () => { + jest.useFakeTimers(); + try { + const lookup = jest.fn(); + const { out, pending } = call( + createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn(() => new Promise(() => undefined)) as never, + getAppConfig: jest + .fn() + .mockResolvedValue(configWith({ ...enabled, batchTimeoutSeconds: 20 })), + lookup: lookup as never, + env: { GH_TOKEN: 'ghp_secret' }, + }), + ); + await jest.advanceTimersByTimeAsync(19_000); + expect(out.json).not.toHaveBeenCalled(); + await jest.advanceTimersByTimeAsync(2_000); + await pending; + expect(out.status).toHaveBeenCalledWith(503); + expect(out.json).toHaveBeenCalledWith(respond503); + expect(lookup).not.toHaveBeenCalled(); + } finally { + jest.useRealTimers(); + } + }); + + it('ends a request whose config read never answers at the default deadline', async () => { + jest.useFakeTimers(); + try { + const lookup = jest.fn(); + const { out, pending } = call( + createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn().mockResolvedValue([lane('a')]) as never, + getAppConfig: jest.fn(() => new Promise(() => undefined)) as never, + lookup: lookup as never, + env: { GH_TOKEN: 'ghp_secret' }, + }), + ); + await jest.advanceTimersByTimeAsync(19_000); + expect(out.json).not.toHaveBeenCalled(); + await jest.advanceTimersByTimeAsync(2_000); + await pending; + expect(out.status).toHaveBeenCalledWith(503); + expect(lookup).not.toHaveBeenCalled(); + } finally { + jest.useRealTimers(); + } + }); + + it('holds the lane read to a configured deadline shorter than the default', async () => { + jest.useFakeTimers(); + try { + const { out, pending } = call( + createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn(() => new Promise(() => undefined)) as never, + getAppConfig: jest + .fn() + .mockResolvedValue(configWith({ ...enabled, batchTimeoutSeconds: 3 })), + lookup: jest.fn() as never, + env: { GH_TOKEN: 'ghp_secret' }, + }), + ); + await jest.advanceTimersByTimeAsync(3_500); + await pending; + expect(out.status).toHaveBeenCalledWith(503); + } finally { + jest.useRealTimers(); + } + }); + + it('starts no lookup when the reads alone used up the deadline', async () => { + const lookup = jest.fn().mockResolvedValue({ ok: true, value: pr }); + const { out, pending } = call( + createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn( + () => new Promise((resolve) => setTimeout(() => resolve([lane('a')]), 150)), + ) as never, + getAppConfig: jest + .fn() + .mockResolvedValue(configWith({ ...enabled, batchTimeoutSeconds: 0.05 })), + lookup: lookup as never, + env: { GH_TOKEN: 'ghp_secret' }, + }), + ); + await pending; + expect(lookup).not.toHaveBeenCalled(); + expect(out.status).toHaveBeenCalledWith(503); + expect(out.json).toHaveBeenCalledWith(respond503); + }); + + it('does not fail unhandled when the lane read rejects after the request has ended', async () => { + const unhandled = jest.fn(); + process.on('unhandledRejection', unhandled); + try { + const { pending } = call( + createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn( + () => new Promise((_, reject) => setTimeout(() => reject(new Error('late')), 30)), + ) as never, + getAppConfig: jest.fn().mockResolvedValue(configWith({ ...enabled, enabled: false })), + lookup: jest.fn() as never, + env: { GH_TOKEN: 'ghp_secret' }, + }), + ); + await pending; + await new Promise((resolve) => setTimeout(resolve, 80)); + expect(unhandled).not.toHaveBeenCalled(); + } finally { + process.off('unhandledRejection', unhandled); + } + }); + + it('does not change the answer of a batch that finishes in time', async () => { + const { run, res } = batch({ settings: { ...enabled, batchTimeoutSeconds: 5 } }); + await run(); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', pullRequest: pr }, + { conversationId: 'b', pullRequest: null }, + ].map((entry, index) => (index === 1 ? { ...entry, pullRequest: pr } : entry)), + }); + }); + }); + + describe('concurrency across requests', () => { + /** Lookups that end only when released, so a request can finish while its lookups still run. */ + const heldLookups = () => { + let active = 0; + let peak = 0; + const releases: Array<() => void> = []; + const lookup = jest.fn( + () => + new Promise((resolve) => { + active += 1; + peak = Math.max(peak, active); + releases.push(() => { + active -= 1; + resolve({ ok: true, value: pr }); + }); + }), + ); + return { lookup, releases, active: () => active, peak: () => peak }; + }; + const ids = ['a', 'b', 'c', 'd']; + const lanes = ids.map((id) => lane(id)); + const settings = { ...enabled, maxConcurrentLookups: 2, batchTimeoutSeconds: 0.08 }; + const flush = () => new Promise((resolve) => setTimeout(resolve, 20)); + + it('keeps the lookups of a timed-out request inside the limit for the next request', async () => { + const held = heldLookups(); + const handler = batch({ lookup: held.lookup, lanes, settings }); + await handler.run({ conversationIds: ids }); + /** The first request has answered, yet its two lookups are still out. */ + expect(held.active()).toBe(2); + const second = handler.run({ conversationIds: ids }); + await flush(); + expect(held.lookup).toHaveBeenCalledTimes(2); + expect(held.peak()).toBe(2); + held.releases.splice(0).forEach((release) => release()); + await second; + held.releases.splice(0).forEach((release) => release()); + expect(held.peak()).toBe(2); + }); + + it('starts no lookup for an entry that gave up while it waited behind another request', async () => { + const held = heldLookups(); + const handler = batch({ lookup: held.lookup, lanes, settings }); + /** The first request ends with its two lookups still holding both slots. */ + await handler.run({ conversationIds: ids }); + expect(held.lookup).toHaveBeenCalledTimes(2); + /** The second request's entries wait for those slots, and give up at its own deadline. */ + await handler.run({ conversationIds: ids }); + expect(held.lookup).toHaveBeenCalledTimes(2); + /** Slots free up only now; the entries that already answered must not start work. */ + held.releases.splice(0).forEach((release) => release()); + await flush(); + expect(held.lookup).toHaveBeenCalledTimes(2); + expect(held.active()).toBe(0); + }); + + it("counts slots per credential, so one tenant's stalled lookups do not block another's", async () => { + const held = heldLookups(); + const configFor = (tokenName: string) => + configWith({ + ...enabled, + token: `\${${tokenName}}`, + maxConcurrentLookups: 1, + batchTimeoutSeconds: 0.2, + }); + /** Each tenant's config names its own token, as principal-scoped configuration can. */ + const getAppConfig = jest.fn(async (options: { userId?: string; role?: string }) => + configFor(options.userId === 'u2' ? 'TOKEN_TWO' : 'TOKEN_ONE'), + ); + const handler = createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn().mockResolvedValue([lane('a')]), + getAppConfig: getAppConfig as never, + lookup: held.lookup as never, + env: { TOKEN_ONE: 'token-one', TOKEN_TWO: 'token-two' }, + }); + const call = async (user: string) => { + const out = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + await handler( + { user: { id: user }, body: { conversationIds: ['a'] } } as unknown as ServerRequest, + out as unknown as Response, + ); + return out; + }; + + /** The first tenant's only slot is now held by a lookup that outlived its request. */ + await call('u1'); + expect(held.active()).toBe(1); + /** Same credential: the next request waits for that slot, gives up, and starts nothing. */ + const same = await call('u1'); + expect(held.lookup).toHaveBeenCalledTimes(1); + expect(same.json).toHaveBeenCalledWith({ + results: [{ conversationId: 'a', error: { code: 'UPSTREAM_ERROR' } }], + }); + /** Another credential has its own slots, so it is not held behind the first tenant. */ + await call('u2'); + expect(held.lookup).toHaveBeenCalledTimes(2); + }); + + it('looks identical lanes up once and spends one slot on them, so an unrelated lane still starts', async () => { + const held = heldLookups(); + /** Four conversations on one branch and one on another, with room for a single lookup at a time. */ + const same = ['a', 'b', 'c', 'd'].map((id) => lane(id, 'o/r', 'feat/shared')); + const other = lane('e', 'o/r', 'feat/other'); + const handler = batch({ + lookup: held.lookup, + lanes: [...same, other], + settings: { ...enabled, maxConcurrentLookups: 1, batchTimeoutSeconds: 5 }, + }); + const done = handler.run({ conversationIds: ['a', 'b', 'c', 'd', 'e'] }); + await flush(); + expect(held.lookup).toHaveBeenCalledTimes(1); + held.releases.splice(0).forEach((release) => release()); + await flush(); + /** The shared branch finished in one lookup, freeing the only slot for the other branch. */ + expect(held.lookup).toHaveBeenCalledTimes(2); + held.releases.splice(0).forEach((release) => release()); + await done; + expect(handler.res.json).toHaveBeenCalledWith({ + results: ['a', 'b', 'c', 'd', 'e'].map((conversationId) => ({ + conversationId, + pullRequest: pr, + })), + }); + }); + + it('joins a lane another request already started instead of taking a slot of its own for it', async () => { + const started: string[] = []; + const releases: Array<() => void> = []; + const lookup = jest.fn( + (input: { branch: string }) => + new Promise((resolve) => { + started.push(input.branch); + releases.push(() => resolve({ ok: true, value: pr })); + }), + ); + const rows = [ + lane('a', 'o/r', 'feat/shared'), + lane('b', 'o/r', 'feat/shared'), + lane('e', 'o/r', 'feat/other'), + ]; + const handler = createConversationPullRequestsHandler({ + getConvosLaneGit: jest.fn(async (_user: string, wanted: string[]) => + rows.filter((row) => wanted.includes(row.conversationId)), + ), + getAppConfig: jest + .fn() + .mockResolvedValue( + configWith({ ...enabled, maxConcurrentLookups: 1, batchTimeoutSeconds: 5 }), + ), + lookup: lookup as never, + env: { GH_TOKEN: 'ghp_secret' }, + }); + const call = (conversationIds: string[]) => { + const out = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const done = handler( + { user: { id: 'u1' }, body: { conversationIds } } as unknown as ServerRequest, + out as unknown as Response, + ); + return { out, done }; + }; + const first = call(['a']); + await flush(); + const second = call(['b', 'e']); + await flush(); + expect(started).toEqual(['feat/shared']); + releases.shift()?.(); + await flush(); + /** The shared lane was answered once for both requests, so the only slot went to the other lane. */ + expect(started).toEqual(['feat/shared', 'feat/other']); + releases.shift()?.(); + await Promise.all([first.done, second.done]); + expect(second.out.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'b', pullRequest: pr }, + { conversationId: 'e', pullRequest: pr }, + ], + }); + }); + + it('gives every conversation on a shared branch the same failure when that lookup fails', async () => { + const lookup = jest + .fn() + .mockResolvedValueOnce({ ok: false, error: { code: 'RATE_LIMITED' } }) + .mockResolvedValue({ ok: true, value: pr }); + const { run, res } = batch({ + lookup, + lanes: [lane('a', 'o/r', 'feat/s'), lane('b', 'o/r', 'feat/s'), lane('c', 'o/r', 'feat/t')], + }); + await run({ conversationIds: ['a', 'b', 'c'] }); + expect(lookup).toHaveBeenCalledTimes(2); + expect(res.json).toHaveBeenCalledWith({ + results: [ + { conversationId: 'a', error: { code: 'RATE_LIMITED' } }, + { conversationId: 'b', error: { code: 'RATE_LIMITED' } }, + { conversationId: 'c', pullRequest: pr }, + ], + }); + }); + + it('does not merge lanes that differ in repository, branch or head', async () => { + const head = 'a'.repeat(40); + const rows = [ + lane('a', 'o/r', 'feat/x'), + lane('b', 'o/r2', 'feat/x'), + lane('c', 'o/r', 'feat/y'), + { conversationId: 'd', laneGit: { branch: 'feat/x', head, repo: 'o/r' } }, + ]; + const { run, lookup } = batch({ + lanes: rows, + settings: { ...enabled, allowedRepositories: ['o/*'] }, + }); + await run({ conversationIds: ['a', 'b', 'c', 'd'] }); + expect(lookup).toHaveBeenCalledTimes(4); + }); + + it('frees the slot when a lookup fails, so later requests are not stuck', async () => { + const lookup = jest + .fn() + .mockRejectedValueOnce(new Error('boom')) + .mockResolvedValue({ ok: true, value: pr }); + const handler = batch({ + lookup, + lanes: [lane('a')], + settings: { ...enabled, maxConcurrentLookups: 1 }, + }); + await handler.run({ conversationIds: ['a'] }); + await handler.run({ conversationIds: ['a'] }); + expect(lookup).toHaveBeenCalledTimes(2); + }); + }); + + it('uses the same lookup input as the single route', async () => { + const single = setup({ laneGit: { branch: 'feat/a', head: null, repo: 'o/r' } }); + await single.run(); + const many = batch({ lanes: [lane('a')] }); + await many.run({ conversationIds: ['a'] }); + expect(many.lookup.mock.calls[0][0]).toEqual(single.lookup.mock.calls[0][0]); + }); +}); diff --git a/packages/api/src/pulls/handler.ts b/packages/api/src/pulls/handler.ts index 42a72097836..e04e8d7b462 100644 --- a/packages/api/src/pulls/handler.ts +++ b/packages/api/src/pulls/handler.ts @@ -1,12 +1,20 @@ +import { createHash } from 'node:crypto'; import { logger } from '@librechat/data-schemas'; -import { EModelEndpoint } from 'librechat-data-provider'; -import type { TAgentsEndpoint, TConversationPullRequestResponse } from 'librechat-data-provider'; +import { EModelEndpoint, PULL_REQUEST_BATCH_MAX } from 'librechat-data-provider'; +import type { + TAgentsEndpoint, + TConversationPullRequest, + TConversationPullRequestsEntry, + TConversationPullRequestResponse, + TConversationPullRequestsResponse, +} from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; import type { Response } from 'express'; import type { GetAppConfigOptions } from '~/app/service'; import type { PullRequestLookup } from './types'; import type { ServerRequest } from '~/types'; import { getAppConfigOptionsFromUser } from '~/app/service'; +import { isAllowedRepository } from './repository'; import { getSafeErrorMetadata } from '~/utils'; const MAX_CONVERSATION_ID_LENGTH = 256; @@ -24,20 +32,198 @@ export function resolveTokenReference( return env[name]?.trim() || null; } -/** `owner/name` or `owner/*`, compared without case as GitHub does. Nothing else matches. */ -export function isAllowedRepository(repo: string, allowed: readonly string[] | undefined): boolean { - const [owner, name] = repo.toLowerCase().split('/'); - return (allowed ?? []).some((entry) => { - const [allowedOwner, allowedName] = entry.toLowerCase().split('/'); - return allowedOwner === owner && (allowedName === '*' || allowedName === name); - }); -} - const validConversationId = (value: string | undefined): value is string => value != null && value.trim() !== '' && value.length <= MAX_CONVERSATION_ID_LENGTH; const NONE: TConversationPullRequestResponse = { pullRequest: null }; +type PullRequestSettings = NonNullable; +type EligibleLane = { branch: string; head: string | null; repo: string }; + +/** Lookups in flight at once for one batch, unless configured. */ +const DEFAULT_BATCH_CONCURRENCY = 4; +/** Longest one batch request may stay open, unless configured. */ +const DEFAULT_BATCH_TIMEOUT_SECONDS = 20; + +const TIMED_OUT = Symbol('timed-out'); + +/** + * Lookups in flight across every batch request of one handler. A lookup that outlives its + * request's deadline keeps running (and keeps its slot) until it really ends, so the next request + * cannot start another on top of it: the configured concurrency holds across requests, not only + * within one. Waiters are served in order. + */ +function createLookupLimiter() { + type Scope = { active: number; waiting: Array<() => void> }; + const scopes = new Map(); + return { + /** Resolves once a slot is free under `limit`; the caller must `release` exactly once. */ + acquire(scopeKey: string, limit: number): Promise { + const scope = scopes.get(scopeKey) ?? { active: 0, waiting: [] }; + scopes.set(scopeKey, scope); + if (scope.active < limit) { + scope.active += 1; + return Promise.resolve(); + } + return new Promise((resolve) => { + scope.waiting.push(() => { + scope.active += 1; + resolve(); + }); + }); + }, + release(scopeKey: string, limit: number): void { + const scope = scopes.get(scopeKey); + if (scope == null) return; + scope.active -= 1; + while (scope.waiting.length > 0 && scope.active < limit) scope.waiting.shift()?.(); + /** An idle scope holds nothing, so credentials that come and go leave no entries behind. */ + if (scope.active === 0 && scope.waiting.length === 0) scopes.delete(scopeKey); + }, + }; +} + +type LookupResultOf = Awaited>; +type SharedLookup = { promise: Promise; interested: number }; + +/** + * One upstream lookup per distinct input across every batch request of a handler. The first + * request that needs a lane starts it and takes its slot; any other request for the same lane, + * however concurrent, joins that promise and takes none. Interest is counted, so a lookup still + * waiting for a slot starts only while some request still wants its answer, and one request + * giving up never cancels what another is waiting for. + */ +function createSharedLookups( + limiter: ReturnType, + lookup: PullRequestLookup, +) { + const shared = new Map(); + return { + /** `leave` must be called by a caller that stops waiting before the answer arrives. */ + join( + key: string, + scope: string, + limit: number, + input: Parameters[0], + ): { promise: SharedLookup['promise']; leave: () => void } { + let entry = shared.get(key); + if (entry == null) { + const created: SharedLookup = { + interested: 0, + promise: Promise.resolve(TIMED_OUT), + }; + created.promise = (async () => { + await limiter.acquire(scope, limit); + try { + if (created.interested === 0) { + /** Nobody is waiting any more; a later request must start its own. */ + if (shared.get(key) === created) shared.delete(key); + return TIMED_OUT; + } + return await lookup(input); + } finally { + limiter.release(scope, limit); + } + })().finally(() => { + if (shared.get(key) === created) shared.delete(key); + }); + shared.set(key, created); + entry = created; + } + const joined = entry; + joined.interested += 1; + let left = false; + return { + promise: joined.promise, + leave: () => { + if (left) return; + left = true; + joined.interested -= 1; + }, + }; + }, + }; +} + +/** A credential's digest, never the token, names the scope a lookup's slot is counted in. */ +const limiterScope = (token: string): string => + createHash('sha256').update(token).digest('hex').slice(0, 32); + +/** Settles with `TIMED_OUT` once `ms` has passed, and never leaves its timer running. */ +async function withDeadline(work: Promise, ms: number): Promise { + if (ms <= 0) return TIMED_OUT; + let timer: ReturnType | undefined; + const expired = new Promise((resolve) => { + timer = setTimeout(() => resolve(TIMED_OUT), ms); + }); + try { + return await Promise.race([work, expired]); + } finally { + clearTimeout(timer); + } +} + +/** A lane the server's token may be used for: it names a branch and an allowed repository. */ +function eligibleLane( + laneGit: ConversationLaneGit | null | undefined, + settings: PullRequestSettings, +): EligibleLane | null { + if (laneGit?.branch == null || laneGit.repo == null) return null; + if (!isAllowedRepository(laneGit.repo, settings.allowedRepositories)) return null; + return { branch: laneGit.branch, head: laneGit.head, repo: laneGit.repo }; +} + +/** Everything one lookup needs from the settings, so the single and batch routes cannot differ. */ +function lookupInput(settings: PullRequestSettings, token: string, lane: EligibleLane) { + return { + repo: lane.repo, + branch: lane.branch, + head: lane.head, + token, + ttlMs: (settings.cacheTtlSeconds ?? 30) * 1000, + cacheMaxEntries: settings.cacheMaxEntries ?? 500, + cacheMaxCredentials: settings.cacheMaxCredentials ?? 256, + allowedRepositories: settings.allowedRepositories ?? [], + limits: { + requestTimeoutMs: (settings.requestTimeoutSeconds ?? 10) * 1000, + lookupTimeoutMs: (settings.lookupTimeoutSeconds ?? 30) * 1000, + maxCheckRunPages: settings.maxCheckRunPages ?? 10, + maxCandidatePullRequests: settings.maxCandidatePullRequests ?? 10, + maxHeadComparisons: settings.maxHeadComparisons ?? 3, + maxCandidatePages: settings.maxCandidatePages ?? 1, + }, + }; +} + +/** Runs `task` over `items` with at most `limit` in flight, keeping the results in order. */ +async function mapWithLimit( + items: readonly T[], + limit: number, + task: (item: T) => Promise, +): Promise { + const results = new Array(items.length); + let next = 0; + const worker = async (): Promise => { + for (let index = next++; index < items.length; index = next++) { + results[index] = await task(items[index]); + } + }; + await Promise.all(Array.from({ length: Math.min(limit, items.length) }, worker)); + return results; +} + +/** Accepts only a non-empty list of valid ids within the batch bound; duplicates collapse. */ +function parseConversationIds(body: unknown): string[] | null { + const ids = (body as { conversationIds?: unknown } | null | undefined)?.conversationIds; + if (!Array.isArray(ids) || ids.length === 0 || ids.length > PULL_REQUEST_BATCH_MAX) return null; + const seen = new Set(); + for (const id of ids) { + if (typeof id !== 'string' || !validConversationId(id)) return null; + seen.add(id); + } + return [...seen]; +} + /** * Serves the pull request of the branch a conversation's code workspace last reported. The * stored branch is read owner-scoped, so another user's conversation is indistinguishable from @@ -77,15 +263,11 @@ export function createConversationPullRequestHandler(deps: { return; } - if (laneGit?.branch == null || laneGit.repo == null) { - res.status(200).json(NONE); - return; - } - /** The repository comes from the worker, so it is never used with the token unless the * administrator named it. A repository that is not allowed looks like one without a pull * request. */ - if (!isAllowedRepository(laneGit.repo, settings.allowedRepositories)) { + const lane = eligibleLane(laneGit, settings); + if (lane == null) { res.status(200).json(NONE); return; } @@ -97,22 +279,7 @@ export function createConversationPullRequestHandler(deps: { return; } - const result = await deps.lookup({ - repo: laneGit.repo, - branch: laneGit.branch, - head: laneGit.head, - token, - ttlMs: (settings.cacheTtlSeconds ?? 30) * 1000, - cacheMaxEntries: settings.cacheMaxEntries ?? 500, - cacheMaxCredentials: settings.cacheMaxCredentials ?? 256, - limits: { - requestTimeoutMs: (settings.requestTimeoutSeconds ?? 10) * 1000, - lookupTimeoutMs: (settings.lookupTimeoutSeconds ?? 30) * 1000, - maxCheckRunPages: settings.maxCheckRunPages ?? 10, - maxCandidatePullRequests: settings.maxCandidatePullRequests ?? 10, - maxHeadComparisons: settings.maxHeadComparisons ?? 3, - }, - }); + const result = await deps.lookup(lookupInput(settings, token, lane)); if (!result.ok) { res.status(503).json({ error: 'Pull request lookup is unavailable', @@ -127,3 +294,167 @@ export function createConversationPullRequestHandler(deps: { } }; } + +/** + * Serves the pull requests of many conversations at once, for the sidebar list. It answers what + * the single route answers for each id, from one owner-scoped read of the stored lanes, so a + * conversation that is missing, expired, another user's or without a lane is simply "no pull + * request". Each entry carries its own result or a stable failure code, so one repository's + * rate limit does not hide the others. Lookups share the single route's cache, in-flight + * sharing and rate-limit cooldown, and run a few at a time. + */ +export function createConversationPullRequestsHandler(deps: { + getConvosLaneGit: ( + user: string, + conversationIds: string[], + ) => Promise>; + getAppConfig: (options: GetAppConfigOptions) => Promise; + lookup: PullRequestLookup; + env: Readonly>; +}) { + const limiter = createLookupLimiter(); + const sharedLookups = createSharedLookups(limiter, deps.lookup); + return async (req: ServerRequest, res: Response): Promise => { + const userId = req.user?.id; + const ids = parseConversationIds(req.body); + if (!userId || ids == null) { + res.status(400).json({ error: 'Invalid conversation list' }); + return; + } + try { + /** + * The clock starts when the request arrives. The configured deadline lives in the config, + * so until the config has loaded the default applies to it; once it has, the lane read and + * everything after it share what is left of the configured deadline, however the two reads + * overlapped. + */ + const startedAt = Date.now(); + const configRead = deps.getAppConfig({ + ...getAppConfigOptionsFromUser(req.user), + skipRuntimeAugmentation: true, + failClosed: true, + }); + const laneRead = deps.getConvosLaneGit(userId, ids); + /** A lane read nobody waits for (feature off, config timed out) must not fail unhandled. */ + laneRead.catch(() => undefined); + const unavailable = () => + res + .status(503) + .json({ error: 'Pull request lookup is unavailable', code: 'UPSTREAM_ERROR' }); + const loaded = await withDeadline(configRead, DEFAULT_BATCH_TIMEOUT_SECONDS * 1000); + if (loaded === TIMED_OUT) { + unavailable(); + return; + } + const appConfig = loaded; + const settings = (appConfig.endpoints?.[EModelEndpoint.agents] as TAgentsEndpoint | undefined) + ?.pullRequests; + const deadline = + startedAt + (settings?.batchTimeoutSeconds ?? DEFAULT_BATCH_TIMEOUT_SECONDS) * 1000; + const none = (): TConversationPullRequestsResponse => ({ + results: ids.map((conversationId) => ({ conversationId, pullRequest: null })), + }); + if (settings?.enabled !== true) { + res.status(200).json(none()); + return; + } + + const lanes = await withDeadline(laneRead, deadline - Date.now()); + if (lanes === TIMED_OUT) { + unavailable(); + return; + } + + const eligible = new Map(); + for (const { conversationId, laneGit } of lanes) { + const lane = eligibleLane(laneGit, settings); + if (lane != null) eligible.set(conversationId, lane); + } + if (eligible.size === 0) { + res.status(200).json(none()); + return; + } + + /** + * A missing token is a fact about the conversations that needed one, not about the request: + * those entries say so, and every other conversation keeps its correct "no pull request". + */ + const token = resolveTokenReference(settings.token, deps.env); + if (token == null) { + logger.warn('[PullRequests] Enabled without a usable token reference'); + res.status(200).json({ + results: ids.map( + (conversationId): TConversationPullRequestsEntry => + eligible.has(conversationId) + ? { conversationId, error: { code: 'NOT_CONFIGURED' } } + : { conversationId, pullRequest: null }, + ), + }); + return; + } + + /** + * One deadline for the whole request, not per lookup: each lookup may run for its own + * configured limit, and a stalled upstream would otherwise hold a full batch open for many + * times that. Past the deadline nothing new starts and the entries still waiting answer + * with an upstream error; a lookup already running finishes in the background under its own + * limit and fills the cache for the next request. + */ + const limit = settings.maxConcurrentLookups ?? DEFAULT_BATCH_CONCURRENCY; + const scope = limiterScope(token); + /** + * Conversations that share a repository, branch and head share one answer, and the lookup + * already coalesces them, so each distinct lookup takes one slot and one deadline, not one + * per conversation. Otherwise duplicates would hold every slot while a single request ran. + */ + const keyOf = (lane: EligibleLane) => `${lane.repo}\0${lane.branch}\0${lane.head ?? ''}`; + const distinct = new Map(); + for (const lane of eligible.values()) distinct.set(keyOf(lane), lane); + type Outcome = + | { ok: true; value: TConversationPullRequest | null } + | { ok: false; code: string }; + const outcomes = new Map(); + await mapWithLimit([...distinct.entries()], limit, async ([key, lane]): Promise => { + /** Checked before anything is joined, so nothing starts once time is up. */ + const remaining = deadline - Date.now(); + if (remaining <= 0) { + outcomes.set(key, { ok: false, code: 'UPSTREAM_ERROR' }); + return; + } + const input = lookupInput(settings, token, lane); + /** The same lane under the same policy is one lookup for every request that wants it. */ + const shareKey = `${scope}\0${JSON.stringify([ + input.repo, + input.branch, + input.head, + input.ttlMs, + input.limits, + input.allowedRepositories, + ])}`; + const joined = sharedLookups.join(shareKey, scope, limit, input); + const result = await withDeadline(joined.promise, remaining); + if (result === TIMED_OUT) { + joined.leave(); + outcomes.set(key, { ok: false, code: 'UPSTREAM_ERROR' }); + return; + } + outcomes.set( + key, + result.ok ? { ok: true, value: result.value } : { ok: false, code: result.error.code }, + ); + }); + const looked = ids.map((conversationId): TConversationPullRequestsEntry => { + const lane = eligible.get(conversationId); + const outcome = lane == null ? undefined : outcomes.get(keyOf(lane)); + if (outcome == null) return { conversationId, pullRequest: null }; + return outcome.ok + ? { conversationId, pullRequest: outcome.value } + : { conversationId, error: { code: outcome.code } }; + }); + res.status(200).json({ results: looked }); + } catch (error) { + logger.error('[PullRequests] Batch handler failed', getSafeErrorMetadata(error)); + res.status(500).json({ error: 'Failed to load the pull requests' }); + } + }; +} diff --git a/packages/api/src/pulls/index.ts b/packages/api/src/pulls/index.ts index f6fca526708..7c5127bb327 100644 --- a/packages/api/src/pulls/index.ts +++ b/packages/api/src/pulls/index.ts @@ -2,3 +2,6 @@ export * from './types'; export * from './github'; export * from './lookup'; export * from './handler'; +export * from './transport'; +export * from './repository'; +export * from './capabilities'; diff --git a/packages/api/src/pulls/lookup.spec.ts b/packages/api/src/pulls/lookup.spec.ts index 319d7c30ff9..ae2ac28a624 100644 --- a/packages/api/src/pulls/lookup.spec.ts +++ b/packages/api/src/pulls/lookup.spec.ts @@ -402,3 +402,44 @@ describe('cache capacity per credential', () => { expect(find).toHaveBeenCalledTimes(before); }); }); + +describe('cache capacity shared by principals with one token', () => { + it('keeps the largest capacity seen, so a smaller one does not trim what a larger keeps', async () => { + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ source: { find } }); + for (const branch of ['a', 'b', 'c', 'd']) { + await lookup({ ...input, branch, cacheMaxEntries: 100 }); + } + await lookup({ ...input, branch: 'e', cacheMaxEntries: 2 }); + const before = find.mock.calls.length; + for (const branch of ['a', 'b', 'c', 'd']) { + await lookup({ ...input, branch, cacheMaxEntries: 2 }); + } + expect(find).toHaveBeenCalledTimes(before); + }); + + it('does not let a larger capacity of one token raise the bound of another', async () => { + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ source: { find } }); + await lookup({ ...input, token: 'big', branch: 'x', cacheMaxEntries: 100 }); + for (const branch of ['a', 'b', 'c']) { + await lookup({ ...input, token: 'small', branch, cacheMaxEntries: 2 }); + } + const before = find.mock.calls.length; + await lookup({ ...input, token: 'small', branch: 'a', cacheMaxEntries: 2 }); + expect(find).toHaveBeenCalledTimes(before + 1); + }); +}); + +describe('the repositories the credential may be used for', () => { + it('passes them to the source, and keeps answers for different lists apart', async () => { + const find = jest.fn().mockResolvedValue(value); + const lookup = createPullRequestLookup({ source: { find } }); + await lookup({ ...input, allowedRepositories: ['o/r'] }); + await lookup({ ...input, allowedRepositories: ['o/r', 'fork/*'] }); + expect(find).toHaveBeenCalledTimes(2); + expect(find.mock.calls[1][0].allowedRepositories).toEqual(['o/r', 'fork/*']); + await lookup({ ...input, allowedRepositories: ['FORK/*', 'o/r'] }); + expect(find).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/api/src/pulls/lookup.ts b/packages/api/src/pulls/lookup.ts index c53fecd324a..ed71c84ed53 100644 --- a/packages/api/src/pulls/lookup.ts +++ b/packages/api/src/pulls/lookup.ts @@ -41,6 +41,9 @@ export function createPullRequestLookup({ /** One bounded partition per credential, so a capacity setting only ever evicts entries * fetched with its own credential. */ const partitions = new Map>(); + /** The largest capacity any caller asked for per credential. Principals can share a token yet + * carry different configured sizes, and the smaller must not trim what the larger keeps. */ + const capacities = new Map(); const inflight = new Map>(); const cooldowns = new Map(); const rateLimited: PullRequestLookupResult = { ok: false, error: { code: 'RATE_LIMITED' } }; @@ -69,6 +72,7 @@ export function createPullRequestLookup({ const idle = partitions.keys().next(); if (idle.done) break; partitions.delete(idle.value); + capacities.delete(idle.value); } return partition; } @@ -83,9 +87,11 @@ export function createPullRequestLookup({ ): 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 }); - while (entries.size > capacity) { + while (entries.size > bound) { const oldest = entries.keys().next(); if (oldest.done) break; entries.delete(oldest.value); @@ -107,6 +113,7 @@ export function createPullRequestLookup({ token, ttlMs, limits, + allowedRepositories, cacheMaxEntries, cacheMaxCredentials, }) => { @@ -119,6 +126,12 @@ export function createPullRequestLookup({ limits?.maxCheckRunPages, limits?.maxCandidatePullRequests, limits?.maxHeadComparisons, + limits?.maxCandidatePages, + /** Which fork a pull request's checks may be read from depends on this list. */ + [...(allowedRepositories ?? [])] + .map((entry) => entry.toLowerCase()) + .sort() + .join('|'), ].join(','); const key = `${scope}\0${repo}#${branch}\0${head ?? ''}\0${policy}`; const cached = touch(scope)?.get(key); @@ -134,7 +147,10 @@ export function createPullRequestLookup({ const run = (async (): Promise => { let result: PullRequestLookupResult; try { - result = { ok: true, value: await source.find({ repo, branch, head, token, limits }) }; + result = { + ok: true, + value: await source.find({ repo, branch, head, token, limits, allowedRepositories }), + }; } catch (error) { if (!(error instanceof PullRequestSourceError)) { logger.warn('[PullRequests] Lookup failed', getSafeErrorMetadata(error)); diff --git a/packages/api/src/pulls/repository.ts b/packages/api/src/pulls/repository.ts new file mode 100644 index 00000000000..3677cd0ba9e --- /dev/null +++ b/packages/api/src/pulls/repository.ts @@ -0,0 +1,8 @@ +/** `owner/name` or `owner/*`, compared without case as GitHub does. Nothing else matches. */ +export function isAllowedRepository(repo: string, allowed: readonly string[] | undefined): boolean { + const [owner, name] = repo.toLowerCase().split('/'); + return (allowed ?? []).some((entry) => { + const [allowedOwner, allowedName] = entry.toLowerCase().split('/'); + return allowedOwner === owner && (allowedName === '*' || allowedName === name); + }); +} diff --git a/packages/api/src/pulls/transport.spec.ts b/packages/api/src/pulls/transport.spec.ts new file mode 100644 index 00000000000..cc04c975faf --- /dev/null +++ b/packages/api/src/pulls/transport.spec.ts @@ -0,0 +1,49 @@ +import type { Dispatcher } from 'undici'; +import { createProxyAwareFetch } from './transport'; + +const dispatcher = { name: 'proxy' } as unknown as Dispatcher; +type Fetch = typeof globalThis.fetch; +const recordingFetch = () => { + const calls: Array<[unknown, RequestInit | undefined]> = []; + const fetchFn = Object.assign( + async (input: unknown, init?: RequestInit) => { + calls.push([input, init]); + return new Response('{}'); + }, + { preconnect: () => undefined }, + ) as unknown as Fetch; + return { fetchFn, calls }; +}; + +describe('createProxyAwareFetch', () => { + it('sends the request through the proxy dispatcher when one is configured', async () => { + const { fetchFn, calls } = recordingFetch(); + const signal = new AbortController().signal; + await createProxyAwareFetch(fetchFn, () => dispatcher)('https://api.github.com/x', { + headers: { a: 'b' }, + signal, + }); + expect(calls).toEqual([ + ['https://api.github.com/x', { headers: { a: 'b' }, signal, dispatcher }], + ]); + }); + + it('goes direct, with the request untouched, when no proxy is configured', async () => { + const { fetchFn, calls } = recordingFetch(); + const init = { headers: { a: 'b' } }; + await createProxyAwareFetch(fetchFn, () => undefined)('https://api.github.com/x', init); + expect(calls).toEqual([['https://api.github.com/x', init]]); + expect((calls[0][1] as { dispatcher?: unknown }).dispatcher).toBeUndefined(); + }); + + it('resolves the dispatcher for each request', async () => { + const { fetchFn, calls } = recordingFetch(); + const getDispatcher = jest.fn().mockReturnValue(undefined); + const wrapped = createProxyAwareFetch(fetchFn, getDispatcher); + await wrapped('https://api.github.com/a'); + getDispatcher.mockReturnValue(dispatcher); + await wrapped('https://api.github.com/b'); + expect(calls[1][1]).toMatchObject({ dispatcher }); + expect((calls[0][1] as { dispatcher?: unknown } | undefined)?.dispatcher).toBeUndefined(); + }); +}); diff --git a/packages/api/src/pulls/transport.ts b/packages/api/src/pulls/transport.ts new file mode 100644 index 00000000000..9d849492dfa --- /dev/null +++ b/packages/api/src/pulls/transport.ts @@ -0,0 +1,21 @@ +import type { Dispatcher } from 'undici'; +import type { PullRequestFetch } from './github'; +import { getEnvProxyDispatcher } from '~/utils/proxy'; + +/** + * Wraps a fetch so every request goes through the deployment's proxy when one is configured + * (`HTTPS_PROXY` and friends, honoring `NO_PROXY`), and goes direct when none is. The dispatcher is + * resolved per request, so an environment that changes at runtime is respected, and the client and + * the dispatcher are both arguments, so a caller or a test supplies its own. + */ +export function createProxyAwareFetch( + fetchFn: typeof globalThis.fetch = globalThis.fetch, + getDispatcher: () => Dispatcher | undefined = getEnvProxyDispatcher, +): PullRequestFetch { + return (input, init) => { + const dispatcher = getDispatcher(); + if (dispatcher == null) return fetchFn(input, init); + const request: RequestInit & { dispatcher?: Dispatcher } = { ...init, dispatcher }; + return fetchFn(input, request); + }; +} diff --git a/packages/api/src/pulls/types.ts b/packages/api/src/pulls/types.ts index e177e67bc6b..c06c1ca1608 100644 --- a/packages/api/src/pulls/types.ts +++ b/packages/api/src/pulls/types.ts @@ -32,6 +32,9 @@ export type PullRequestLookupLimits = { maxCandidatePullRequests?: number; /** Most candidates compared with the recorded commit before the search gives up. */ maxHeadComparisons?: number; + /** Pages of candidates read per state, so a branch name reused more than one page of times + * can still reach an older pull request. One page is what the feature shipped with. */ + maxCandidatePages?: number; }; export type PullRequestFindInput = { @@ -42,6 +45,9 @@ export type PullRequestFindInput = { head?: string | null; token: string; limits?: PullRequestLookupLimits; + /** The repositories the token may be used for. A pull request from a fork has its checks read + * from the fork only when the fork is named here, so the token never reaches an unlisted one. */ + allowedRepositories?: readonly string[]; }; /** Finds the pull request for a branch. Null is a documented absence, not a failure. */ diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index 69092d97ee4..985fcf075eb 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -171,6 +171,8 @@ export const backgroundTasksCancel = (conversationId: string) => export const conversationPullRequest = (conversationId: string) => `${conversationsRoot}/${encodeURIComponent(conversationId)}/pull-request`; +export const conversationPullRequests = () => `${conversationsRoot}/pull-requests`; + export const genTitle = (conversationId: string) => `${conversationsRoot}/gen_title/${encodeURIComponent(conversationId)}`; diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index 99fb5de8b05..b84bdf48c0d 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -1132,9 +1132,12 @@ describe('agent pull request config', () => { lookupTimeoutSeconds: 30, maxCheckRunPages: 10, maxCandidatePullRequests: 10, + maxCandidatePages: 1, maxHeadComparisons: 3, cacheMaxEntries: 500, cacheMaxCredentials: 256, + maxConcurrentLookups: 4, + batchTimeoutSeconds: 20, }); }); @@ -1185,6 +1188,8 @@ describe('agent pull request config', () => { ['maxCheckRunPages', 1.5], ['maxCandidatePullRequests', 0], ['maxCandidatePullRequests', 101], + ['maxCandidatePages', 0], + ['maxCandidatePages', 11], ['maxHeadComparisons', -1], ['maxHeadComparisons', 21], ['maxHeadComparisons', 1.5], @@ -1194,6 +1199,12 @@ describe('agent pull request config', () => { ['cacheMaxCredentials', 0], ['cacheMaxCredentials', 10_001], ['cacheMaxCredentials', 1.5], + ['maxConcurrentLookups', 0], + ['maxConcurrentLookups', 17], + ['maxConcurrentLookups', 1.5], + ['batchTimeoutSeconds', 0], + ['batchTimeoutSeconds', 121], + ['batchTimeoutSeconds', 1.5], ])('rejects %s of %s', (field, value) => { expect(parse({ [field]: value }).success).toBe(false); }); @@ -1205,8 +1216,11 @@ describe('agent pull request config', () => { lookupTimeoutSeconds: 120, maxCheckRunPages: 50, maxCandidatePullRequests: 100, + maxCandidatePages: 10, maxHeadComparisons: 20, cacheMaxEntries: 100_000, + maxConcurrentLookups: 16, + batchTimeoutSeconds: 120, }).success, ).toBe(true); expect(parse({ maxHeadComparisons: 0 }).success).toBe(true); diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 9a86d953651..5a806006a88 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -62,6 +62,10 @@ import { ComponentTypes, SettingTypes, OptionTypes } from './generate'; import { STATEFUL_CODE_ENVIRONMENTS } from './stateful-code'; import { specsConfigSchema, TSpecsConfig } from './models'; import { fileConfigSchema } from './file-config'; +import { + PULL_REQUEST_BATCH_VERSION, + PULL_REQUEST_BATCH_TIMEOUT_MAX_SECONDS, +} from './types/pullRequest'; import { isActionTool } from './types/tools'; import { apiBaseUrl } from './api-endpoints'; import { FileSources } from './types/files'; @@ -1876,6 +1880,17 @@ export const agentsEndpointSchema = baseEndpointSchema * Deployment-wide: when principals configure different values, the largest applies. * Raise it when many tenants each configure their own token. */ cacheMaxCredentials: z.number().int().min(1).max(10_000).optional().default(256), + /** Pull request lookups a single sidebar request runs at once. Each lookup is a few + * GitHub requests, so this bounds how hard one request leans on the token. */ + maxConcurrentLookups: z.number().int().min(1).max(16).optional().default(4), + /** Seconds one sidebar request may stay open in total, whatever its lookups are doing. */ + batchTimeoutSeconds: z + .number() + .int() + .min(1) + .max(PULL_REQUEST_BATCH_TIMEOUT_MAX_SECONDS) + .optional() + .default(20), /** Longest one GitHub request may take. Raise it behind a slow proxy. */ requestTimeoutSeconds: z.number().int().min(1).max(60).optional().default(10), /** Longest a whole lookup, every request together, may hold the header request. */ @@ -1885,6 +1900,9 @@ export const agentsEndpointSchema = baseEndpointSchema /** Pull requests listed per state when matching a branch's history to the commit a chat * last ran at. Raise it for branch names that are reused many times. */ maxCandidatePullRequests: z.number().int().min(1).max(100).optional().default(10), + /** Pages of those candidates read per state, so a branch name reused more often than one + * page holds can still reach an older match. Each page is one request. */ + 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), }) @@ -3137,6 +3155,11 @@ export type TStartupConfig = { insightsEnabled?: boolean; /** `endpoints.agents.pullRequests.enabled`; the header does not ask for a pull request without it. */ pullRequestsEnabled?: boolean; + /** Present with `pullRequestsEnabled` once this server has the batch route the sidebar uses. */ + pullRequestsBatchVersion?: typeof PULL_REQUEST_BATCH_VERSION; + /** `endpoints.agents.pullRequests.maxConcurrentLookups`, so the single-route fallback of an + * upgrade in progress keeps to the limit the operator configured. */ + pullRequestsMaxConcurrentLookups?: number; /** Manual context compaction, gated by the same `summarization.enabled` * switch that governs the automatic detour. */ compactionEnabled?: boolean; diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 0898bed4ab0..72687aef8dd 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -1229,8 +1229,16 @@ export function cancelBackgroundTasks( export function getConversationPullRequest( conversationId: string, + options?: { signal?: AbortSignal }, ): Promise { - return request.get(endpoints.conversationPullRequest(conversationId)); + return request.get(endpoints.conversationPullRequest(conversationId), options); +} + +export function getConversationPullRequests( + conversationIds: string[], + options?: { signal?: AbortSignal }, +): Promise { + return request.post(endpoints.conversationPullRequests(), { conversationIds }, options); } export function getPrompt(id: string): Promise<{ prompt: t.TPrompt }> { diff --git a/packages/data-provider/src/request.ts b/packages/data-provider/src/request.ts index acc84735020..21af6ed95e9 100644 --- a/packages/data-provider/src/request.ts +++ b/packages/data-provider/src/request.ts @@ -19,8 +19,9 @@ async function _getResponse( return await axios.get(url, { ...options }); } -async function _post(url: string, data?: any) { +async function _post(url: string, data?: any, options?: AxiosRequestConfig) { const response = await axios.post(url, JSON.stringify(data), { + ...options, headers: { 'Content-Type': 'application/json' }, }); return response.data; diff --git a/packages/data-provider/src/types/pullRequest.ts b/packages/data-provider/src/types/pullRequest.ts index 0796962e567..02ef2cd42bb 100644 --- a/packages/data-provider/src/types/pullRequest.ts +++ b/packages/data-provider/src/types/pullRequest.ts @@ -22,6 +22,29 @@ export type TConversationPullRequest = { checks: PullRequestChecks; }; +/** Most conversations one batch lookup may name. It bounds the fan-out to GitHub of a single + * request; the lookups themselves stay governed by the configured limits and cache. */ +export const PULL_REQUEST_BATCH_MAX = 50; + +/** Version of the batch lookup route. The server advertises it with `pullRequestsEnabled`, and a + * client that lists conversations asks for batches only when it matches, so during a rolling + * upgrade it never sends a batch to a replica that does not have the route yet. */ +export const PULL_REQUEST_BATCH_VERSION = 1 as const; + +/** The most a deployment may set `batchTimeoutSeconds` to. The server answers by its own deadline, + * so a client that waits just past this never gives up before the server does. */ +export const PULL_REQUEST_BATCH_TIMEOUT_MAX_SECONDS = 120; + +export type TConversationPullRequestsRequest = { conversationIds: string[] }; + +/** One conversation's answer: its pull request (null when it has none), or a stable failure code. */ +export type TConversationPullRequestsEntry = + | { conversationId: string; pullRequest: TConversationPullRequest | null } + | { conversationId: string; error: { code: string } }; + +/** An array rather than an object keyed by id, so no client-supplied id is ever a property name. */ +export type TConversationPullRequestsResponse = { results: TConversationPullRequestsEntry[] }; + /** `pullRequest` is null when the conversation has no pull request, which is not an error. */ export type TConversationPullRequestResponse = { pullRequest: TConversationPullRequest | null; diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index 6215f558d72..ee5a4edeaed 100644 --- a/packages/data-schemas/src/methods/conversation.spec.ts +++ b/packages/data-schemas/src/methods/conversation.spec.ts @@ -10041,6 +10041,23 @@ describe('laneGit review hardening', () => { await expect(write(conversationId, 2, mac)).resolves.toBe(true); }); + it('returns the epoch with the admitted decision, from the same document read', async () => { + const conversationId = await seedAttached([mac]); + await expect( + methods.readAdmittedConvoCodeEnvironmentDecision('lane-user', conversationId), + ).resolves.toMatchObject({ codeWorkspaces: [mac], codeAttachmentEpoch: 0 }); + await move(conversationId, [team]); + await expect( + methods.readAdmittedConvoCodeEnvironmentDecision('lane-user', conversationId), + ).resolves.toMatchObject({ codeWorkspaces: [team], codeAttachmentEpoch: 1 }); + await move(conversationId, [mac]); + const admitted = await methods.readAdmittedConvoCodeEnvironmentDecision( + 'lane-user', + conversationId, + ); + expect(admitted).toMatchObject({ codeWorkspaces: [mac], codeAttachmentEpoch: 2 }); + }); + it('is not advanced by admitting a generation', async () => { const conversationId = await seedAttached([mac]); await methods.readAdmittedConvoCodeEnvironmentDecision('lane-user', conversationId); @@ -10088,6 +10105,101 @@ describe('laneGit review hardening', () => { }); }); + it("does not read another user's conversation as the root", async () => { + const rootId = uuidv4(); + await Conversation.collection.insertOne({ + conversationId: rootId, + user: 'someone-else', + title: 'Not yours', + endpoint: 'agents', + messages: [], + codeAttachmentEpoch: 9, + }); + const threadId = await seedThread(rootId); + await expect(methods.getConvoLaneContext('lane-user', threadId)).resolves.toMatchObject({ + codeAttachmentEpoch: 0, + }); + }); + + it('does not read an expired root, which the write path treats as gone', async () => { + const rootId = await seedAttached([mac]); + await move(rootId, [team]); + await Conversation.collection.updateOne( + { conversationId: rootId }, + { $set: { expiredAt: new Date(Date.now() - 60_000) } }, + ); + const threadId = await seedThread(rootId); + await expect(methods.getConvoLaneContext('lane-user', threadId)).resolves.toMatchObject({ + codeAttachmentEpoch: 0, + }); + }); + + it("does not read another tenant's conversation as the root", async () => { + const rootId = uuidv4(); + await Conversation.collection.insertOne({ + conversationId: rootId, + user: 'lane-user', + title: 'Other tenant', + endpoint: 'agents', + messages: [], + codeAttachmentEpoch: 7, + tenantId: 'tenant-b', + }); + const threadId = uuidv4(); + await Conversation.collection.insertOne({ + conversationId: threadId, + user: 'lane-user', + title: 'Thread', + endpoint: 'agents', + messages: [], + tenantId: 'tenant-a', + subagentThread: { rootConversationId: rootId, parentConversationId: rootId }, + }); + const context = await tenantStorage.run({ tenantId: 'tenant-a' }, () => + methods.getConvoLaneContext('lane-user', threadId), + ); + expect(context).toMatchObject({ codeAttachmentEpoch: 0 }); + }); + + it('reads a root in the same tenant', async () => { + const rootId = uuidv4(); + await Conversation.collection.insertOne({ + conversationId: rootId, + user: 'lane-user', + title: 'Same tenant', + endpoint: 'agents', + messages: [], + codeAttachmentEpoch: 4, + tenantId: 'tenant-a', + }); + const threadId = uuidv4(); + await Conversation.collection.insertOne({ + conversationId: threadId, + user: 'lane-user', + title: 'Thread', + endpoint: 'agents', + messages: [], + tenantId: 'tenant-a', + subagentThread: { rootConversationId: rootId, parentConversationId: rootId }, + }); + const context = await tenantStorage.run({ tenantId: 'tenant-a' }, () => + methods.getConvoLaneContext('lane-user', threadId), + ); + expect(context).toMatchObject({ codeAttachmentEpoch: 4 }); + }); + + it('reads the thread and its root in one database operation', async () => { + const rootId = await seedAttached([mac]); + const threadId = await seedThread(rootId); + const spy = jest.spyOn(Conversation, 'findOne'); + const aggregate = jest.spyOn(Conversation, 'aggregate'); + await methods.getConvoLaneContext('lane-user', threadId); + expect(spy).not.toHaveBeenCalled(); + expect(aggregate).toHaveBeenCalledTimes(1); + spy.mockRestore(); + aggregate.mockRestore(); + }); + it('reads zero when the root is gone, so the write is refused rather than guessed', async () => { const threadId = await seedThread(uuidv4()); await expect(methods.getConvoLaneContext('lane-user', threadId)).resolves.toMatchObject({ @@ -10109,3 +10221,65 @@ describe('laneGit review hardening', () => { }); }); }); + +describe('getConvosLaneGit', () => { + const head = 'a'.repeat(40); + const seedLane = async ( + user: string, + laneGit?: Record, + extra: Record = {}, + ) => { + const conversationId = uuidv4(); + await Conversation.collection.insertOne({ + conversationId, + user, + title: 'Lane', + endpoint: 'agents', + messages: [], + ...(laneGit ? { laneGit: { seq: 3, ...laneGit } } : {}), + ...extra, + }); + return conversationId; + }; + + it("returns the lanes of the owner's conversations in one query, without the sequence number", async () => { + const a = await seedLane('bulk-user', { branch: 'feat/a', head, repo: 'o/r' }); + const b = await seedLane('bulk-user', { branch: 'feat/b', head: null }); + const found = await methods.getConvosLaneGit('bulk-user', [a, b]); + expect(found).toHaveLength(2); + expect(found).toEqual( + expect.arrayContaining([ + { conversationId: a, laneGit: { branch: 'feat/a', head, repo: 'o/r' } }, + { conversationId: b, laneGit: { branch: 'feat/b', head: null } }, + ]), + ); + expect(JSON.stringify(found)).not.toContain('seq'); + }); + + it("leaves out a conversation with no lane, an unknown id and another owner's conversation alike", async () => { + const mine = await seedLane('bulk-user', { branch: 'feat/a', head }); + const noLane = await seedLane('bulk-user'); + const theirs = await seedLane('someone-else', { branch: 'secret', head }); + const found = await methods.getConvosLaneGit('bulk-user', [mine, noLane, theirs, uuidv4()]); + expect(found.map((row) => row.conversationId)).toEqual([mine]); + }); + + it('leaves out an expired temporary chat', async () => { + const expired = await seedLane( + 'bulk-user', + { branch: 'feat/x', head }, + { expiredAt: new Date(Date.now() - 60_000) }, + ); + const live = await seedLane( + 'bulk-user', + { branch: 'feat/y', head }, + { expiredAt: new Date(Date.now() + 3_600_000) }, + ); + const found = await methods.getConvosLaneGit('bulk-user', [expired, live]); + expect(found.map((row) => row.conversationId)).toEqual([live]); + }); + + it('asks nothing for an empty list', async () => { + await expect(methods.getConvosLaneGit('bulk-user', [])).resolves.toEqual([]); + }); +}); diff --git a/packages/data-schemas/src/methods/conversation.ts b/packages/data-schemas/src/methods/conversation.ts index 1391d71b4e6..0bfbcf7928e 100644 --- a/packages/data-schemas/src/methods/conversation.ts +++ b/packages/data-schemas/src/methods/conversation.ts @@ -79,6 +79,9 @@ function stripActorCheckpointFields(record: Record): void { /** Written only by the lane methods, so no save, unset or import may reach them. */ const LANE_PRIVATE_FIELDS = ['laneGit', 'laneGitSeq', 'codeAttachmentEpoch'] as const; +/** The lane state a reader may see; the sequence number that fences writes stays private. */ +export type ConvoLaneGitView = Omit, 'seq'>; + /** What a lane recorder needs to place and fence its writes, read once when its tool is created. */ export type ConvoLaneContext = { subagentThread?: IConversation['subagentThread'] | null; @@ -367,6 +370,10 @@ export interface ConversationMethods { user: string, conversationId: string, ): Promise, 'seq'> | null>; + getConvosLaneGit( + user: string, + conversationIds: string[], + ): Promise>; addConvoToolApprovalAllows(input: { user: string; conversationId: string; @@ -378,7 +385,7 @@ export interface ConversationMethods { conversationId: string, ): Promise | null>; replaceConvoCodeEnvironmentDecision( params: { @@ -3078,22 +3085,59 @@ export function createConversationMethods( conversationId: string, ): Promise { const Conversation = mongoose.models.Conversation as Model; - const stored = await Conversation.findOne( - { user, conversationId, ...activeExpirationFilter() }, - 'subagentThread +codeAttachmentEpoch', - ).lean>(); - if (stored == null) return null; - const rootId = stored.subagentThread?.rootConversationId; + const now = new Date(); + /** + * The thread and its visible root in one operation, so the epoch cannot move between two + * reads. The tenant plugin scopes only the base collection of an aggregate, never the joined + * one, so the join is narrowed here to the same owner, tenant and visibility as the base row: + * a root outside them reads as absent, exactly as it did when it was read on its own. + */ + const [row] = await Conversation.aggregate<{ + subagentThread?: IConversation['subagentThread'] | null; + codeAttachmentEpoch?: number; + _root?: Array<{ codeAttachmentEpoch?: number }>; + }>([ + { $match: { user, conversationId, ...activeExpirationFilter() } }, + { + $lookup: { + from: 'conversations', + localField: 'subagentThread.rootConversationId', + foreignField: 'conversationId', + as: '_root', + }, + }, + { + $project: { + subagentThread: 1, + codeAttachmentEpoch: 1, + _root: { + $filter: { + input: '$_root', + as: 'r', + cond: { + $and: [ + { $eq: ['$$r.user', user] }, + { $eq: [{ $ifNull: ['$$r.tenantId', null] }, { $ifNull: ['$tenantId', null] }] }, + { + $or: [ + { $eq: [{ $ifNull: ['$$r.expiredAt', null] }, null] }, + { $gt: ['$$r.expiredAt', now] }, + ], + }, + ], + }, + }, + }, + }, + }, + { $limit: 1 }, + ]); + if (row == null) return null; + const rootId = row.subagentThread?.rootConversationId; /** A thread's report is written to its visible root, so it is fenced by the root's epoch. */ - const fenced = - rootId && rootId !== conversationId - ? await Conversation.findOne( - { user, conversationId: rootId, ...activeExpirationFilter() }, - '+codeAttachmentEpoch', - ).lean>() - : stored; + const fenced = rootId && rootId !== conversationId ? row._root?.[0] : row; return { - subagentThread: stored.subagentThread ?? null, + subagentThread: row.subagentThread ?? null, codeAttachmentEpoch: fenced?.codeAttachmentEpoch ?? 0, }; } @@ -3214,17 +3258,51 @@ export function createConversationMethods( return { branch, head, ...(repo ? { repo } : {}) }; } + /** + * The last reported lane state of each of the owner's conversations that has one, in one query. + * A conversation that is missing, expired, another owner's or has no lane is simply absent, so a + * caller cannot tell them apart. + */ + async function getConvosLaneGit( + user: string, + conversationIds: string[], + ): Promise> { + if (conversationIds.length === 0) return []; + const Conversation = mongoose.models.Conversation as Model; + const stored = await Conversation.find({ + user, + conversationId: { $in: conversationIds }, + laneGit: { $exists: true, $ne: null }, + ...activeExpirationFilter(), + }) + .select('conversationId laneGit') + .lean>>(); + return stored.flatMap(({ conversationId, laneGit }) => { + if (laneGit == null || conversationId == null) return []; + const { branch, head, repo } = laneGit; + return [{ conversationId, laneGit: { branch, head, ...(repo ? { repo } : {}) } }]; + }); + } + + /** + * The decision a run is admitted on, with the attachment epoch of that same document. The epoch + * is `0` for a conversation never moved, so a lane report can be fenced by exactly the snapshot + * the workspace came from instead of a second read that a move could slip between. + */ async function readAdmittedConvoCodeEnvironmentDecision(user: string, conversationId: string) { const Conversation = mongoose.models.Conversation as Model; - return withoutMeiliIndexing( + const admitted = await withoutMeiliIndexing( Conversation.findOneAndUpdate( { user, conversationId }, { $inc: { codeEnvironmentRevision: 1 } }, { new: true, timestamps: false }, ), ) - .select('conversationId codeEnvironmentMode codeWorkspaces') + .select('conversationId codeEnvironmentMode codeWorkspaces +codeAttachmentEpoch') .lean(); + return admitted == null + ? admitted + : { ...admitted, codeAttachmentEpoch: admitted.codeAttachmentEpoch ?? 0 }; } /** @@ -4452,6 +4530,7 @@ export function createConversationMethods( reserveConvoLaneGitSeq, setConvoLaneGit, getConvoLaneGit, + getConvosLaneGit, getConvoLaneContext, readAdmittedConvoCodeEnvironmentDecision, replaceConvoCodeEnvironmentDecision,