diff --git a/client/src/data-provider/Schedules/admission.ts b/client/src/data-provider/Schedules/admission.ts new file mode 100644 index 0000000000..18358af8f7 --- /dev/null +++ b/client/src/data-provider/Schedules/admission.ts @@ -0,0 +1,169 @@ +/* Scheduled chats */ +import { QueryKeys, dataService } from 'librechat-data-provider'; +import type { TUser, TConversation } from 'librechat-data-provider'; +import type { QueryClient } from '@tanstack/react-query'; +import { extendActiveJobsGrace, queueTitleGeneration } from '~/data-provider/SSE/queries'; +import { upsertConvoInAllQueries, getResponseStatus } from '~/utils'; + +/** + * Waits before each probe, so the first is not spent on a conversation that + * cannot exist yet. Admission is normally a second or two, but a dispatch that + * keeps failing is retried, and the budget has to outlast that or the watch + * gives up on a run the server is still going to admit: the trigger engine + * allows eight attempts backing off from a second and doubling, which is about + * two minutes of retries in the worst case. Roughly five minutes of probes + * covers that window twice over in eleven requests. + * + * It does not cover everything, deliberately. A dispatch refused with a long + * `Retry-After` can be re-driven hours later, and holding a watch open for that + * is worse than what already happens without one — the chat arrives with the + * next list refetch. A later announcement of the same run restarts the watch. + */ +const ADMISSION_DELAYS_MS = [ + 750, 1_500, 3_000, 6_000, 10_000, 15_000, 30_000, 45_000, 60_000, 60_000, 60_000, +]; + +/** 404 is the run not having started yet, which is the whole reason for waiting. + * 408 and 429 ask to be retried by definition, and a proxy in front of the app + * can raise either without the run being involved at all. */ +const RETRYABLE_CLIENT_STATUS = new Set([404, 408, 429]); + +/** + * Whether a failed probe has answered the question for good. A 5xx, or a request + * that never got a response, has not — those are the probe failing rather than the + * run, and the durable delivery is still free to admit afterwards. Only a client + * error that is not asking to be retried settles it: waiting longer will not make + * this client able to read that conversation. + */ +const isTerminalProbeFailure = (error: unknown): boolean => { + const status = getResponseStatus(error); + return status != null && status >= 400 && status < 500 && !RETRYABLE_CLIENT_STATUS.has(status); +}; + +const wait = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)); + +/** Who this cache belongs to. Both sign-in and sign-out empty it wholesale, so a + * watch outliving either has no business writing to whatever replaced it. An + * `undefined` is "not loaded yet" or "emptied", and the two read the same. */ +const cacheOwner = (queryClient: QueryClient): string | undefined => + queryClient.getQueryData([QueryKeys.user])?.id; + +/** + * Runs with a watch in flight or already landed. A run can be announced twice — + * by the click that started it, and by the schedules poll that then lists it — + * and must be watched once. A watch that gave up is forgotten, so a later + * announcement of a run admitted after the budget can try again. + */ +const tracked = new Set(); + +/** @internal Test seam. */ +export function resetTrackedRuns(): void { + tracked.clear(); +} + +/** + * The session fence, in one place. The server enforces ownership on the probe + * itself — a request sent under one account can never return another's chat — + * so the only dangerous response is one SENT under the account that clicked and + * LANDING after someone else signed in. That fixes the shape of the fence: + * + * - The owner is bound only from a reading taken BEFORE a probe goes out. A + * reading taken after one cannot tell a cache that finished loading from a + * cache that changed hands, so it is never adopted as the baseline. + * - Once bound, any different reading — another id, or none — ends the watch. + * - A successful probe with no baseline yet is not written and not abandoned: + * the budget is kept, and the next iteration binds before it probes. + */ +async function admit(queryClient: QueryClient, conversationId: string): Promise { + let startedFor = cacheOwner(queryClient); + for (const delay of ADMISSION_DELAYS_MS) { + await wait(delay); + const before = cacheOwner(queryClient); + if (startedFor === undefined) { + startedFor = before; + } else if (before !== startedFor) { + return false; + } + let conversation: TConversation; + try { + conversation = await dataService.getConversationById(conversationId); + } catch (error) { + if (isTerminalProbeFailure(error)) { + return false; + } + continue; + } + if (startedFor === undefined) { + continue; + } + if (cacheOwner(queryClient) !== startedFor) { + return false; + } + /** The conversation route serves an archived chat; the list query filters one + * out. Another tab archiving this run between its first write and this probe + * would otherwise put a row in the sidebar that no refetch agrees belongs + * there — and the archive was deliberate, so there is nothing else to do. */ + if (conversation.isArchived === true) { + return true; + } + /** Warms the key the chat route reads, so opening the row it just added does + * not have to ask again. Absent-only: anything already cached under this id + * was put there by the chat itself and knows more than one list read does. */ + queryClient.setQueryData( + [QueryKeys.conversation, conversationId], + (current) => current ?? conversation, + ); + upsertConvoInAllQueries(queryClient, conversation); + /** The same bookkeeping the foreground path does when a chat lands in a + * project: the project's count and recent-activity ordering live on the + * project rows, which the conversation caches above do not touch. */ + if (conversation.chatProjectId) { + queryClient.invalidateQueries([QueryKeys.projects]); + queryClient.invalidateQueries([QueryKeys.project, conversation.chatProjectId]); + } + /** A generation is live by definition here. The active-job list is the only + * thing that can mark the row as running, and it stops polling while nothing + * is listed — so re-arm it at admission, which is when a job exists, rather + * than at the click, which is before one does. */ + extendActiveJobsGrace(); + queryClient.invalidateQueries([QueryKeys.activeJobs]); + /** A conversation admitted mid-run has not been titled yet, so the row lands + * as "New Chat". The foreground paths hand that to the title queue, which + * owns the timing the server was configured for; a scheduled run has no + * foreground path, so hand it over here rather than leaving the placeholder + * until an unrelated refetch. */ + queueTitleGeneration(conversationId); + return true; + } + return false; +} + +/** + * Puts a scheduled run's chat in the sidebar once the server actually has it. + * + * The generation that writes the conversation starts after the run is announced + * — after run-now answers, after the schedules poll first lists the occurrence — + * so at the moment of the announcement there is nothing to refetch. The + * conversation route answers 404 until that write commits and applies the same + * visibility rule the list query does, which makes it an exact admission probe: + * a 200 means the chat is listable now, and carries the server's own row rather + * than a guess assembled from the client's cached schedule. + * + * Deliberately detached from the caller: the sidebar has to gain the chat whether + * or not the panel the run was announced in is still mounted. Bounded, and it + * writes only what the server returned — a delivery that never admits leaves + * every cache exactly as it found it. + */ +export async function trackScheduledRun( + queryClient: QueryClient, + conversationId: string, +): Promise { + if (!conversationId || tracked.has(conversationId)) { + return; + } + tracked.add(conversationId); + const landed = await admit(queryClient, conversationId); + if (!landed) { + tracked.delete(conversationId); + } +} diff --git a/client/src/data-provider/Schedules/mutations.ts b/client/src/data-provider/Schedules/mutations.ts index a16e48e2ef..6fbc4f875b 100644 --- a/client/src/data-provider/Schedules/mutations.ts +++ b/client/src/data-provider/Schedules/mutations.ts @@ -8,6 +8,7 @@ import type { TScheduleRunNowResponse, } from 'librechat-data-provider'; import type { UseMutationOptions } from '@tanstack/react-query'; +import { trackScheduledRun } from './admission'; export const useCreateScheduleMutation = ( options?: UseMutationOptions, @@ -73,6 +74,10 @@ export const useRunScheduleNowMutation = ( (id: string) => dataService.runScheduleNow(id), { ...options, + onSuccess: (data, ...rest) => { + void trackScheduledRun(queryClient, data.conversationId); + options?.onSuccess?.(data, ...rest); + }, // Invalidate on SETTLED, not just success: several run-now 409 paths are still // server-side mutations (a balance skip updates lastRun/counters and can // auto-disable; agent/permission/invalid-schedule skips disable the schedule diff --git a/client/src/data-provider/__tests__/scheduleRunNow.test.tsx b/client/src/data-provider/__tests__/scheduleRunNow.test.tsx new file mode 100644 index 0000000000..2a87f25926 --- /dev/null +++ b/client/src/data-provider/__tests__/scheduleRunNow.test.tsx @@ -0,0 +1,438 @@ +import React from 'react'; +import { QueryKeys } from 'librechat-data-provider'; +import { act, renderHook } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import type { TConversation, TScheduleRunNowResponse } from 'librechat-data-provider'; +import type { InfiniteData } from '@tanstack/react-query'; +import type { ReactNode } from 'react'; +import { + ACTIVE_JOBS_POLL_MS, + resetActiveJobsGrace, + getActiveJobsRefetchInterval, +} from '../SSE/queries'; +import { useRunScheduleNowMutation } from '../Schedules/mutations'; +import { resetTrackedRuns } from '../Schedules/admission'; +import * as sseQueries from '../SSE/queries'; + +const mockRunScheduleNow = jest.fn, [string]>(); +const mockGetConversationById = jest.fn, [string]>(); + +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + runScheduleNow: (id: string) => mockRunScheduleNow(id), + getConversationById: (id: string) => mockGetConversationById(id), + }, + }; +}); + +type ConversationPages = InfiniteData<{ + conversations: TConversation[]; + nextCursor: string | null; +}>; + +type ListParams = { projectId?: string }; + +/** Past the first two probes, and still inside the active-job grace window that + * admission opens — a test asserting that window must not advance the clock + * beyond it and read its own fast-forward as the poll having lapsed. */ +const PAST_ADMISSION_MS = 3_000; +/** Past the whole probe budget, so a test can run the watch to exhaustion. */ +const ADMISSION_BUDGET_MS = 600_000; + +/** `isNotFoundError` reads the status off an axios error, so the marker matters: + * a bare `{ response: { status } }` is not one, and the watch would read a 404 as + * an unrelated failure and stop — passing the never-admits tests for the wrong + * reason while failing the ones that retry. */ +const httpError = (status: number, message: string) => + Object.assign(new Error(message), { isAxiosError: true, response: { status } }); + +const archivedConversation = (): TConversation => + ({ ...serverConversation(), isArchived: true }) as TConversation; + +const serverConversation = (chatProjectId?: string): TConversation => + ({ + conversationId: 'run-convo-1', + title: 'Overnight activity summary', + endpoint: 'agents', + agent_id: 'agent-1', + ...(chatProjectId != null ? { chatProjectId } : {}), + createdAt: '2026-09-04T09:00:01.000Z', + updatedAt: '2026-09-04T09:00:01.000Z', + }) as TConversation; + +function createQueryClient() { + return new QueryClient({ + defaultOptions: { + /** These tests run the fake clock past the whole probe budget, which is + * longer than the default five-minute cacheTime — without this the list + * query is garbage collected mid-test and its emptiness reads as the + * watch having removed something. */ + queries: { retry: false, cacheTime: Infinity }, + mutations: { retry: false }, + }, + }); +} + +const createWrapper = (queryClient: QueryClient) => + function Wrapper({ children }: { children: ReactNode }) { + return {children}; + }; + +const signIn = (queryClient: QueryClient, id: string) => { + queryClient.setQueryData([QueryKeys.user], { id }); +}; + +const seedList = (queryClient: QueryClient, params?: ListParams) => { + queryClient.setQueryData( + params ? [QueryKeys.allConversations, params] : [QueryKeys.allConversations], + { + pages: [ + { + conversations: [{ conversationId: 'existing-convo' } as TConversation], + nextCursor: null, + }, + ], + pageParams: [], + }, + ); +}; + +const readList = (queryClient: QueryClient, params?: ListParams) => + queryClient.getQueryData( + params ? [QueryKeys.allConversations, params] : [QueryKeys.allConversations], + )?.pages[0].conversations ?? []; + +const runNow = async (queryClient: QueryClient) => { + const { result } = renderHook(() => useRunScheduleNowMutation(), { + wrapper: createWrapper(queryClient), + }); + await act(async () => { + await result.current.mutateAsync('schedule-1'); + }); +}; + +const settleAdmission = async (ms = PAST_ADMISSION_MS) => { + await act(async () => { + await jest.advanceTimersByTimeAsync(ms); + }); +}; + +describe('run-now conversation tracking', () => { + beforeEach(() => { + /** `resetAllMocks`, not `clearAllMocks`: a test whose `mockRejectedValueOnce` + * queue is not fully drained — which is what a failing retry assertion looks + * like — leaves the remainder queued for the next test, where it silently + * answers the first probe. That coupling let an assertion here pass against a + * missing guard. */ + jest.resetAllMocks(); + jest.useFakeTimers(); + resetActiveJobsGrace(); + resetTrackedRuns(); + mockRunScheduleNow.mockResolvedValue({ + scheduleId: 'schedule-1', + conversationId: 'run-convo-1', + status: 'started', + }); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('adds the run to the sidebar once the server has the conversation', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById + .mockRejectedValueOnce(httpError(404, 'Not Found')) + .mockResolvedValue(serverConversation()); + + await runNow(queryClient); + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + + await settleAdmission(); + + const [first, ...rest] = readList(queryClient); + expect(first.conversationId).toBe('run-convo-1'); + expect(first.title).toBe('Overnight activity summary'); + expect(rest.map((convo) => convo.conversationId)).toEqual(['existing-convo']); + queryClient.clear(); + }); + + it('files the row by the server row, not by a client snapshot of the schedule', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient, { projectId: 'project-server' }); + seedList(queryClient, { projectId: 'project-stale' }); + mockGetConversationById.mockResolvedValue(serverConversation('project-server')); + + await runNow(queryClient); + await settleAdmission(); + + expect(readList(queryClient, { projectId: 'project-server' })[0].conversationId).toBe( + 'run-convo-1', + ); + expect(readList(queryClient, { projectId: 'project-stale' })[0].conversationId).toBe( + 'existing-convo', + ); + queryClient.clear(); + }); + + it('warms the chat route cache so the new row opens without a 404', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById.mockResolvedValue(serverConversation()); + + await runNow(queryClient); + await settleAdmission(); + + expect( + queryClient.getQueryData([QueryKeys.conversation, 'run-convo-1'])?.title, + ).toBe('Overnight activity summary'); + queryClient.clear(); + }); + + it('re-arms the active-job poll only once a generation is known to exist', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById + .mockRejectedValueOnce(httpError(404, 'Not Found')) + .mockResolvedValue(serverConversation()); + + await runNow(queryClient); + expect(getActiveJobsRefetchInterval({ activeJobIds: [] })).toBe(false); + + await settleAdmission(); + + expect(getActiveJobsRefetchInterval({ activeJobIds: [] })).toBe(ACTIVE_JOBS_POLL_MS); + queryClient.clear(); + }); + + it('leaves every cache untouched when the delivery never admits', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById.mockRejectedValue(httpError(404, 'Not Found')); + + await runNow(queryClient); + await settleAdmission(ADMISSION_BUDGET_MS); + + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + expect(queryClient.getQueryData([QueryKeys.conversation, 'run-convo-1'])).toBeUndefined(); + expect(getActiveJobsRefetchInterval({ activeJobIds: [] })).toBe(false); + queryClient.clear(); + }); + + it('keeps waiting through a probe that fails for its own reasons', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById + .mockRejectedValueOnce(httpError(502, 'Bad Gateway')) + .mockRejectedValueOnce(new Error('Network Error')) + .mockResolvedValue(serverConversation()); + + await runNow(queryClient); + await settleAdmission(6_000); + + expect(readList(queryClient)[0].conversationId).toBe('run-convo-1'); + queryClient.clear(); + }); + + it('hands the admitted chat to the title queue, since it lands as "New Chat"', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById.mockResolvedValue(serverConversation()); + /** A spy that calls through: the queue is module state with no reader of its + * own, and the point is that the tracker joins it at all. */ + const queueTitle = jest.spyOn(sseQueries, 'queueTitleGeneration'); + + await runNow(queryClient); + await settleAdmission(); + + expect(queueTitle).toHaveBeenCalledWith('run-convo-1'); + queueTitle.mockRestore(); + queryClient.clear(); + }); + + it('does not file the chat into a cache that changed hands mid-watch', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + + /** The switch lands while the probe is IN FLIGHT, which is the window that + * matters: sign-out and sign-in both empty the cache, and the response then + * arrives holding a chat that belongs to whoever was here before. */ + mockGetConversationById.mockImplementation(async () => { + queryClient.removeQueries([QueryKeys.user]); + seedList(queryClient); + signIn(queryClient, 'user-b'); + return serverConversation(); + }); + + await runNow(queryClient); + await settleAdmission(); + + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + expect(queryClient.getQueryData([QueryKeys.conversation, 'run-convo-1'])).toBeUndefined(); + queryClient.clear(); + }); + + it('keeps waiting through a client error that asks to be retried', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById + .mockRejectedValueOnce(httpError(429, 'Too Many Requests')) + .mockRejectedValueOnce(httpError(408, 'Request Timeout')) + .mockResolvedValue(serverConversation()); + + await runNow(queryClient); + await settleAdmission(6_000); + + expect(readList(queryClient)[0].conversationId).toBe('run-convo-1'); + queryClient.clear(); + }); + + it('does not put an archived chat in the active sidebar', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById.mockResolvedValue(archivedConversation()); + + await runNow(queryClient); + await settleAdmission(); + + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + queryClient.clear(); + }); + + it('still tracks a run started before the user query resolved', async () => { + const queryClient = createQueryClient(); + seedList(queryClient); + mockGetConversationById.mockResolvedValue(serverConversation()); + + /** Deliberately signed out of the CACHE at the click, not signed in later: + * run-now answering 200 already proved a session exists, so the query + * resolving is a load rather than a change of hands. */ + await runNow(queryClient); + signIn(queryClient, 'user-a'); + await settleAdmission(); + + expect(readList(queryClient)[0].conversationId).toBe('run-convo-1'); + queryClient.clear(); + }); + + it('keeps its budget when admitted before the owner is known, and lands once it is', async () => { + const queryClient = createQueryClient(); + seedList(queryClient); + mockGetConversationById.mockResolvedValue(serverConversation()); + + /** The first probe succeeds with nobody to attribute it to: the user query is + * still loading. That must not end the watch — the same account resolves a + * moment later, and the next probe is the one that gets to write. */ + await runNow(queryClient); + await settleAdmission(1_000); + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + + signIn(queryClient, 'user-a'); + await settleAdmission(3_000); + + expect(readList(queryClient)[0].conversationId).toBe('run-convo-1'); + expect(mockGetConversationById).toHaveBeenCalledTimes(2); + queryClient.clear(); + }); + + it('refreshes the project rows when the run was filed under a project', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + queryClient.setQueryData([QueryKeys.projects], { pages: [], pageParams: [] }); + queryClient.setQueryData([QueryKeys.project, 'project-a'], { _id: 'project-a' }); + mockGetConversationById.mockResolvedValue(serverConversation('project-a')); + + await runNow(queryClient); + await settleAdmission(); + + expect(queryClient.getQueryState([QueryKeys.projects])?.isInvalidated).toBe(true); + expect(queryClient.getQueryState([QueryKeys.project, 'project-a'])?.isInvalidated).toBe(true); + queryClient.clear(); + }); + + it('watches a run once however many times it is announced', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById.mockResolvedValue(serverConversation()); + + await runNow(queryClient); + await runNow(queryClient); + await settleAdmission(); + + expect(mockGetConversationById).toHaveBeenCalledTimes(1); + expect(readList(queryClient).filter((c) => c.conversationId === 'run-convo-1')).toHaveLength(1); + queryClient.clear(); + }); + + it('refuses the write when the cache is claimed while the owner is still unknown', async () => { + const queryClient = createQueryClient(); + seedList(queryClient); + /** Never known at the click, and someone else's session arrives while the + * probe is in flight — "loading" and "changed hands" read identically here, + * so the response that was sent for the first account must not be written. + * Every probe after that goes out under the second account, and the server + * enforces ownership on the read, so those answer 404: the watch drains its + * budget and nothing lands. */ + mockGetConversationById + .mockImplementationOnce(async () => { + signIn(queryClient, 'user-b'); + return serverConversation(); + }) + .mockRejectedValue(httpError(404, 'Not Found')); + + await runNow(queryClient); + await settleAdmission(ADMISSION_BUDGET_MS); + + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + expect(queryClient.getQueryData([QueryKeys.conversation, 'run-convo-1'])).toBeUndefined(); + queryClient.clear(); + }); + + it('stops on a failure that settles the matter', async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + mockGetConversationById.mockRejectedValue(httpError(403, 'Forbidden')); + + await runNow(queryClient); + await settleAdmission(ADMISSION_BUDGET_MS); + + expect(mockGetConversationById).toHaveBeenCalledTimes(1); + expect(readList(queryClient).map((convo) => convo.conversationId)).toEqual(['existing-convo']); + queryClient.clear(); + }); + + it("outlasts the trigger engine's own retry window", async () => { + const queryClient = createQueryClient(); + signIn(queryClient, 'user-a'); + seedList(queryClient); + /** The engine allows eight attempts backing off from a second and doubling, + * so a dispatch that keeps failing can admit around two minutes in. */ + mockGetConversationById.mockRejectedValue(httpError(404, 'Not Found')); + + await runNow(queryClient); + await settleAdmission(150_000); + mockGetConversationById.mockResolvedValue(serverConversation()); + await settleAdmission(120_000); + + expect(readList(queryClient)[0].conversationId).toBe('run-convo-1'); + queryClient.clear(); + }); +});