diff --git a/client/src/components/Chat/Steering/recovery.spec.ts b/client/src/components/Chat/Steering/recovery.spec.ts new file mode 100644 index 00000000000..d9783ec56af --- /dev/null +++ b/client/src/components/Chat/Steering/recovery.spec.ts @@ -0,0 +1,66 @@ +import { createStore } from 'jotai'; +import { + blockRecovery, + canRestoreRecovery, + recoveryDisposition, + recoveryDispositionsFamily, +} from './recovery'; + +const item = { recoverySteerId: 'source' }; + +describe('steer recovery safety state', () => { + beforeEach(() => { + sessionStorage.clear(); + recoveryDispositionsFamily.remove('conversation'); + }); + + it('preserves terminal cancellation and dismissal over a late start rejection', () => { + expect(blockRecovery({ source: 'cancelled' }, 'source')).toEqual({ source: 'cancelled' }); + expect(blockRecovery({ source: 'dismissed' }, 'source')).toEqual({ source: 'dismissed' }); + expect(blockRecovery({ source: 'cancelling' }, 'source')).toEqual({ source: 'cancelling' }); + expect(blockRecovery({}, 'source')).toEqual({ source: 'blocked' }); + }); + + it('does not allow stale submission restoration after cancellation or dismissal', () => { + expect(canRestoreRecovery({ source: 'cancelled' }, item)).toBe(false); + expect(canRestoreRecovery({ source: 'dismissed' }, item)).toBe(false); + expect(canRestoreRecovery({ source: 'blocked' }, item)).toBe(true); + expect(recoveryDisposition({ source: 'blocked' }, {})).toBeUndefined(); + }); + + it('retains tab-local safety decisions across reload, without message content', () => { + const first = createStore(); + first.set(recoveryDispositionsFamily('conversation'), { source: 'dismissed' }); + expect(sessionStorage.getItem('steer-recovery:conversation')).toBe('{"source":"dismissed"}'); + recoveryDispositionsFamily.remove('conversation'); + const reloaded = createStore(); + expect(reloaded.get(recoveryDispositionsFamily('conversation'))).toEqual({ + source: 'dismissed', + }); + expect(reloaded.get(recoveryDispositionsFamily('another-conversation'))).toEqual({}); + }); + + it('retains in-memory safety when browser storage is unavailable', () => { + const write = jest.spyOn(Storage.prototype, 'setItem').mockImplementation(() => { + throw new Error('quota'); + }); + try { + const store = createStore(); + expect(() => + store.set(recoveryDispositionsFamily('conversation'), { source: 'blocked' }), + ).not.toThrow(); + expect(store.get(recoveryDispositionsFamily('conversation'))).toEqual({ source: 'blocked' }); + } finally { + write.mockRestore(); + } + }); + + it('treats an interrupted cancellation as held rather than locked or sendable after reload', () => { + const first = createStore(); + first.set(recoveryDispositionsFamily('conversation'), { source: 'cancelling' }); + recoveryDispositionsFamily.remove('conversation'); + expect(createStore().get(recoveryDispositionsFamily('conversation'))).toEqual({ + source: 'blocked', + }); + }); +}); diff --git a/client/src/components/Chat/Steering/recovery.ts b/client/src/components/Chat/Steering/recovery.ts new file mode 100644 index 00000000000..ec9bf070a6d --- /dev/null +++ b/client/src/components/Chat/Steering/recovery.ts @@ -0,0 +1,64 @@ +import { atomFamily, atomWithStorage, createJSONStorage } from 'jotai/utils'; +import type { QueuedMessage } from '~/store/families'; + +export type RecoveryDisposition = 'blocked' | 'cancelling' | 'cancelled' | 'dismissed'; +export type RecoveryDispositions = Partial>; + +/** Receipt-scoped safety decisions survive remounts and reloads in this tab. + * No message content is stored. A crashed cancellation stays held, not sendable. */ +export const recoveryDispositionsFamily = atomFamily((conversationId: string) => + atomWithStorage( + `steer-recovery:${conversationId}`, + {}, + createJSONStorage( + () => ({ + getItem(key) { + try { + return sessionStorage.getItem(key); + } catch { + return null; + } + }, + setItem(key, value) { + try { + sessionStorage.setItem(key, value); + } catch { + /* Keep the in-memory safety decision. */ + } + }, + removeItem(key) { + try { + sessionStorage.removeItem(key); + } catch { + /* Storage may be disabled. */ + } + }, + }), + { + reviver: (_key, value) => (value === 'cancelling' ? 'blocked' : value), + }, + ), + { getOnInit: true }, + ), +); + +export function recoveryDisposition( + dispositions: RecoveryDispositions, + item: Pick, +): RecoveryDisposition | undefined { + return item.recoverySteerId == null ? undefined : dispositions?.[item.recoverySteerId]; +} + +export function canRestoreRecovery( + dispositions: RecoveryDispositions, + item: Pick, +): boolean { + const disposition = recoveryDisposition(dispositions, item); + return disposition !== 'cancelled' && disposition !== 'dismissed'; +} + +export function blockRecovery(dispositions: RecoveryDispositions, steerId: string) { + return dispositions?.[steerId] != null + ? dispositions + : { ...dispositions, [steerId]: 'blocked' as const }; +} diff --git a/client/src/components/Chat/Subagents/ParentSubagentsProvider.test.tsx b/client/src/components/Chat/Subagents/ParentSubagentsProvider.test.tsx index 59eade9bda0..d722f4e432a 100644 --- a/client/src/components/Chat/Subagents/ParentSubagentsProvider.test.tsx +++ b/client/src/components/Chat/Subagents/ParentSubagentsProvider.test.tsx @@ -58,9 +58,11 @@ describe('ParentSubagentsProvider', () => { , ); - expect(mockUseParentSubagentsQuery).toHaveBeenCalledWith('parent-conversation', { - enabled: true, - }); + expect(mockUseParentSubagentsQuery).toHaveBeenCalledWith( + 'parent-conversation', + { enabled: true }, + false, + ); expect(context?.byMessageId.get('parent-message')).toEqual([eventChild]); expect(context?.byThreadId.get('tool-thread')).toEqual(toolChild); let result: ParentSubagentIndex | undefined; @@ -70,4 +72,18 @@ describe('ParentSubagentsProvider', () => { expect(refetch).toHaveBeenCalledTimes(1); expect(result).toEqual(refreshed); }); + + it('marks an active parent run for prompt discovery', () => { + mockUseParentSubagentsQuery.mockReturnValue({ data: undefined, refetch: jest.fn() }); + render( + +
+ , + ); + expect(mockUseParentSubagentsQuery).toHaveBeenCalledWith( + 'parent-conversation', + { enabled: true }, + true, + ); + }); }); diff --git a/client/src/components/Chat/Subagents/ParentSubagentsProvider.tsx b/client/src/components/Chat/Subagents/ParentSubagentsProvider.tsx index 82ec322f94f..d04c75f252e 100644 --- a/client/src/components/Chat/Subagents/ParentSubagentsProvider.tsx +++ b/client/src/components/Chat/Subagents/ParentSubagentsProvider.tsx @@ -21,13 +21,15 @@ const ParentSubagentsContext = createContext(defaul export function ParentSubagentsProvider({ conversationId, enabled, + isSubmitting = false, children, }: { conversationId: string; enabled: boolean; + isSubmitting?: boolean; children: React.ReactNode; }) { - const { data, refetch } = useParentSubagentsQuery(conversationId, { enabled }); + const { data, refetch } = useParentSubagentsQuery(conversationId, { enabled }, isSubmitting); const refresh = useCallback(async () => { const result = await refetch(); return result.data; diff --git a/client/src/components/Chat/__tests__/ChatView.spec.tsx b/client/src/components/Chat/__tests__/ChatView.spec.tsx index 80d03eada04..e44ebf33020 100644 --- a/client/src/components/Chat/__tests__/ChatView.spec.tsx +++ b/client/src/components/Chat/__tests__/ChatView.spec.tsx @@ -135,4 +135,13 @@ describe('ChatView composer column', () => { expect(composerColumn).not.toHaveClass('overflow-y-auto'); expect(composerColumn).not.toHaveClass('scrollbar-gutter-stable'); }); + + test('layers composer overlays above positioned tool glyphs in the message column', () => { + const { container } = render(); + + const composerColumn = container.querySelector('.scrollbar-gutter-spacer'); + + expect(composerColumn).toHaveClass('[view-transition-name:chat-form]'); + expect(composerColumn).toHaveClass('relative', 'z-10'); + }); }); diff --git a/client/src/components/Chat/__tests__/Header.stacking.spec.tsx b/client/src/components/Chat/__tests__/Header.stacking.spec.tsx new file mode 100644 index 00000000000..d573da116bd --- /dev/null +++ b/client/src/components/Chat/__tests__/Header.stacking.spec.tsx @@ -0,0 +1,48 @@ +import React from 'react'; +import '@testing-library/jest-dom'; +import { render } from '@testing-library/react'; +import Header from '../Header'; + +jest.mock('react-router-dom', () => ({ + useParams: () => ({ conversationId: 'convo-1' }), +})); +jest.mock('recoil', () => ({ useRecoilValue: () => false })); +jest.mock('librechat-data-provider', () => ({ + getConfigDefaults: () => ({ interface: {} }), + Constants: { NEW_CONVO: 'new' }, + PermissionTypes: { BOOKMARKS: 'bookmarks', MULTI_CONVO: 'multi_convo', TEMPORARY_CHAT: 'temp' }, + Permissions: { USE: 'use' }, +})); +jest.mock('~/data-provider', () => ({ useGetStartupConfig: () => ({ data: undefined }) })); +jest.mock('~/hooks', () => ({ useHasAccess: () => false })); +jest.mock('~/store', () => ({ + __esModule: true, + default: { sidebarExpanded: {}, isSubmittingFamily: () => ({}) }, +})); +jest.mock('~/utils', () => ({ + cn: (...classes: Array) => classes.filter(Boolean).join(' '), +})); +jest.mock('../Menus', () => ({ + OpenSidebar: () => null, + PresetsMenu: () => null, + NewChat: () => null, + HeaderMenu: () => null, +})); +jest.mock('../TemporaryChat', () => ({ + TemporaryChat: () => null, + TemporaryChatIndicator: () => null, +})); +jest.mock('../Trace', () => ({ useTraceControl: () => ({ show: false }) })); +jest.mock('../Menus/Endpoints/ModelSelector', () => () => null); +jest.mock('../ExportAndShareMenu', () => () => null); +jest.mock('../SubagentThreadLink', () => () => null); +jest.mock('../Menus/BookmarkMenu', () => () => null); +jest.mock('../AddMultiConvo', () => () => null); + +describe('Header stacking', () => { + test('keeps header controls above the z-10 composer approval review', () => { + const { container } = render(
); + + expect(container.firstElementChild).toHaveClass('absolute', 'top-0', 'z-20'); + }); +}); diff --git a/client/src/data-provider/Agents/queries.ts b/client/src/data-provider/Agents/queries.ts index d51328aa252..0e6c1d06c7c 100644 --- a/client/src/data-provider/Agents/queries.ts +++ b/client/src/data-provider/Agents/queries.ts @@ -6,6 +6,7 @@ import type { UseInfiniteQueryOptions, } from '@tanstack/react-query'; import type t from 'librechat-data-provider'; +import { retryTransientQuery } from '../retry'; import { isEphemeralAgent } from '~/common'; /** @@ -74,8 +75,13 @@ export const useListAgentsQuery = ( params: t.AgentListParams = defaultAgentParams, config?: UseQueryOptions, ): QueryObserverResult => { - const queryClient = useQueryClient(); - const endpointsConfig = queryClient.getQueryData([QueryKeys.endpoints]); + /** The shell owns fetching endpoints. Observe its query, but do not start a second + * request or couple this shared agent hook to the shell's Recoil gate. */ + const { data: endpointsConfig } = useQuery( + [QueryKeys.endpoints], + () => dataService.getAIEndpoints(), + { enabled: false }, + ); const enabled = !!endpointsConfig?.[EModelEndpoint.agents]; return useQuery( @@ -83,10 +89,10 @@ export const useListAgentsQuery = ( () => fetchAllAgentPages(params), { staleTime: 1000 * 5, - refetchOnWindowFocus: false, - refetchOnReconnect: false, - refetchOnMount: false, - retry: false, + refetchOnWindowFocus: true, + refetchOnReconnect: true, + refetchOnMount: true, + retry: retryTransientQuery, ...config, enabled: config?.enabled !== undefined ? config.enabled && enabled : enabled, }, @@ -109,10 +115,11 @@ export const useGetAgentByIdQuery = ( agent_id: agent_id as string, }), { - refetchOnWindowFocus: false, - refetchOnReconnect: false, - refetchOnMount: false, - retry: false, + staleTime: 1000 * 5, + refetchOnWindowFocus: true, + refetchOnReconnect: true, + refetchOnMount: true, + retry: retryTransientQuery, ...config, enabled: isValidAgentId && (config?.enabled ?? true), }, diff --git a/client/src/data-provider/CodeEnvironments.ts b/client/src/data-provider/CodeEnvironments.ts index 6dad49dab8e..3213accea33 100644 --- a/client/src/data-provider/CodeEnvironments.ts +++ b/client/src/data-provider/CodeEnvironments.ts @@ -1,5 +1,11 @@ -import { useMutation, useQueries, useQuery, useQueryClient } from '@tanstack/react-query'; import { DynamicQueryKeys, MutationKeys, QueryKeys, dataService } from 'librechat-data-provider'; +import { + useIsFetching, + useMutation, + useQueries, + useQuery, + useQueryClient, +} from '@tanstack/react-query'; import type { TConversation, CodeEnvironmentUserSettings, @@ -9,7 +15,9 @@ import type { TCodeEnvironmentStatusResponse, TCodeEnvironmentsResponse, } from 'librechat-data-provider'; +import type { QueryFilters } from '@tanstack/react-query'; import { CONVERSATION_LIST_KEYS, updateConvoInAllQueries } from '~/utils'; +import { retryTransientQuery } from './retry'; export type CodeEnvironmentPairingResponse = TCodeEnvironmentPairingResponse; @@ -21,30 +29,41 @@ export function useCodeEnvironmentsQuery(enabled = true) { ); } +function codeEnvironmentStatusOptions(id: string, enabled: boolean) { + return { + queryKey: DynamicQueryKeys.codeEnvironmentStatus(id), + queryFn: (): Promise => + dataService.getCodeEnvironmentStatus(id), + enabled: enabled && id.length > 0, + staleTime: 10_000, + refetchInterval: 10_000, + refetchIntervalInBackground: false, + retry: retryTransientQuery, + }; +} + export function useCodeEnvironmentStatusQuery(id: string, enabled = true) { - return useQuery( - DynamicQueryKeys.codeEnvironmentStatus(id), - () => dataService.getCodeEnvironmentStatus(id), - { - enabled: enabled && id.length > 0, - refetchInterval: 10_000, - refetchIntervalInBackground: false, - retry: false, - }, - ); + return useQuery(codeEnvironmentStatusOptions(id, enabled)); } export function useCodeEnvironmentStatusQueries(ids: string[], enabled = true) { - return useQueries({ - queries: ids.map((id) => ({ - queryKey: DynamicQueryKeys.codeEnvironmentStatus(id), - queryFn: () => dataService.getCodeEnvironmentStatus(id), - enabled: enabled && id.length > 0, - refetchInterval: 10_000, - refetchIntervalInBackground: false, - retry: false, - })), - }); + return useQueries({ queries: ids.map((id) => codeEnvironmentStatusOptions(id, enabled)) }); +} + +/** Refresh discovery, not the conversation draft or expanded agent editor records. */ +const isWorkspaceDiscovery: NonNullable = ({ queryKey }) => + queryKey[0] === QueryKeys.endpoints || + queryKey[0] === QueryKeys.agents || + queryKey[0] === QueryKeys.codeEnvironments || + (queryKey[0] === QueryKeys.agent && queryKey.length === 2); + +export function useCodeWorkspaceRefresh() { + const queryClient = useQueryClient(); + const isRefreshing = useIsFetching({ predicate: isWorkspaceDiscovery }) > 0; + return { + isRefreshing, + refresh: () => queryClient.invalidateQueries({ predicate: isWorkspaceDiscovery }), + }; } export function usePairCodeEnvironmentMutation() { diff --git a/client/src/data-provider/Subagents/queries.polling.test.tsx b/client/src/data-provider/Subagents/queries.polling.test.tsx new file mode 100644 index 00000000000..13f49ae0578 --- /dev/null +++ b/client/src/data-provider/Subagents/queries.polling.test.tsx @@ -0,0 +1,120 @@ +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 { ParentSubagentIndex } from 'librechat-data-provider'; +import { useParentSubagentsQuery } from './queries'; + +const mockGetParentSubagents = jest.fn(); + +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + getParentSubagents: (...args: unknown[]) => mockGetParentSubagents(...args), + }, + }; +}); + +const index: ParentSubagentIndex = { + parentConversationId: 'parent', + children: [], + childrenTruncated: false, +}; + +function setup(isSubmitting = false) { + const client = new QueryClient(); + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ); + const hook = renderHook(() => useParentSubagentsQuery('parent', undefined, isSubmitting), { + wrapper, + }); + return { ...hook, client }; +} + +describe('parent subagent discovery polling', () => { + beforeEach(() => { + jest.useFakeTimers(); + mockGetParentSubagents.mockReset(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('starts at the existing cadence and settles at one read per minute when empty', async () => { + mockGetParentSubagents.mockResolvedValue(index); + const { client, unmount } = setup(); + await act(async () => jest.advanceTimersByTimeAsync(0)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(1); + await act(async () => jest.advanceTimersByTimeAsync(10_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(2); + await act(async () => jest.advanceTimersByTimeAsync(59_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(2); + await act(async () => jest.advanceTimersByTimeAsync(1_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(3); + unmount(); + client.clear(); + }); + + it('wakes immediately on child activity and resumes active-child refresh', async () => { + mockGetParentSubagents + .mockResolvedValueOnce(index) + .mockResolvedValueOnce(index) + .mockResolvedValue({ ...index, children: [{ status: 'running' }] }); + const { client, unmount } = setup(); + await act(async () => jest.advanceTimersByTimeAsync(10_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(2); + + await act(async () => { + await client.invalidateQueries([QueryKeys.parentSubagents, 'parent']); + }); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(3); + await act(async () => jest.advanceTimersByTimeAsync(2_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(4); + unmount(); + client.clear(); + }); + + it('does not retry or poll a missing parent on an idle tab', async () => { + jest.spyOn(console, 'error').mockImplementation(() => undefined); + mockGetParentSubagents.mockRejectedValue({ response: { status: 404 } }); + const { client, unmount } = setup(); + await act(async () => jest.advanceTimersByTimeAsync(0)); + await act(async () => jest.advanceTimersByTimeAsync(120_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(1); + unmount(); + client.clear(); + }); + + it('recovers when an initially missing parent persists during the first run', async () => { + jest.spyOn(console, 'error').mockImplementation(() => undefined); + mockGetParentSubagents.mockRejectedValueOnce({ response: { status: 404 } }); + mockGetParentSubagents.mockResolvedValue(index); + const { client, unmount, result } = setup(true); + await act(async () => jest.advanceTimersByTimeAsync(0)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(1); + await act(async () => jest.advanceTimersByTimeAsync(10_000)); + await act(async () => jest.advanceTimersByTimeAsync(1)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(2); + expect(result.current.data?.children).toEqual([]); + unmount(); + client.clear(); + }); + + it('retries a new parent during a running first turn but stops after readiness expires', async () => { + jest.spyOn(console, 'error').mockImplementation(() => undefined); + mockGetParentSubagents.mockRejectedValue({ response: { status: 404 } }); + const { client, unmount } = setup(true); + await act(async () => jest.advanceTimersByTimeAsync(0)); + await act(async () => jest.advanceTimersByTimeAsync(10_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(2); + await act(async () => jest.advanceTimersByTimeAsync(120_000)); + expect(mockGetParentSubagents).toHaveBeenCalledTimes(7); + unmount(); + client.clear(); + }); +}); diff --git a/client/src/data-provider/Subagents/queries.test.ts b/client/src/data-provider/Subagents/queries.test.ts index ddfc1ba5561..9076a471615 100644 --- a/client/src/data-provider/Subagents/queries.test.ts +++ b/client/src/data-provider/Subagents/queries.test.ts @@ -120,36 +120,60 @@ describe('subagent thread refresh policy', () => { expect(refetch).not.toHaveBeenCalled(); }); - it('refreshes active parent children quickly and discovers idle actors at a bounded cadence', () => { - expect( - parentSubagentsRefetchInterval({ - parentConversationId: 'parent-conversation', - childrenTruncated: false, - children: [{ status: 'running' }], - } as ParentSubagentIndex), - ).toBe(2_000); - expect( - parentSubagentsRefetchInterval({ + it('retains fast live refresh but sparsifies quiet parent discovery', () => { + const parent = (status?: ParentSubagentIndex['children'][number]['status']) => + ({ parentConversationId: 'parent-conversation', childrenTruncated: false, - children: [{ status: 'dispatched' }], - } as ParentSubagentIndex), - ).toBe(10_000); + children: status == null ? [] : [{ status }], + }) as ParentSubagentIndex; + + expect(parentSubagentsRefetchInterval(parent('running'))).toBe(2_000); + expect(parentSubagentsRefetchInterval(parent('dispatched'))).toBe(10_000); + expect(parentSubagentsRefetchInterval(parent(), false, null, 1)).toBe(10_000); + expect(parentSubagentsRefetchInterval(parent(), false, null, 2)).toBe(60_000); + expect(parentSubagentsRefetchInterval(parent('completed'), false, null, 3)).toBe(60_000); + expect(parentSubagentsRefetchInterval(parent(), true, null, 20)).toBe(10_000); expect(parentSubagentsRefetchInterval(undefined)).toBe(10_000); + }); + + it('stops idle 404 polling and gives a running first turn bounded readiness retries', () => { + const missing = { response: { status: 404 } }; + expect(parentSubagentsRefetchInterval(undefined, false, missing, 0, Date.now() + 60_000)).toBe( + false, + ); + expect(parentSubagentsRefetchInterval(undefined, true, missing, 0, Date.now() + 60_000)).toBe( + 10_000, + ); + expect(parentSubagentsRefetchInterval(undefined, true, missing, 0, Date.now())).toBe(false); + expect(parentSubagentsRefetchInterval(undefined, false, { response: { status: 503 } })).toBe( + 60_000, + ); + }); + it('keeps idle fallback and suppresses automatic 404 retries in the live query', () => { mockUseQuery.mockReturnValue({ data: undefined, error: null, refetch: jest.fn() }); renderHook(() => useParentSubagentsQuery('parent-conversation')); expect(mockUseQuery.mock.calls.at(-1)?.[0]).toEqual(['parentSubagents', 'parent-conversation']); - expect(mockUseQuery.mock.calls.at(-1)?.[2]).toEqual( + const options = mockUseQuery.mock.calls.at(-1)?.[2]; + expect(options).toEqual( expect.objectContaining({ enabled: true, refetchOnWindowFocus: true, - refetchInterval: parentSubagentsRefetchInterval, refetchIntervalInBackground: false, staleTime: 5_000, }), ); + expect(options.retry(0, { response: { status: 404 } })).toBe(false); + expect(options.retry(0, { response: { status: 500 } })).toBe(true); + expect(options.retry(3, { response: { status: 500 } })).toBe(false); + expect( + options.refetchInterval(undefined, { state: { dataUpdateCount: 0, error: { status: 404 } } }), + ).toBe(false); + expect( + options.refetchInterval({ children: [] }, { state: { dataUpdateCount: 2, error: null } }), + ).toBe(60_000); }); }); diff --git a/client/src/data-provider/Subagents/queries.ts b/client/src/data-provider/Subagents/queries.ts index 290b1e2f869..be9202727ae 100644 --- a/client/src/data-provider/Subagents/queries.ts +++ b/client/src/data-provider/Subagents/queries.ts @@ -10,20 +10,53 @@ import type { } from 'librechat-data-provider'; import type { UseQueryOptions, QueryObserverResult } from '@tanstack/react-query'; +const responseStatus = (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; +}; + export const ACTIVE_THREAD_REFRESH_MS = 2_000; const IDLE_PARENT_REFRESH_MS = 10_000; +const QUIET_PARENT_REFRESH_MS = 60_000; const CHILD_READY_POLL_WINDOW_MS = 60_000; -export const parentSubagentsRefetchInterval = (index: ParentSubagentIndex | undefined): number => - index?.children.some((child) => child.status === 'running') === true - ? ACTIVE_THREAD_REFRESH_MS - : IDLE_PARENT_REFRESH_MS; +/** Streamed child lifecycle events wake local discovery; the capped fallback finds + * children published elsewhere when this tab has no active parent stream. */ +export const parentSubagentsRefetchInterval = ( + index: ParentSubagentIndex | undefined, + isSubmitting = false, + error?: unknown, + dataUpdateCount = 0, + readinessDeadline = 0, +): number | false => { + if (responseStatus(error) === 404) { + return isSubmitting && Date.now() < readinessDeadline ? IDLE_PARENT_REFRESH_MS : false; + } + if (index?.children.some((child) => child.status === 'running') === true) { + return ACTIVE_THREAD_REFRESH_MS; + } + if (isSubmitting || index?.children.some((child) => child.status === 'dispatched')) { + return IDLE_PARENT_REFRESH_MS; + } + if (index == null && error == null) { + return IDLE_PARENT_REFRESH_MS; + } + // Recheck the first empty result soon: publication may still be committing. + return dataUpdateCount <= 1 && error == null ? IDLE_PARENT_REFRESH_MS : QUIET_PARENT_REFRESH_MS; +}; export const useParentSubagentsQuery = ( parentConversationId: string, config?: UseQueryOptions, -) => - useQuery( + isSubmitting = false, +) => { + const readinessKey = `${parentConversationId}\u0000${isSubmitting}`; + const readiness = useMemo( + () => ({ key: readinessKey, deadline: Date.now() + CHILD_READY_POLL_WINDOW_MS }), + [readinessKey], + ); + return useQuery( [QueryKeys.parentSubagents, parentConversationId], () => dataService.getParentSubagents(parentConversationId), { @@ -33,11 +66,20 @@ export const useParentSubagentsQuery = ( parentConversationId !== Constants.PENDING_CONVO, staleTime: 5_000, refetchOnWindowFocus: true, - refetchInterval: parentSubagentsRefetchInterval, + retry: (failureCount, error) => responseStatus(error) !== 404 && failureCount < 3, + refetchInterval: (index, query) => + parentSubagentsRefetchInterval( + index, + isSubmitting, + query.state.error, + query.state.dataUpdateCount, + readiness.deadline, + ), refetchIntervalInBackground: false, ...config, }, ); +}; const isTerminal = (status: SubagentThreadView['status']): boolean => status === 'completed' || @@ -80,12 +122,6 @@ export const subagentThreadRefetchInterval = ( return isTerminal(view.status) ? false : ACTIVE_THREAD_REFRESH_MS; }; -const responseStatus = (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; -}; - export const isSubagentReadinessPending = ( error: unknown, readinessDeadline: number, diff --git a/client/src/data-provider/__tests__/moveCodeEnvironment.test.tsx b/client/src/data-provider/__tests__/moveCodeEnvironment.test.tsx index 2fbc1eb17df..f56b71a5181 100644 --- a/client/src/data-provider/__tests__/moveCodeEnvironment.test.tsx +++ b/client/src/data-provider/__tests__/moveCodeEnvironment.test.tsx @@ -18,50 +18,57 @@ const mac = { environmentId: 'mac', workspaceId: 'primary' }; const vm = { environmentId: 'vm', workspaceId: 'projects' }; describe('useMoveConversationCodeEnvironmentMutation', () => { - it('keeps the moved decision when a conversation read from before the move resolves later', async () => { - const queryClient = new QueryClient({ - defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, - }); - const stored = { - conversationId: 'convo-1', - codeEnvironmentMode: 'attached', - codeWorkspaces: [mac], - } as TConversation; - queryClient.setQueryData([QueryKeys.conversation, 'convo-1'], stored); - jest.mocked(dataService.moveConversationCodeEnvironment).mockResolvedValue({ - conversationId: 'convo-1', - codeEnvironmentMode: 'attached', - codeWorkspaces: [vm], - }); - let resolveStaleRead: (conversation: TConversation) => void = () => undefined; - const wrapper = ({ children }: { children: ReactNode }) => - createElement(QueryClientProvider, { client: queryClient }, children); + it.each([vm, { ...mac, workspaceId: 'replacement' }])( + 'keeps the moved decision %j when a pre-move conversation read resolves later', + async (target) => { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + const stored = { + conversationId: 'convo-1', + codeEnvironmentMode: 'attached', + codeWorkspaces: [mac], + } as TConversation; + queryClient.setQueryData([QueryKeys.conversation, 'convo-1'], stored); + jest.mocked(dataService.moveConversationCodeEnvironment).mockResolvedValue({ + conversationId: 'convo-1', + codeEnvironmentMode: 'attached', + codeWorkspaces: [target], + }); + let resolveStaleRead: (conversation: TConversation) => void = () => undefined; + const wrapper = ({ children }: { children: ReactNode }) => + createElement(QueryClientProvider, { client: queryClient }, children); - const { result } = renderHook( - () => ({ - conversation: useQuery( - [QueryKeys.conversation, 'convo-1'], - () => - new Promise((resolve) => { - resolveStaleRead = resolve; - }), - { staleTime: 0 }, - ), - move: useMoveConversationCodeEnvironmentMutation(), - }), - { wrapper }, - ); - await waitFor(() => expect(result.current.conversation.isFetching).toBe(true)); + const { result } = renderHook( + () => ({ + conversation: useQuery( + [QueryKeys.conversation, 'convo-1'], + () => + new Promise((resolve) => { + resolveStaleRead = resolve; + }), + { staleTime: 0 }, + ), + move: useMoveConversationCodeEnvironmentMutation(), + }), + { wrapper }, + ); + await waitFor(() => expect(result.current.conversation.isFetching).toBe(true)); - await act(async () => { - await result.current.move.mutateAsync({ conversationId: 'convo-1', from: [mac], to: [vm] }); - }); - await act(async () => { - resolveStaleRead(stored); - }); + await act(async () => { + await result.current.move.mutateAsync({ + conversationId: 'convo-1', + from: [mac], + to: [target], + }); + }); + await act(async () => { + resolveStaleRead(stored); + }); - expect(queryClient.getQueryData([QueryKeys.conversation, 'convo-1'])).toEqual( - expect.objectContaining({ codeEnvironmentMode: 'attached', codeWorkspaces: [vm] }), - ); - }); + expect(queryClient.getQueryData([QueryKeys.conversation, 'convo-1'])).toEqual( + expect.objectContaining({ codeEnvironmentMode: 'attached', codeWorkspaces: [target] }), + ); + }, + ); }); diff --git a/client/src/data-provider/__tests__/workspaceFreshness.spec.tsx b/client/src/data-provider/__tests__/workspaceFreshness.spec.tsx new file mode 100644 index 00000000000..a59af1dabcf --- /dev/null +++ b/client/src/data-provider/__tests__/workspaceFreshness.spec.tsx @@ -0,0 +1,220 @@ +import { createElement } from 'react'; +import { act, renderHook, waitFor } from '@testing-library/react'; +import { QueryClient, QueryClientProvider, focusManager, useQueries } from '@tanstack/react-query'; +import { + dataService, + DynamicQueryKeys, + EModelEndpoint, + PermissionBits, + QueryKeys, +} from 'librechat-data-provider'; +import type { + Agent, + AgentListResponse, + TCodeEnvironmentStatusResponse, +} from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import { + useCodeWorkspaceRefresh, + useCodeEnvironmentStatusQueries, + useCodeEnvironmentStatusQuery, +} from '../CodeEnvironments'; +import { useGetAgentByIdQuery, useListAgentsQuery } from '../Agents/queries'; + +jest.mock('librechat-data-provider', () => { + const actual = + jest.requireActual('librechat-data-provider'); + return { ...actual, dataService: { ...actual.dataService } }; +}); + +const ready: TCodeEnvironmentStatusResponse = { + environmentId: 'vm', + status: 'ready', + workspaces: [{ id: 'project' }], +}; +const agent = { + id: 'agent_primary', + code_environment_id: 'vm', + code_workspace_id: 'project', +} as Agent; +const agentList = { object: 'list', data: [agent], has_more: false } as AgentListResponse; + +function setup() { + const client = new QueryClient({ + defaultOptions: { queries: { retryDelay: 0, cacheTime: Infinity } }, + }); + client.setQueryData([QueryKeys.endpoints], { [EModelEndpoint.agents]: {} }); + const wrapper = ({ children }: { children: ReactNode }) => + createElement(QueryClientProvider, { client }, children); + return { client, wrapper }; +} + +describe('workspace query lifecycle', () => { + afterEach(() => focusManager.setFocused(undefined)); + + it.each(['single', 'multiple'] as const)( + 'retries a transient %s status request without a reload', + async (kind) => { + const request = jest + .spyOn(dataService, 'getCodeEnvironmentStatus') + .mockRejectedValueOnce({ status: 503 }) + .mockResolvedValue(ready); + const { client, wrapper } = setup(); + const { result, unmount } = renderHook( + () => + kind === 'single' + ? useCodeEnvironmentStatusQuery('vm') + : useCodeEnvironmentStatusQueries(['vm'])[0], + { wrapper }, + ); + await waitFor(() => expect(result.current.data).toEqual(ready)); + expect(request).toHaveBeenCalledTimes(2); + unmount(); + client.clear(); + }, + ); + + it.each([401, 403, 404, 429])('does not retry status HTTP %s', async (status) => { + const request = jest + .spyOn(dataService, 'getCodeEnvironmentStatus') + .mockRejectedValue({ status }); + const { client, wrapper } = setup(); + const { result, unmount } = renderHook(() => useCodeEnvironmentStatusQuery('vm'), { wrapper }); + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(request).toHaveBeenCalledTimes(1); + unmount(); + client.clear(); + }); + + it('bounds transient retries instead of fetching indefinitely', async () => { + const request = jest + .spyOn(dataService, 'getCodeEnvironmentStatus') + .mockRejectedValue({ status: 503 }); + const { client, wrapper } = setup(); + const { result, unmount } = renderHook(() => useCodeEnvironmentStatusQuery('vm'), { wrapper }); + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(request).toHaveBeenCalledTimes(3); + unmount(); + client.clear(); + }); + + it('refreshes discovery but not saved conversation or expanded editor snapshots', async () => { + const { client, wrapper } = setup(); + const keys = [ + [QueryKeys.endpoints], + [QueryKeys.agents], + [QueryKeys.agent, agent.id], + DynamicQueryKeys.codeEnvironmentStatus('vm'), + [QueryKeys.agent, agent.id, 'expanded'], + [QueryKeys.conversation, 'saved-chat'], + ]; + const requests = keys.map(() => jest.fn().mockResolvedValue({ loaded: true })); + const { result, unmount } = renderHook( + () => { + useQueries({ + queries: keys.map((queryKey, i) => ({ + queryKey, + queryFn: requests[i], + staleTime: Infinity, + })), + }); + return useCodeWorkspaceRefresh(); + }, + { wrapper }, + ); + await waitFor(() => expect(result.current.isRefreshing).toBe(false)); + const before = requests.map((request) => request.mock.calls.length); + await act(async () => { + await result.current.refresh(); + }); + requests.forEach((request, i) => { + expect(request).toHaveBeenCalledTimes(before[i] + (i < 4 ? 1 : 0)); + }); + unmount(); + client.clear(); + }); + + it('shares a fresh status between composer, submit hook and settings observers', async () => { + const request = jest.spyOn(dataService, 'getCodeEnvironmentStatus').mockResolvedValue(ready); + const { client, wrapper } = setup(); + const first = renderHook(() => useCodeEnvironmentStatusQuery('vm'), { wrapper }); + await waitFor(() => expect(first.result.current.isSuccess).toBe(true)); + const second = renderHook(() => useCodeEnvironmentStatusQueries(['vm']), { wrapper }); + await act(async () => {}); + expect(second.result.current[0].data).toEqual(ready); + expect(request).toHaveBeenCalledTimes(1); + first.unmount(); + second.unmount(); + client.clear(); + }); + + it('revalidates stale agent defaults on mount and focus', async () => { + const request = jest.spyOn(dataService, 'getAgentById').mockResolvedValue(agent); + const { client, wrapper } = setup(); + client.setQueryData( + [QueryKeys.agent, agent.id], + { ...agent, code_workspace_id: 'old' }, + { updatedAt: 1 }, + ); + const { result, unmount } = renderHook(() => useGetAgentByIdQuery(agent.id), { wrapper }); + await waitFor(() => expect(result.current.data?.code_workspace_id).toBe('project')); + act(() => { + client.setQueryData( + [QueryKeys.agent, agent.id], + { ...agent, code_workspace_id: 'old' }, + { updatedAt: 1 }, + ); + focusManager.setFocused(false); + focusManager.setFocused(true); + }); + await waitFor(() => expect(request).toHaveBeenCalledTimes(2)); + unmount(); + client.clear(); + }); + + it('starts agent discovery from cached endpoint configuration without Recoil or a duplicate request', async () => { + const request = jest.spyOn(dataService, 'listAgents').mockResolvedValue(agentList); + const getEndpoints = jest.spyOn(dataService, 'getAIEndpoints'); + const { client, wrapper } = setup(); + client.removeQueries([QueryKeys.endpoints]); + const { result, unmount } = renderHook(() => useListAgentsQuery(), { wrapper }); + expect(request).not.toHaveBeenCalled(); + act(() => { + client.setQueryData([QueryKeys.endpoints], { [EModelEndpoint.agents]: {} }); + }); + await waitFor(() => expect(result.current.data?.data).toEqual([agent])); + expect(request).toHaveBeenCalledTimes(1); + expect(getEndpoints).not.toHaveBeenCalled(); + unmount(); + client.clear(); + }); + + it('revalidates stale reachable-agent metadata on mount', async () => { + const request = jest.spyOn(dataService, 'listAgents').mockResolvedValue(agentList); + const { client, wrapper } = setup(); + const params = { requiredPermission: PermissionBits.VIEW, limit: 100 }; + client.setQueryData([QueryKeys.agents, params], { ...agentList, data: [] }, { updatedAt: 1 }); + const { result, unmount } = renderHook(() => useListAgentsQuery(params), { wrapper }); + await waitFor(() => expect(result.current.data?.data).toEqual([agent])); + expect(request).toHaveBeenCalledTimes(1); + unmount(); + client.clear(); + }); + + it('retains a previous status while a background request retries', async () => { + const request = jest + .spyOn(dataService, 'getCodeEnvironmentStatus') + .mockRejectedValueOnce(new Error('Network interrupted')) + .mockResolvedValue(ready); + const { client, wrapper } = setup(); + client.setQueryData(DynamicQueryKeys.codeEnvironmentStatus('vm'), ready, { updatedAt: 1 }); + const { result, unmount } = renderHook(() => useCodeEnvironmentStatusQuery('vm'), { wrapper }); + expect(result.current.data).toEqual(ready); + expect(result.current.isFetching).toBe(true); + await waitFor(() => expect(request).toHaveBeenCalledTimes(2)); + await waitFor(() => expect(result.current.isFetching).toBe(false)); + expect(result.current.isSuccess).toBe(true); + unmount(); + client.clear(); + }); +}); diff --git a/client/src/data-provider/retry.ts b/client/src/data-provider/retry.ts new file mode 100644 index 00000000000..cc7e2657e96 --- /dev/null +++ b/client/src/data-provider/retry.ts @@ -0,0 +1,7 @@ +import { getResponseStatus } from '~/utils/errors'; + +/** Retry transport and server failures, never an authorization or resource decision. */ +export function retryTransientQuery(failureCount: number, error: unknown): boolean { + const status = getResponseStatus(error); + return failureCount < 2 && (status == null || status >= 500); +} diff --git a/client/src/hooks/Agents/__tests__/useCodeWorkspace.test.tsx b/client/src/hooks/Agents/__tests__/useCodeWorkspace.test.tsx index b7ce0bf74a9..c2a98c1d129 100644 --- a/client/src/hooks/Agents/__tests__/useCodeWorkspace.test.tsx +++ b/client/src/hooks/Agents/__tests__/useCodeWorkspace.test.tsx @@ -99,6 +99,14 @@ describe('useCodeWorkspace', () => { }, ); + it('reports unavailable when a ready worker advertises no workspaces', () => { + mockStatus()[0].data.workspaces = []; + const { result } = renderHook(() => useCodeWorkspace(conversation())); + expect(result.current.state).toBe('unavailable'); + expect(result.current.canSubmit).toBe(true); + expect(result.current.resolveSubmission(undefined, 'attached')).toBeUndefined(); + }); + it('selects one unambiguous initial workspace', () => { const { result } = renderHook(() => useCodeWorkspace(conversation())); @@ -627,7 +635,7 @@ describe('useCodeWorkspace', () => { expect(result.current.relocation).toBeUndefined(); }); - it('does not offer a move when the sealed machine lost its workspace', () => { + it('does not offer recovery to an API that only supports environment moves', () => { const removed = { environmentId: 'personal-vm', workspaceId: 'removed-project' }; const { result } = renderHook(() => useCodeWorkspace(sealed([removed]))); @@ -636,6 +644,120 @@ describe('useCodeWorkspace', () => { expect(result.current.relocation).toBeUndefined(); }); + describe('missing workspace recovery', () => { + const missing = { environmentId: 'personal-vm', workspaceId: 'deleted-project' }; + const replacement = { environmentId: 'personal-vm', workspaceId: 'project-a' }; + + beforeEach(() => { + mockStartupConfig.mockReturnValue({ + codeEnvironmentMoveVersion: 1, + codeWorkspaceRecoveryVersion: 1, + }); + }); + + it.each(['attached', undefined] as const)( + 'offers explicit recovery without silently selecting a replacement for mode %s', + (codeEnvironmentMode) => { + const { result, rerender } = renderHook( + ({ codeWorkspaces }) => + useCodeWorkspace({ ...sealed(codeWorkspaces), codeEnvironmentMode }), + { initialProps: { codeWorkspaces: [missing] } }, + ); + + expect(result.current.state).toBe('relocatable'); + expect(result.current.canSubmit).toBe(false); + expect(result.current.selections).toBeUndefined(); + expect(result.current.resolveSubmission([missing], 'attached')).toBeUndefined(); + expect(result.current.relocation).toEqual({ + conversationId: 'existing', + from: [missing], + previous: [], + retained: [], + targets: [expect.objectContaining({ state: 'missing', selected: undefined })], + }); + + rerender({ codeWorkspaces: [replacement] }); + expect(result.current.state).toBe('ready'); + expect(result.current.locked).toBe(true); + expect(result.current.canSubmit).toBe(true); + expect(result.current.relocation).toBeUndefined(); + }, + ); + + it.each([ + { codeEnvironmentMoveVersion: 1 }, + { codeEnvironmentMoveVersion: 1, codeWorkspaceRecoveryVersion: 2 }, + { codeWorkspaceRecoveryVersion: 1 }, + { codeEnvironmentMoveVersion: 2, codeWorkspaceRecoveryVersion: 1 }, + ])('never offers recovery without both supported capabilities: %j', (config) => { + mockStartupConfig.mockReturnValue(config); + const { result } = renderHook(() => useCodeWorkspace(sealed([missing]))); + expect(result.current.state).toBe('missing'); + expect(result.current.relocation).toBeUndefined(); + expect(result.current.canSubmit).toBe(false); + }); + + it('preserves ordinary environment moves on a recovery-capable API', () => { + const { result } = renderHook(() => useCodeWorkspace(sealed([mac]))); + expect(result.current.state).toBe('relocatable'); + expect(result.current.relocation?.targets[0].state).toBe('choose'); + }); + + it('offers the empty recovery picker without inventing a replacement', () => { + mockStatus()[0].data.workspaces = []; + const { result } = renderHook(() => useCodeWorkspace(sealed([missing]))); + expect(result.current.relocation?.targets[0].workspaces).toEqual([]); + expect(result.current.canSubmit).toBe(false); + }); + + it.each(['ready', 'loading', 'unavailable', 'unsupported', 'choose', 'missing'] as const)( + 'handles a second environment in state %s without dropping its sealed selection', + (state) => { + mockAgentPermissions().agent.subagents = { enabled: true, agent_ids: ['child'] }; + mockAgentsMap.mockReturnValue({ + child: { + id: 'child', + stateful_code_sessions: true, + code_environment_id: 'team-vm', + tools: [Tools.execute_code], + }, + }); + mockAgentsConfig().agentsConfig.statefulCodeSessions.environments.push({ + id: 'team-vm', + name: 'Team VM', + type: 'attached', + baseURL: 'https://team.example.com', + }); + const kept = { environmentId: 'team-vm', workspaceId: 'shared' }; + mockStatus().push({ + data: { + environmentId: 'team-vm', + status: state === 'unavailable' ? 'unavailable' : 'ready', + workspaces: + state === 'unsupported' + ? undefined + : [{ id: state === 'missing' ? 'replacement' : 'shared' }], + }, + isLoading: state === 'loading', + isError: false, + }); + const from = state === 'choose' ? [missing] : [missing, kept]; + const { result } = renderHook(() => useCodeWorkspace(sealed(from))); + + expect(result.current.canSubmit).toBe(false); + if (['loading', 'unavailable', 'unsupported'].includes(state)) { + expect(result.current.relocation).toBeUndefined(); + return; + } + expect(result.current.relocation?.from).toEqual(from); + expect(result.current.relocation?.retained).toEqual(state === 'ready' ? [kept] : []); + expect( + result.current.relocation?.targets.map(({ environment }) => environment.id), + ).toEqual(state === 'ready' ? ['personal-vm'] : ['personal-vm', 'team-vm']); + }, + ); + }); + it('does not offer an attached machine to a chat that continues without one', () => { const { result } = renderHook(() => useCodeWorkspace({ diff --git a/client/src/hooks/Agents/__tests__/workspaceLifecycle.spec.tsx b/client/src/hooks/Agents/__tests__/workspaceLifecycle.spec.tsx new file mode 100644 index 00000000000..5f4bf838d51 --- /dev/null +++ b/client/src/hooks/Agents/__tests__/workspaceLifecycle.spec.tsx @@ -0,0 +1,124 @@ +import { createElement } from 'react'; +import { act, renderHook, waitFor } from '@testing-library/react'; +import { dataService, DynamicQueryKeys } from 'librechat-data-provider'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import type { TCodeEnvironmentStatusResponse, TConversation } from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import useCodeWorkspace from '../useCodeWorkspace'; + +jest.mock('librechat-data-provider', () => { + const actual = + jest.requireActual('librechat-data-provider'); + return { ...actual, dataService: { ...actual.dataService } }; +}); +jest.mock('~/data-provider', () => ({ + useCodeEnvironmentStatusQueries: jest.requireActual('~/data-provider/CodeEnvironments') + .useCodeEnvironmentStatusQueries, + useGetStartupConfig: () => ({ data: { codeEnvironmentDecisionVersion: 1 } }), +})); +jest.mock('../workspacePreferences', () => ({ + useWorkspacePreferences: () => ({ get: () => undefined, remember: jest.fn() }), +})); +jest.mock('~/hooks/Roles/useHasAccess', () => () => true); +jest.mock('~/Providers', () => ({ useAgentsMapContext: () => ({}) })); +jest.mock('../useAgentToolPermissions', () => (id: string) => ({ + agent: id + ? { id, stateful_code_sessions: true, code_environment_id: 'vm', tools: ['execute_code'] } + : undefined, +})); +jest.mock('../useGetAgentsConfig', () => () => ({ + agentsConfig: { + capabilities: ['execute_code', 'stateful_code_sessions'], + statefulCodeSessions: { environments: [{ id: 'vm', name: 'VM', type: 'attached' }] }, + }, +})); + +const selected = { environmentId: 'vm', workspaceId: 'project' }; +const ready: TCodeEnvironmentStatusResponse = { + environmentId: 'vm', + status: 'ready', + workspaces: [{ id: 'project' }], +}; +const saved = { + conversationId: 'saved-chat', + endpoint: 'agents', + agent_id: 'agent_primary', + codeEnvironmentMode: 'attached', + codeWorkspaces: [selected], +} as TConversation; + +function setup(conversation: TConversation = saved) { + const client = new QueryClient({ + defaultOptions: { queries: { retryDelay: 0, cacheTime: Infinity } }, + }); + const wrapper = ({ children }: { children: ReactNode }) => + createElement(QueryClientProvider, { client }, children); + return { client, ...renderHook(() => useCodeWorkspace(conversation), { wrapper }) }; +} + +describe('workspace status recovery', () => { + it('restores the saved selection when a delayed initial status request resolves', async () => { + let resolve = (_status: TCodeEnvironmentStatusResponse) => {}; + jest.spyOn(dataService, 'getCodeEnvironmentStatus').mockReturnValue( + new Promise((done) => { + resolve = done; + }), + ); + const { result, client, unmount } = setup(); + expect(result.current.state).toBe('loading'); + expect(result.current.canSubmit).toBe(false); + await act(async () => resolve(ready)); + await waitFor(() => expect(result.current.state).toBe('ready')); + expect(result.current.selections).toEqual([selected]); + expect(result.current.locked).toBe(true); + expect(result.current.resolveSubmission(saved.codeWorkspaces, 'attached')).toEqual({ + codeEnvironmentMode: 'attached', + codeWorkspaces: [selected], + }); + unmount(); + client.clear(); + }); + + it('recovers an exhausted background failure without replacing the saved selection', async () => { + const request = jest.spyOn(dataService, 'getCodeEnvironmentStatus').mockResolvedValue(ready); + const { result, client, unmount } = setup(); + await waitFor(() => expect(result.current.state).toBe('ready')); + request.mockRejectedValue(new Error('offline')); + await act(async () => { + await client.invalidateQueries(DynamicQueryKeys.codeEnvironmentStatus('vm')); + }); + await waitFor(() => expect(result.current.state).toBe('unavailable')); + expect(result.current.canSubmit).toBe(false); + expect(result.current.environments[0].selected).toEqual(selected); + request.mockResolvedValue(ready); + await act(async () => { + await client.invalidateQueries(DynamicQueryKeys.codeEnvironmentStatus('vm')); + }); + await waitFor(() => expect(result.current.state).toBe('ready')); + expect(result.current.selections).toEqual([selected]); + unmount(); + client.clear(); + }); + + it('recovers an empty first response into the sole workspace without a reload', async () => { + const request = jest + .spyOn(dataService, 'getCodeEnvironmentStatus') + .mockResolvedValue({ ...ready, workspaces: [] }); + const { result, client, unmount } = setup({ + ...saved, + conversationId: 'new', + codeEnvironmentMode: undefined, + codeWorkspaces: undefined, + }); + await waitFor(() => expect(result.current.state).toBe('unavailable')); + expect(result.current.resolveSubmission(undefined, 'attached')).toBeUndefined(); + request.mockResolvedValue(ready); + await act(async () => { + await client.invalidateQueries(DynamicQueryKeys.codeEnvironmentStatus('vm')); + }); + await waitFor(() => expect(result.current.state).toBe('ready')); + expect(result.current.selections).toEqual([selected]); + unmount(); + client.clear(); + }); +}); diff --git a/client/src/hooks/Agents/useCodeWorkspace.ts b/client/src/hooks/Agents/useCodeWorkspace.ts index a104c5f44a8..67c7c8f5a3e 100644 --- a/client/src/hooks/Agents/useCodeWorkspace.ts +++ b/client/src/hooks/Agents/useCodeWorkspace.ts @@ -9,6 +9,7 @@ import { AgentCapabilities, CODE_ENVIRONMENT_DECISION_VERSION, CODE_ENVIRONMENT_MOVE_VERSION, + CODE_WORKSPACE_RECOVERY_VERSION, PermissionTypes, Permissions, } from 'librechat-data-provider'; @@ -48,9 +49,8 @@ export interface CodeWorkspaceEnvironmentResult { } /** - * A saved chat whose attached decision no longer covers every environment its agents use, most - * often because an agent was pointed at a different machine after the chat was created. The - * decision stays sealed against implicit changes; only its owner's explicit move replaces it. + * A saved chat whose attached decision no longer covers its agents' environments or whose + * selected workspace is missing. Only its owner's explicit move replaces the sealed decision. */ export interface CodeWorkspaceRelocation { conversationId: string; @@ -60,9 +60,9 @@ export interface CodeWorkspaceRelocation { previous: Array< Pick & Partial> >; - /** Sealed selections the agents still use; a move carries them over unchanged. */ + /** Registered sealed selections the agents still use; a move carries them over unchanged. */ retained: CodeWorkspaceSelection[]; - /** Environments the agents now use that the decision does not cover. */ + /** Environments needing a new selection, including those whose workspace disappeared. */ targets: CodeWorkspaceEnvironmentResult[]; } @@ -135,6 +135,9 @@ export default function useCodeWorkspace( startupConfig?.codeEnvironmentDecisionVersion === CODE_ENVIRONMENT_DECISION_VERSION; const supportsEnvironmentMoves = startupConfig?.codeEnvironmentMoveVersion === CODE_ENVIRONMENT_MOVE_VERSION; + const supportsWorkspaceRecovery = + supportsEnvironmentMoves && + startupConfig?.codeWorkspaceRecoveryVersion === CODE_WORKSPACE_RECOVERY_VERSION; const preferences = useWorkspacePreferences(conversation?.agent_id); const { agentsConfig, endpointsConfig } = useGetAgentsConfig(); const canRunCode = useHasAccess({ @@ -305,6 +308,7 @@ export default function useCodeWorkspace( } else if (status.data.workspaces == null) state = 'unsupported'; else if (selected != null) state = 'ready'; else if (stored != null) state = 'missing'; + else if (workspaces.length === 0) state = 'unavailable'; return { environment, state, workspaces, selected }; }); @@ -374,7 +378,8 @@ export default function useCodeWorkspace( if ( supportsEnvironmentMoves && locked && - state === 'choose' && + (state === 'choose' || (supportsWorkspaceRecovery && state === 'missing')) && + environmentResults.every(({ state }) => ['ready', 'choose', 'missing'].includes(state)) && inferredMode === 'attached' && conversation?.conversationId != null && storedSelections != null && @@ -393,7 +398,9 @@ export default function useCodeWorkspace( retained: environmentResults.flatMap((result) => result.state === 'ready' && result.selected != null ? [result.selected] : [], ), - targets: environmentResults.filter((result) => result.state === 'choose'), + targets: environmentResults.filter( + (result) => result.state === 'choose' || result.state === 'missing', + ), }; state = 'relocatable'; } diff --git a/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx b/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx index d14d0b6a7e6..40246872fa5 100644 --- a/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useQueueDrain.spec.tsx @@ -10,6 +10,7 @@ import type { QueuedMessage, SettledQueuedTurnReceipt, } from '~/store/families'; +import { recoveryDispositionsFamily } from '~/components/Chat/Steering/recovery'; import { revealedQueuedTurnFamily } from '~/store/steer'; import useQueueDrain from '../useQueueDrain'; import store from '~/store'; @@ -101,9 +102,83 @@ const runEnd = (overrides: Partial = {}): RunEnd => ({ describe('useQueueDrain', () => { beforeEach(() => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), {}); mockMarkFilesUsage.mockClear(); }); + it('moves a held new-conversation head without spending its terminal boundary', async () => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: 'blocked' }); + const held: QueuedMessage = { + id: 'held', + text: 'review before sending', + createdAt: 1, + recoverySteerId: 'source', + }; + const next: QueuedMessage = { id: 'ordinary', text: 'send after review', createdAt: 2 }; + const { ask, setters } = setup(({ set }) => { + set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [held]); + set(store.queuedMessagesByConvoId(CONVO_ID), [next]); + set(store.isSubmittingFamily(INDEX), false); + }); + const end = runEnd({ startedAsNewConvo: true }); + act(() => setters.setRunEnd!(end)); + await waitFor(() => expect(setters.newConvoQueue).toEqual([])); + expect(setters.queue).toEqual([held, next]); + expect(setters.runEnd).toEqual(end); + expect(ask).not.toHaveBeenCalled(); + + act(() => setters.setQueue!([next])); + await waitFor(() => expect(ask).toHaveBeenCalledTimes(1)); + expect(ask).toHaveBeenCalledWith({ text: next.text }, emptyOverrides); + expect(setters.queue).toEqual([]); + expect(setters.runEnd).toBeNull(); + }); + + it('does not turn an aborted run into a send after a held row is dismissed', async () => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: 'blocked' }); + const held: QueuedMessage = { + id: 'held', + text: 'review me', + createdAt: 1, + recoverySteerId: 'source', + }; + const next: QueuedMessage = { id: 'ordinary', text: 'wait for next run', createdAt: 2 }; + const { ask, setters } = setup(({ set }) => { + set(store.queuedMessagesByConvoId(CONVO_ID), [held, next]); + set(store.isSubmittingFamily(INDEX), false); + }); + act(() => setters.setRunEnd!(runEnd({ outcome: 'aborted' }))); + await waitFor(() => expect(setters.runEnd).toBeNull()); + act(() => setters.setQueue!([next])); + expect(ask).not.toHaveBeenCalled(); + expect(setters.queue).toEqual([next]); + }); + + it.each(['blocked', 'cancelling', 'cancelled', 'dismissed'] as const)( + 'does not drain a %s recovery while its run-end boundary remains available', + async (disposition) => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: disposition }); + const item = { + id: 'leftover', + text: 'original words', + createdAt: 1, + recoverySteerId: 'source', + clientRequestId: 'same-attempt', + }; + const firstEnd = runEnd(); + const { ask, setters } = setup(({ set }) => { + set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + set(store.isSubmittingFamily(INDEX), false); + set(store.runEndByIndex(INDEX), firstEnd); + }); + await waitFor(() => expect(setters.runEnd).toEqual(firstEnd)); + act(() => setters.setRunEnd?.(runEnd({ generationCreatedAt: 42 }))); + expect(setters.runEnd).toEqual(firstEnd); + expect(ask).not.toHaveBeenCalled(); + expect(setters.queue).toEqual([item]); + }, + ); + it('drains exactly one queued message on clean completion', async () => { const { ask, setters } = setup(({ set }) => { set(store.queuedMessagesByConvoId(CONVO_ID), [ diff --git a/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx b/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx index 98adf2d9244..e94880b4b99 100644 --- a/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useSteerConvert.spec.tsx @@ -1,6 +1,8 @@ import React from 'react'; +import { getDefaultStore } from 'jotai'; import { act, renderHook } from '@testing-library/react'; import { RecoilRoot, useRecoilValue, useSetRecoilState, type MutableSnapshot } from 'recoil'; +import { recoveryDispositionsFamily } from '~/components/Chat/Steering/recovery'; import useSteerConvert from '../useSteerConvert'; import store from '~/store'; @@ -48,6 +50,7 @@ function setup(initialize?: (snapshot: MutableSnapshot) => void) { describe('useSteerConvert', () => { beforeEach(() => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), {}); mockFileMap = {}; }); @@ -173,6 +176,40 @@ describe('useSteerConvert', () => { expect(result.current.applied).toContain('server-replacement-id'); }); + it.each(['cancelled', 'dismissed'] as const)( + 'does not resurrect a %s recovery after remount or redelivery', + (disposition) => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: disposition }); + const { result } = setup(); + act(() => + result.current.convert(CONVO_ID, [{ steerId: 'source', text: 'old words' }], { + generationProtocolVersion: 2, + allowPreviouslyConvertedIds: ['source'], + }), + ); + expect(result.current.queue).toEqual([]); + }, + ); + + it('does not make a held recovery sendable when a legacy claim arrives', () => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: 'blocked' }); + const item = { + id: 'leftover', + text: 'original words', + recoverySteerId: 'source', + clientRequestId: 'attempt', + createdAt: 1, + }; + const { result } = setup(({ set }) => set(store.queuedMessagesByConvoId(CONVO_ID), [item])); + act(() => + result.current.convert(CONVO_ID, [{ steerId: 'source', text: 'original words' }], { + generationProtocolVersion: 1, + }), + ); + expect(result.current.queue).toContainEqual(item); + expect(result.current.queue.every((row) => row.recoverySteerId === 'source')).toBe(true); + }); + it('recovers a v1 leftover as an ordinary local follow-up without receipt binding', () => { const { result } = setup(); act(() => { diff --git a/client/src/hooks/Chat/__tests__/useSteering.spec.tsx b/client/src/hooks/Chat/__tests__/useSteering.spec.tsx index 2e62cbe58dd..6f63a875633 100644 --- a/client/src/hooks/Chat/__tests__/useSteering.spec.tsx +++ b/client/src/hooks/Chat/__tests__/useSteering.spec.tsx @@ -6,6 +6,7 @@ import { Constants, ContentTypes, EModelEndpoint, LocalStorageKeys } from 'libre import type { TConversation, TFile, TMessage } from 'librechat-data-provider'; import type { QueuedMessage } from '~/store/families'; import { clearAllDrafts, getPendingDraftId, getNewConversationDraftId } from '~/utils'; +import { recoveryDispositionsFamily } from '~/components/Chat/Steering/recovery'; import useSteering, { mergeQueuedTurnFileMetadata } from '../useSteering'; import { revealedQueuedTurnFamily } from '~/store/steer'; import useQueueDrain from '../useQueueDrain'; @@ -145,6 +146,7 @@ function useQueue(convoId: string) { describe('useSteering', () => { beforeEach(() => { + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), {}); jest.clearAllMocks(); getDefaultStore().set(revealedQueuedTurnFamily(CONVO_ID), null); mockMessages = undefined; @@ -153,6 +155,175 @@ describe('useSteering', () => { mockFileMap = {}; }); + it.each([true, false])( + 'reserves a recovery against queue drain while cancelling (removed=%s)', + async (removed) => { + const item: QueuedMessage = { + id: 'leftover', + text: 'original words', + createdAt: 1, + recoverySteerId: 'source', + recoveryClientSteerId: 'client-source', + clientRequestId: 'attempt', + }; + let resolveCancel: (value: { removed: boolean }) => void = () => undefined; + mockCancelSteer.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveCancel = resolve; + }), + ); + const ask = jest.fn(); + const wrapper = ({ children }: { children: React.ReactNode }) => ( + { + set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + set(store.isSubmittingFamily(0), false); + })} + > + {children} + + ); + const { result } = renderHook( + () => { + const steering = useSteering({ + consumeDraft: jest.fn(), + index: 0, + conversationId: CONVO_ID, + conversation: agentsConversation, + isSubmitting: false, + answerModeActive: false, + sendNow: jest.fn(), + stopGenerating: jest.fn(), + }); + useQueueDrain(0, CONVO_ID, ask); + return { + steering, + queue: useQueue(CONVO_ID), + setEnd: useSetRecoilState(store.runEndByIndex(0)), + }; + }, + { wrapper }, + ); + let cancellation: Promise = Promise.resolve(false); + act(() => { + cancellation = result.current.steering.discardQueued(item); + }); + act(() => + result.current.setEnd({ + conversationId: CONVO_ID, + outcome: 'completed', + endedAt: 200, + generationCreatedAt: 41, + }), + ); + expect(ask).not.toHaveBeenCalled(); + expect(result.current.queue).toEqual([item]); + let discarded = false; + await act(async () => { + resolveCancel({ removed }); + discarded = await cancellation; + }); + expect(discarded).toBe(removed); + expect(ask).not.toHaveBeenCalled(); + expect(getDefaultStore().get(recoveryDispositionsFamily(CONVO_ID))).toEqual({ + source: removed ? 'cancelled' : 'blocked', + }); + if (removed) { + expect(result.current.queue).toEqual([item]); + } else { + expect(result.current.queue).toEqual([item]); + act(() => + result.current.setEnd({ + conversationId: CONVO_ID, + outcome: 'completed', + endedAt: 300, + generationCreatedAt: 42, + }), + ); + expect(ask).not.toHaveBeenCalled(); + act(() => result.current.steering.dismissRecovery(item)); + expect(result.current.queue).toEqual([]); + expect(getDefaultStore().get(recoveryDispositionsFamily(CONVO_ID))).toEqual({ + source: 'dismissed', + }); + } + }, + ); + + it.each(['blocked', 'cancelled'] as const)( + 'releases exactly one completed run-end after dismissing a %s recovery', + async (disposition) => { + const held: QueuedMessage = { + id: 'held', + text: 'might already be delivered', + createdAt: 1, + recoverySteerId: 'source', + clientRequestId: 'recovery-attempt', + }; + const first: QueuedMessage = { id: 'next', text: 'send after dismissal', createdAt: 2 }; + const second: QueuedMessage = { id: 'later', text: 'wait for next run', createdAt: 3 }; + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: disposition }); + const ask = jest.fn(); + const wrapper = ({ children }: { children: React.ReactNode }) => ( + { + set(store.queuedMessagesByConvoId(CONVO_ID), [held, first, second]); + set(store.isSubmittingFamily(0), false); + })} + > + {children} + + ); + const { result } = renderHook( + () => { + const steering = useSteering({ + consumeDraft: jest.fn(), + index: 0, + conversationId: CONVO_ID, + conversation: agentsConversation, + isSubmitting: false, + answerModeActive: false, + sendNow: jest.fn(), + stopGenerating: jest.fn(), + }); + useQueueDrain(0, CONVO_ID, ask); + return { + steering, + queue: useQueue(CONVO_ID), + end: useRecoilValue(store.runEndByIndex(0)), + setEnd: useSetRecoilState(store.runEndByIndex(0)), + }; + }, + { wrapper }, + ); + const terminal = { + conversationId: CONVO_ID, + outcome: 'completed' as const, + endedAt: 200, + generationCreatedAt: 41, + }; + act(() => result.current.setEnd(terminal)); + await waitFor(() => expect(result.current.end).toEqual(terminal)); + expect(ask).not.toHaveBeenCalled(); + expect(result.current.queue).toEqual([held, first, second]); + + act(() => result.current.steering.dismissRecovery(held)); + await waitFor(() => expect(ask).toHaveBeenCalledTimes(1)); + expect(ask).toHaveBeenCalledWith( + { text: first.text }, + expect.objectContaining({ overrideQueuedMessageOrigin: expect.any(Object) }), + ); + expect(result.current.end).toBeNull(); + expect(result.current.queue).toEqual([second]); + expect(getDefaultStore().get(recoveryDispositionsFamily(CONVO_ID))).toEqual({ + source: 'dismissed', + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(ask).toHaveBeenCalledTimes(1); + }, + ); + it('keeps optimistic delivery metadata when a legacy receipt omits it', () => { expect( mergeQueuedTurnFileMetadata( @@ -2778,7 +2949,7 @@ describe('useSteering', () => { expect(mockMutate).not.toHaveBeenCalled(); }); - it('atomically discards a recovered source and downgrades its row in place', async () => { + it('confirms a discarded recovery but keeps it held until the guarded UI action succeeds', async () => { mockCancelSteer.mockResolvedValueOnce({ removed: true, generationProtocolVersion: 2, @@ -2827,26 +2998,12 @@ describe('useSteering', () => { steerId: 'server-leftover', clientSteerId: 'client-leftover', }); - expect(result.current.queue).toEqual([ - before, - { - id: 'queued-leftover', - text: 'edit this next', - createdAt: 1, - expectedPredecessorCreatedAt: 41, - files: [ - { - file_id: 'file-1', - filepath: '/uploads/file-1.txt', - type: 'text/plain', - }, - ], - quotes: ['keep quote'], - manualSkills: ['keep skill'], - priority: true, - }, - after, - ]); + expect(result.current.queue).toEqual([before, recovered, after]); + expect(getDefaultStore().get(recoveryDispositionsFamily(CONVO_ID))).toEqual({ + 'server-leftover': 'cancelled', + }); + act(() => result.current.steering.sendQueuedNow(recovered)); + expect(mockMutate).not.toHaveBeenCalled(); }); it('keeps a recovered queue row when its parked source cannot be discarded', async () => { @@ -2899,6 +3056,22 @@ describe('useSteering', () => { expect(result.current.queue).toEqual([recovered]); }); + it('refuses a direct manual send of an unreconciled recovery', () => { + const item = { + id: 'held', + text: 'may already be delivered', + createdAt: 1, + recoverySteerId: 'source', + }; + getDefaultStore().set(recoveryDispositionsFamily(CONVO_ID), { source: 'blocked' }); + const { result, sendNow } = setupWithState({ isSubmitting: false }, ({ set }) => { + set(store.queuedMessagesByConvoId(CONVO_ID), [item]); + }); + act(() => result.current.steering.sendQueuedNow(item)); + expect(sendNow).not.toHaveBeenCalled(); + expect(result.current.queue).toEqual([item]); + }); + it('sendQueuedNow steers whenever steering is available, even under the queue preference', () => { const { result } = setupWithState({}, ({ set }) => { set(store.duringRunDefaultAction, 'queue'); diff --git a/client/src/hooks/Chat/useQueueDrain.ts b/client/src/hooks/Chat/useQueueDrain.ts index b0e4c21fa55..48d3b4f1eec 100644 --- a/client/src/hooks/Chat/useQueueDrain.ts +++ b/client/src/hooks/Chat/useQueueDrain.ts @@ -1,9 +1,14 @@ import { useEffect, useMemo } from 'react'; -import { useAtomValue } from 'jotai'; +import { useAtomValue, useStore } from 'jotai'; import { Constants } from 'librechat-data-provider'; import { useRecoilValue, useRecoilCallback } from 'recoil'; import type { DrainAfterAbort, QueuedMessage, QueuedMessageOrigin, RunEnd } from '~/store/families'; import type { TAskFunction } from '~/common'; +import { + recoveryDispositionsFamily, + recoveryDisposition, + canRestoreRecovery, +} from '~/components/Chat/Steering/recovery'; import { selectQueuedTurnReveal } from '~/hooks/Chat/useQueuedTurnReveal'; import { useMarkFilesUsageMutation } from '~/data-provider'; import { revealedQueuedTurnFamily } from '~/store/steer'; @@ -80,6 +85,7 @@ export default function useQueueDrain( ask: TAskFunction, revealQueuedTurn?: (item: QueuedMessage, end: RunEnd) => void, ) { + const jotaiStore = useStore(); const runEnd = useRecoilValue(store.runEndByIndex(index)); const parkedRunEnd = useRecoilValue( store.pendingRunEndByConvoId(activeConversationId ?? Constants.NEW_CONVO), @@ -101,6 +107,10 @@ export default function useQueueDrain( store.settledQueuedTurnReceiptsByConvoId(activeConversationId ?? Constants.NEW_CONVO), ); const hasServerOwnedQueue = [...ownQueue, ...newConvoQueue].some((item) => item.server != null); + // A held head can leave the queue without changing the terminal signal. + // Observe its identity so that dismissal wakes the parked boundary. + const ownHeadId = ownQueue[0]?.id ?? null; + const newConvoHeadId = newConvoQueue[0]?.id ?? null; /** The row the reveal would pick, and whether one is already revealed: a * revealed head that is cancelled or dies before admission leaves the * server-owned queue non-empty and its terminal evidence out of the @@ -334,11 +344,26 @@ export default function useQueueDrain( return reveal == null ? null : { kind: 'reveal', item: reveal, end }; } - // Consume only after server authority has yielded the boundary — a - // hard double-fire guard even if the effect re-runs before propagation. + const head = merged[0]; + const held = + head != null && + recoveryDisposition(jotaiStore.get(recoveryDispositionsFamily(conversationId)), head) != + null; + if (shouldDrain && held) { + // The held source cannot spend this completion. Keep its one-shot + // boundary so a successor can drain if the user dismisses the hold. + if (shouldMigrate && newConvoQueue.length > 0) { + set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), []); + set(store.queuedMessagesByConvoId(conversationId), merged); + } + return null; + } + + // Consume only after server authority and held recoveries yield the + // boundary. A later queue update cannot spend the same end twice. consumeEnd(); - const next = shouldDrain ? (merged[0] ?? null) : null; + const next = shouldDrain ? (head ?? null) : null; const remainder = next ? merged.slice(1) : merged; if (shouldMigrate && newConvoQueue.length > 0) { @@ -362,17 +387,20 @@ export default function useQueueDrain( } : null; }, - [index, activeConversationId], + [index, activeConversationId, jotaiStore], ); const restoreQueued = useRecoilCallback( ({ set }) => (convoId: string, item: QueuedMessage) => { set(store.queuedMessagesByConvoId(convoId), (prev) => - prev.some((queued) => queued.id === item.id) ? prev : [item, ...prev], + !canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(convoId)), item) || + prev.some((queued) => queued.id === item.id) + ? prev + : [item, ...prev], ); }, - [], + [jotaiStore], ); useEffect(() => { @@ -436,6 +464,8 @@ export default function useQueueDrain( }, [ runEnd, parkedRunEnd, + ownHeadId, + newConvoHeadId, isSubmitting, activeConversationId, parkForeignRunEnd, diff --git a/client/src/hooks/Chat/useSteerConvert.ts b/client/src/hooks/Chat/useSteerConvert.ts index 2cc711f2235..e3633543105 100644 --- a/client/src/hooks/Chat/useSteerConvert.ts +++ b/client/src/hooks/Chat/useSteerConvert.ts @@ -1,5 +1,6 @@ import { useCallback, useRef } from 'react'; import { v4 } from 'uuid'; +import { useStore } from 'jotai'; import { useRecoilCallback } from 'recoil'; import type { TPendingSteer } from 'librechat-data-provider'; import type { QueuedMessage, QueuedMessageOrigin } from '~/store/families'; @@ -11,6 +12,10 @@ import { insertQueuedOrigin, hydrateFileDeliveryMetadata, } from '~/utils'; +import { + recoveryDispositionsFamily, + canRestoreRecovery, +} from '~/components/Chat/Steering/recovery'; import { fetchStreamStatus, getGenerationProtocolVersion } from '~/data-provider'; import { useFileMapContext } from '~/Providers'; import store from '~/store'; @@ -53,6 +58,7 @@ interface SteerConvertOptions { * server-side removal. */ export default function useSteerConvert() { + const jotaiStore = useStore(); const fileMap = useFileMapContext(); const fileMapRef = useRef(fileMap); fileMapRef.current = fileMap; @@ -110,13 +116,19 @@ export default function useSteerConvert() { * already created a receipt-bound item before the claim reached an * old replica, that source no longer exists. Downgrade the existing * item in place to an ordinary local follow-up. */ - const existing = bindRecoverySource + const existing: QueuedMessage[] = bindRecoverySource ? prev : prev.map((item) => { const matchesClaimedSource = (item.recoverySteerId != null && steerIds.has(item.recoverySteerId)) || (item.recoveryClientSteerId != null && steerIds.has(item.recoveryClientSteerId)); - if (!matchesClaimedSource) { + if ( + !matchesClaimedSource || + (item.recoverySteerId != null && + jotaiStore.get(recoveryDispositionsFamily(conversationId))[ + item.recoverySteerId + ] != null) + ) { return item; } const { @@ -130,11 +142,14 @@ export default function useSteerConvert() { const fresh = steers .filter( (steer) => - allowedRedeliveries.has(steer.steerId) || - (!settledSteerIds.has(steer.steerId) && - (steer.clientSteerId == null || !settledSteerIds.has(steer.clientSteerId))), + canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(conversationId)), { + recoverySteerId: steer.steerId, + }) && + (allowedRedeliveries.has(steer.steerId) || + (!settledSteerIds.has(steer.steerId) && + (steer.clientSteerId == null || !settledSteerIds.has(steer.clientSteerId)))), ) - .map((steer) => { + .map((steer): { item: QueuedMessage; queuedOrigin?: QueuedMessageOrigin } => { const local = localChipFor(steer); const source = local ?? steer; const queuedOrigin = source.queuedOrigin; @@ -143,17 +158,20 @@ export default function useSteerConvert() { local?.files, fileMapRef.current, ); - const recoveryFields = bindRecoverySource - ? { - // One UUID is stable for this queued attempt and all of - // its POST retries. A later failed generation re-converts - // the durable source and receives a new key, so the old - // started idempotency tombstone cannot make it unsendable. - clientRequestId: v4(), - recoverySteerId: steer.steerId, - ...(steer.clientSteerId && { recoveryClientSteerId: steer.clientSteerId }), - } - : {}; + const held = + jotaiStore.get(recoveryDispositionsFamily(conversationId))[steer.steerId] != null; + const recoveryFields = + bindRecoverySource || held + ? { + // One UUID is stable for this queued attempt and all of + // its POST retries. A later failed generation re-converts + // the durable source and receives a new key, so the old + // started idempotency tombstone cannot make it unsendable. + clientRequestId: v4(), + recoverySteerId: steer.steerId, + ...(steer.clientSteerId && { recoveryClientSteerId: steer.clientSteerId }), + } + : {}; const item = queuedOrigin != null ? { ...queuedOrigin.item, ...recoveryFields, ...(files && { files }) } @@ -172,7 +190,15 @@ export default function useSteerConvert() { queuedOrigin: queuedOrigin != null ? { ...queuedOrigin, item } : undefined, }; }) - .filter(({ item }) => !existing.some((queued) => queued.id === item.id)); + .filter( + ({ item }) => + !existing.some( + (queued) => + queued.id === item.id || + (item.recoverySteerId != null && + queued.recoverySteerId === item.recoverySteerId), + ), + ); if (fresh.length === 0) { return existing; } @@ -193,7 +219,7 @@ export default function useSteerConvert() { return merged; }); }, - [], + [jotaiStore], ); return useCallback( diff --git a/client/src/hooks/Chat/useSteering.ts b/client/src/hooks/Chat/useSteering.ts index 02867fcb73c..afaa77b790d 100644 --- a/client/src/hooks/Chat/useSteering.ts +++ b/client/src/hooks/Chat/useSteering.ts @@ -42,6 +42,12 @@ import { hydrateFileDeliveryMetadata, mergeRestagedQuotes, } from '~/utils'; +import { + recoveryDispositionsFamily, + recoveryDisposition, + canRestoreRecovery, + blockRecovery, +} from '~/components/Chat/Steering/recovery'; import useSteerConvert from '~/hooks/Chat/useSteerConvert'; import { revealedQueuedTurnFamily } from '~/store/steer'; import { useLatestMessage } from '~/hooks/Messages'; @@ -1266,36 +1272,6 @@ export default function useSteering({ [queueKey], ); - /** Once a parked source is discarded it must never be retried as a recovery - * attempt. Downgrade the row in place so a guarded Edit that finds a newer - * draft can leave the same words, context, identity, and queue position as - * an ordinary local follow-up. */ - const downgradeQueuedRecovery = useRecoilCallback( - ({ snapshot, set }) => - (id: string): boolean => { - const queue = snapshot.getLoadable(store.queuedMessagesByConvoId(queueKey)).getValue(); - let found = false; - const next = queue.map((item) => { - if (item.id !== id) { - return item; - } - found = true; - const { - clientRequestId: _clientRequestId, - recoverySteerId: _recoverySteerId, - recoveryClientSteerId: _recoveryClientSteerId, - ...ordinary - } = item; - return ordinary; - }); - if (found) { - set(store.queuedMessagesByConvoId(queueKey), next); - } - return found; - }, - [queueKey], - ); - /** Settle a queued row's terminal recovery source before an Edit/Remove. * Ordinary rows have no server copy. A v2 leftover first uses its durable * receipt to atomically discard the parked copy, then becomes an ordinary @@ -1347,16 +1323,32 @@ export default function useSteering({ return true; } if (item.recoveryClientSteerId == null || !hasRealConvoId) { + jotaiStore.set(recoveryDispositionsFamily(queueKey), (previous) => + blockRecovery(previous, item.recoverySteerId!), + ); showToast({ message: localize('com_ui_steer_cancel_failed'), status: 'error', }); return false; } + const dispositions = recoveryDispositionsFamily(queueKey); + const disposition = recoveryDisposition(jotaiStore.get(dispositions), item); + if (disposition === 'cancelled') { + return true; + } + if (disposition === 'cancelling') { + return false; + } + if (!canRestoreRecovery(jotaiStore.get(dispositions), item)) { + return false; + } + const steerId = item.recoverySteerId; + jotaiStore.set(dispositions, (previous) => ({ ...previous, [steerId]: 'cancelling' })); try { const { removed } = await cancelSteer({ conversationId, - steerId: item.recoverySteerId, + steerId, clientSteerId: item.recoveryClientSteerId, }); if (removed !== true) { @@ -1366,13 +1358,20 @@ export default function useSteering({ }); return false; } - return downgradeQueuedRecovery(item.id); + jotaiStore.set(dispositions, (previous) => ({ ...previous, [steerId]: 'cancelled' })); + // Keep the binding held until the caller's guarded Edit/Remove succeeds. + // A newer composer draft must not turn cancelled words into an auto-send. + return true; } catch { showToast({ message: localize('com_ui_steer_cancel_failed'), status: 'error', }); return false; + } finally { + jotaiStore.set(dispositions, (previous) => + previous[steerId] === 'cancelling' ? { ...previous, [steerId]: 'blocked' } : previous, + ); } }, [ @@ -1380,7 +1379,6 @@ export default function useSteering({ cancelAgentQueuedTurn, applyQueuedTurnReceipts, conversationId, - downgradeQueuedRecovery, downgradeServerQueuedTurn, hasRealConvoId, localize, @@ -1390,6 +1388,24 @@ export default function useSteering({ ], ); + const dismissRecovery = useCallback( + (item: QueuedMessage) => { + const dispositions = recoveryDispositionsFamily(queueKey); + if ( + item.recoverySteerId == null || + !['blocked', 'cancelled'].includes( + recoveryDisposition(jotaiStore.get(dispositions), item) ?? '', + ) + ) { + return; + } + const steerId = item.recoverySteerId; + jotaiStore.set(dispositions, (previous) => ({ ...previous, [steerId]: 'dismissed' })); + removeQueued(item.id); + }, + [jotaiStore, queueKey, removeQueued], + ); + /** Capture-then-remove, including the item's neighbours, so any refused send * or rejected steer can restore the ORIGINAL item in place even if the run * drains an adjacent entry while the request is in flight. */ @@ -1456,9 +1472,13 @@ export default function useSteering({ ({ set }) => (origin: QueuedMessageOrigin) => { releaseQueuedOrigin(origin); - set(store.queuedMessagesByConvoId(queueKey), (prev) => insertQueuedOrigin(prev, origin)); + set(store.queuedMessagesByConvoId(queueKey), (prev) => + canRestoreRecovery(jotaiStore.get(recoveryDispositionsFamily(queueKey)), origin.item) + ? insertQueuedOrigin(prev, origin) + : prev, + ); }, - [queueKey, releaseQueuedOrigin], + [queueKey, releaseQueuedOrigin, jotaiStore], ); /** @@ -2003,6 +2023,9 @@ export default function useSteering({ * boundary; it only means something on the live-run path. */ const sendLocalQueuedNow = useCallback( (item: QueuedMessage, opts?: { preempt?: boolean }) => { + if (recoveryDisposition(jotaiStore.get(recoveryDispositionsFamily(queueKey)), item) != null) { + return; + } /** In answer mode (and any other submission-owned non-steerable state) * there is no immediate path. Refuse before touching queue state so a * stale/direct caller cannot perform the old remove-and-restore no-op. */ @@ -2071,6 +2094,8 @@ export default function useSteering({ }, [ takeQueued, + jotaiStore, + queueKey, duringRunActive, canSteer, submitSteer, @@ -2231,6 +2256,7 @@ export default function useSteering({ enqueue, removeQueued, discardQueued, + dismissRecovery, sendQueuedNow, interruptAndSend, interruptSteer, @@ -2257,6 +2283,7 @@ export default function useSteering({ enqueue, removeQueued, discardQueued, + dismissRecovery, sendQueuedNow, interruptAndSend, interruptSteer, diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index ad672d32381..d480fc74b17 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -10,6 +10,7 @@ import { request, } from 'librechat-data-provider'; import type { TMessage, TSubmission } from 'librechat-data-provider'; +import { recoveryDispositionsFamily } from '~/components/Chat/Steering/recovery'; import { pendingApprovalActionFamily } from '~/components/Chat/approval/state'; type SSEEventListener = (e: Partial & { responseCode?: number }) => void; @@ -359,6 +360,7 @@ const advanceRetryTimer = async (ms: number) => { describe('useResumableSSE', () => { beforeEach(() => { + getDefaultStore().set(recoveryDispositionsFamily(CONV_ID), {}); mockSSEInstances.length = 0; localStorage.clear(); mockErrorHandler.mockClear(); @@ -2965,6 +2967,58 @@ describe('useResumableSSE', () => { unmount(); }); + it.each(['RECOVERY_PAYLOAD_MISMATCH', 'INVALID_RECOVERY_REQUEST'])( + 'holds a rejected recovery (%s) without changing its words or binding', + async (code) => { + (request.post as jest.Mock).mockRejectedValueOnce({ + response: { status: 409, data: { code } }, + }); + const item = { + id: 'leftover', + text: 'original words', + createdAt: 1, + recoverySteerId: 'source', + recoveryClientSteerId: 'client-source', + clientRequestId: 'attempt', + quotes: ['original excerpt'], + files: [{ file_id: 'original-file' }], + }; + const submission = buildSubmission({ + recoverySteerId: 'source', + clientRequestId: 'attempt', + queuedMessageOrigin: { item, beforeIds: [], afterIds: [] }, + }); + const { unmount } = renderHook(() => useResumableSSE(submission, buildChatHelpers())); + await waitFor(() => expect(mockSetSubmission).toHaveBeenCalledWith(null)); + expect(getDefaultStore().get(recoveryDispositionsFamily(CONV_ID))).toEqual({ + source: 'blocked', + }); + expect(mockRestoreQueuedSubmission).toHaveBeenCalledWith(submission); + expect(mockConvertSteersToQueued).not.toHaveBeenCalled(); + expect(request.post).toHaveBeenCalledTimes(1); + unmount(); + }, + ); + + it('does not permanently hold a rate-limited recovery', async () => { + (request.post as jest.Mock).mockRejectedValueOnce({ response: { status: 429 } }); + const submission = buildSubmission({ recoverySteerId: 'source', clientRequestId: 'attempt' }); + const { unmount } = renderHook(() => useResumableSSE(submission, buildChatHelpers())); + await waitFor(() => expect(mockSetSubmission).toHaveBeenCalledWith(null)); + expect(getDefaultStore().get(recoveryDispositionsFamily(CONV_ID))).toEqual({}); + unmount(); + }); + + it('does not dispatch a recovery cancelled between dequeuing and startup', async () => { + getDefaultStore().set(recoveryDispositionsFamily(CONV_ID), { source: 'cancelled' }); + const submission = buildSubmission({ recoverySteerId: 'source', clientRequestId: 'attempt' }); + const { unmount } = renderHook(() => useResumableSSE(submission, buildChatHelpers())); + await waitFor(() => expect(mockSetSubmission).toHaveBeenCalledWith(null)); + expect(request.post).not.toHaveBeenCalled(); + expect(mockSSEInstances).toHaveLength(0); + unmount(); + }); + it('restores an exact queued row after a definitive pre-create rejection', async () => { (request.post as jest.Mock).mockRejectedValueOnce({ response: { status: 429, data: { message: 'too many requests' } }, diff --git a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts index 7e0df6cdaf1..4a712d10955 100644 --- a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts +++ b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts @@ -3240,6 +3240,47 @@ describe('useStepHandler', () => { ...overrides, }); + it('signals parent-index discovery on child lifecycle events, not every progress delta', () => { + const onSubagentIndexChange = jest.fn(); + const { result } = renderHook( + () => useStepHandler({ ...createHookParams(), onSubagentIndexChange }), + { wrapper: subagentStoreWrapper }, + ); + const submission = createSubmission(); + + act(() => { + for (const phase of ['start', 'run_step_delta', 'stop', 'error'] as const) { + result.current.stepHandler( + { event: StepEvents.ON_SUBAGENT_UPDATE, data: makeUpdate({ phase }) }, + submission, + ); + } + }); + expect(onSubagentIndexChange.mock.calls).toEqual([['conv-1'], ['conv-1'], ['conv-1']]); + + act(() => { + result.current.stepHandler( + { event: StepEvents.ON_SUBAGENT_UPDATE, data: makeUpdate({ phase: 'start' }) }, + createSubmission({ + userMessage: createUserMessage({ conversationId: String(Constants.NEW_CONVO) }), + initialResponse: createResponseMessage({ conversationId: String(Constants.NEW_CONVO) }), + }), + ); + }); + expect(onSubagentIndexChange).toHaveBeenCalledTimes(3); + + act(() => { + result.current.stepHandler( + { event: StepEvents.ON_SUBAGENT_UPDATE, data: makeUpdate({ phase: 'start' }) }, + createSubmission({ + userMessage: createUserMessage({ conversationId: String(Constants.NEW_CONVO) }), + initialResponse: createResponseMessage({ conversationId: 'saved-parent' }), + }), + ); + }); + expect(onSubagentIndexChange).toHaveBeenLastCalledWith('saved-parent'); + }); + it('correlates updates to a tool call via parentToolCallId (deterministic path)', () => { const { result, getProgress } = renderStepHandlerWithReader(); const { submission } = seedResponseWithSubagentToolCalls(result, ['call_A']); diff --git a/client/src/hooks/SSE/useEventHandlers.ts b/client/src/hooks/SSE/useEventHandlers.ts index c1d801b0780..7894ca89e9f 100644 --- a/client/src/hooks/SSE/useEventHandlers.ts +++ b/client/src/hooks/SSE/useEventHandlers.ts @@ -552,6 +552,12 @@ export default function useEventHandlers({ queryClient.invalidateQueries({ queryKey: [key], refetchType: 'all' }); } }, [queryClient]); + const onSubagentIndexChange = useCallback( + (conversationId: string) => { + void queryClient.invalidateQueries([QueryKeys.parentSubagents, conversationId]); + }, + [queryClient], + ); const { stepHandler, clearStepMaps, @@ -568,6 +574,7 @@ export default function useEventHandlers({ setIsSubmitting, lastAnnouncementTimeRef, onSkillAuthoringComplete, + onSubagentIndexChange, }); const attachmentHandler = useAttachmentHandler(queryClient); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index a944aab3b37..a2170f70f2a 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -82,6 +82,11 @@ import { supportsGenerationProtocolV2, GENERATION_PROTOCOL_VERSION, } from '~/data-provider'; +import { + recoveryDispositionsFamily, + canRestoreRecovery, + blockRecovery, +} from '~/components/Chat/Steering/recovery'; import useEventHandlers, { buildCreatedInitialResponse, keepLocalCodeApprovalMode, @@ -91,6 +96,7 @@ import useSteerConvert from '~/hooks/Chat/useSteerConvert'; import { useAuthContext } from '~/hooks/AuthContext'; import { useFileMapContext } from '~/Providers'; import useUsageHandler from './useUsageHandler'; +import useLocalize from '~/hooks/useLocalize'; import store from '~/store'; type ChatHelpers = Pick< @@ -824,6 +830,7 @@ export default function useResumableSSE( runIndex = 0, ) { const jotaiStore = useStore(); + const localize = useLocalize(); const queryClient = useQueryClient(); const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex)); @@ -948,10 +955,15 @@ export default function useResumableSSE( return; } set(store.queuedMessagesByConvoId(conversationId), (prev) => - insertQueuedOrigin(prev, origin, expectedPredecessorCreatedAt), + canRestoreRecovery( + jotaiStore.get(recoveryDispositionsFamily(conversationId)), + origin.item, + ) + ? insertQueuedOrigin(prev, origin, expectedPredecessorCreatedAt) + : prev, ); }, - [], + [jotaiStore], ); /** Removes the pending chip once its steer is injected (the inline content @@ -3941,6 +3953,23 @@ export default function useResumableSSE( const readinessDeadline = Date.now() + START_GENERATION_READINESS_TIMEOUT_MS; while (!signal?.aborted) { + const recoverySteerId = getRecoverySteerId(currentSubmission); + const conversationId = currentSubmission.conversation?.conversationId; + if ( + recoverySteerId != null && + conversationId && + jotaiStore.get(recoveryDispositionsFamily(conversationId))[recoverySteerId] != null + ) { + restoreQueuedSubmission(currentSubmission); + errorHandler({ + data: getStreamStartFailureData(localize('com_ui_steer_recovery_held')), + submission: currentSubmission as EventSubmission, + }); + setShowStopButton(false); + setIsSubmitting(false); + setSubmission(null); + return null; + } requestAttempts += 1; try { const data = await postGenerationRequest(url, payload, { signal }); @@ -4072,6 +4101,19 @@ export default function useResumableSSE( const errorData = startError?.response?.data; const responseStatus = startError?.response?.status; if (responseStatus != null && responseStatus >= 400 && responseStatus < 500) { + const recoverySteerId = getRecoverySteerId(currentSubmission); + const conversationId = currentSubmission.conversation?.conversationId; + const recoveryRejected = + errorData != null && + typeof errorData === 'object' && + 'code' in errorData && + (errorData.code === 'RECOVERY_PAYLOAD_MISMATCH' || + errorData.code === 'INVALID_RECOVERY_REQUEST'); + if (recoveryRejected && recoverySteerId != null && conversationId) { + jotaiStore.set(recoveryDispositionsFamily(conversationId), (previous) => + blockRecovery(previous, recoverySteerId), + ); + } // The server rejected admission before exposing a generation. Restore // the exact queue row/position; ambiguous transport/5xx outcomes must // first reconcile durable state instead of risking a duplicate start. @@ -4125,6 +4167,8 @@ export default function useResumableSSE( clearStepMaps, convertSteersToQueued, errorHandler, + jotaiStore, + localize, restoreQueuedSubmission, setIsSubmitting, setShowStopButton, diff --git a/client/src/hooks/SSE/useStepHandler.ts b/client/src/hooks/SSE/useStepHandler.ts index 7e70de85721..f97c5b603a3 100644 --- a/client/src/hooks/SSE/useStepHandler.ts +++ b/client/src/hooks/SSE/useStepHandler.ts @@ -57,6 +57,7 @@ type TUseStepHandler = { * invalidation) so this hook stays free of query-client coupling. */ onSkillAuthoringComplete?: () => void; + onSubagentIndexChange?: (conversationId: string) => void; }; type TStepEvent = @@ -159,6 +160,7 @@ export default function useStepHandler({ announcePolite, lastAnnouncementTimeRef, onSkillAuthoringComplete, + onSubagentIndexChange, }: TUseStepHandler) { const subagentStore = useStore(); const toolCallIdMap = useRef(new Map()); @@ -1376,6 +1378,18 @@ export default function useStepHandler({ responseMessageId = submission?.initialResponse?.messageId ?? ''; } applySubagentUpdate(stepEvent.data, responseMessageId); + if ( + stepEvent.data.phase === 'start' || + stepEvent.data.phase === 'stop' || + stepEvent.data.phase === 'error' + ) { + const conversationId = [ + submission?.userMessage?.conversationId, + submission?.initialResponse?.conversationId, + submission?.conversation?.conversationId, + ].find((id) => id && id !== Constants.NEW_CONVO && id !== Constants.PENDING_CONVO); + if (conversationId) onSubagentIndexChange?.(conversationId); + } } else if (stepEvent.event === StepEvents.ON_SUMMARIZE_START) { announcePolite({ message: 'summarize_started', isStatus: true }); } else if (stepEvent.event === StepEvents.ON_SUMMARIZE_DELTA) { @@ -1494,6 +1508,7 @@ export default function useStepHandler({ calculateContentIndex, getCurrentMessages, applySubagentUpdate, + onSubagentIndexChange, setSandboxStarting, clearSandboxStarting, applyPtcToolCall, diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index db1d8a3aedc..a99faaf2dd0 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -28,14 +28,16 @@ "com_ui_code_workspace_without_attached": "No attached workspace", "com_ui_code_workspace_without_attached_info": "Chat without using files or commands on an attached machine.", "com_ui_code_workspace_locked_recovery": "Reconnect or restore this workspace, or start a new chat.", + "com_ui_code_workspace_recover": "Recover workspace", + "com_ui_code_workspace_recover_info": "A workspace saved with this chat is no longer registered. Choose an available workspace for each environment below to continue this chat. Chat history is kept, but previous workspace files are not restored or copied.", "com_ui_code_workspace_move": "Move chat", "com_ui_code_workspace_move_to": "Move chat to {{0}}", "com_ui_code_workspace_move_info": "This chat used {{0}}, but its agent now uses {{1}}. Files from earlier turns stay on the previous machine.", "com_ui_code_workspace_move_info_removed": "This chat's agent no longer uses {{0}}. Move the chat to continue without it.", "com_ui_code_workspace_move_info_added": "This chat's agent now also uses {{0}}. Choose a workspace there to continue.", - "com_ui_code_workspace_move_stale": "This chat's workspace changed in another window. Reload and try again.", - "com_ui_code_workspace_move_busy": "Wait for the current response to finish, then move this chat.", - "com_ui_code_workspace_move_error": "Couldn't move this chat. Try again.", + "com_ui_code_workspace_move_stale": "This chat's workspace selection is no longer current. Reload and try again.", + "com_ui_code_workspace_move_busy": "Wait for the current response to finish before changing this chat's workspace.", + "com_ui_code_workspace_move_error": "Couldn't update this chat's workspace. Try again.", "com_ui_code_workspaces_selected": "{{0}} code workspaces", "com_ui_code_approval_ask": "Ask before changes", "com_ui_code_approval_ask_description": "Ask before editing files or running commands", @@ -442,7 +444,7 @@ "com_error_ban": "Your account is temporarily banned for violating the terms of service.", "com_error_code_workspace_invalid": "The workspace selection is invalid; select an available workspace.", "com_error_code_workspace_locked": "This chat already uses a different code environment. Continue with its original environment or start a new chat.", - "com_error_code_workspace_missing": "The selected workspace is no longer registered. Restore it on the attached machine or start a new chat.", + "com_error_code_workspace_missing": "The selected workspace is no longer registered. Restore it on the attached machine, choose a replacement if recovery is available, or start a new chat.", "com_error_code_workspace_required": "Choose an attached workspace before sending.", "com_error_code_workspace_unavailable": "The selected code workspace is unavailable. Reconnect its machine or choose an available workspace, then try again.", "com_error_code_workspace_unsupported": "The worker does not support workspace selection.", @@ -2474,6 +2476,11 @@ "com_ui_steer_arm_lost_race": "The interrupt could not be armed. The steering message may already have moved on or may still be waiting for the next tool step", "com_ui_steer_arm_unconfirmed": "Couldn't confirm whether the interrupt was armed. The steering message may interrupt or may wait for the next tool step", "com_ui_steer_cancel": "Cancel steering message", + "com_ui_steer_recovery_held": "Recovery paused", + "com_ui_steer_recovery_copy_refused": "Could not copy the recovery into the composer. Your draft or active chat may have changed; the recovery remains paused.", + "com_ui_steer_recovery_review": "Automatic recovery has stopped. This message may already have been delivered. Check the conversation before sending it again. Dismissing this recovery only removes it from this tab; it does not cancel server work.", + "com_ui_steer_copy_to_composer": "Copy to composer for review", + "com_ui_steer_dismiss_recovery": "Dismiss recovery in this tab", "com_ui_steer_cancel_failed": "Could not cancel the steering message — it may still reach the agent", "com_ui_steer_delivered": "Delivered", "com_ui_steer_delivered_info": "Delivered. Your message will join the response at the next safe moment.", diff --git a/client/src/routes/ChatRoute.tsx b/client/src/routes/ChatRoute.tsx index c0329ea7887..27e9f437a65 100644 --- a/client/src/routes/ChatRoute.tsx +++ b/client/src/routes/ChatRoute.tsx @@ -142,8 +142,8 @@ export default function ChatRoute() { const assistantListMap = useAssistantListMap(); /** The map comes from Root's shared context (one mapping pass app-wide); the * select-less observer only tracks settle state. Only a loaded list may - * invalidate a stored agent pick: on a transient catalog failure (retries are - * disabled) the map stays unknown, the pick stays trusted, and the gate below + * invalidate a stored agent pick: on a transient catalog failure after retries, + * the map stays unknown, the pick stays trusted, and the gate below * releases so the landing never hangs on the error. */ const agentsMap: TAgentsMap | undefined = useAgentsMapContext(); const agentsQuery = useListAgentsQuery( diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index fbd658c5fec..51db43d721c 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -698,10 +698,13 @@ function replyResponses(text) { /** * Attaches synthetic usage_metadata on a final empty chunk (the OpenAI * streaming pattern) so token-usage SSE events flow end to end in mock runs. + * Input is counted over the complete prompt a real provider bills — system + * instructions included — since the context snapshot calibrates against it. */ class UsageEmittingFakeChatModel extends FakeChatModel { - constructor({ resolveInvocation, resolveOnStream, sleep, ...options }) { + constructor({ graph, resolveInvocation, resolveOnStream, sleep, ...options }) { super({ ...options, sleep }); + this.graph = graph; this.resolveInvocation = resolveInvocation; this.resolveOnStream = resolveOnStream; this.streamSleep = sleep ?? CHUNK_DELAY_MS; @@ -818,7 +821,13 @@ class UsageEmittingFakeChatModel extends FakeChatModel { outputChars += typeof chunk.text === 'string' ? chunk.text.length : 0; yield chunk; } - const inputChars = (messages ?? []).reduce( + const { messages: promptMessages } = await getStreamAgentView({ + graph: this.graph, + messages: messages ?? [], + options, + runManager, + }); + const inputChars = promptMessages.reduce( (sum, message) => sum + getContentText(message?.content).length, 0, ); @@ -869,6 +878,7 @@ function overrideModel({ if (!thrownError) { const model = new UsageEmittingFakeChatModel({ + graph, responses, sleep: sleep ?? CHUNK_DELAY_MS, emitCustomEvent: true, @@ -2792,12 +2802,12 @@ function provisioningToolResponses({ text, toolNames }) { const command = Array.from({ length: 120 }, (_, index) => `printf 'line-${index}-☃\\n'`).join( '\n', ); - const args = - codeTool.name === 'bash_tool' - ? { command } - : codeTool.name === 'execute_code' - ? { lang: 'bash', code: command } - : codeTool.args; + let args = codeTool.args; + if (codeTool.name === 'bash_tool') { + args = { command }; + } else if (codeTool.name === 'execute_code') { + args = { lang: 'bash', code: command }; + } return { responses: ['', `E2E highlighted code complete: ${highlightLabel}`], sleep: highlightLabel === 'cancel' ? HIGHLIGHT_CANCEL_CHUNK_DELAY_MS : SLOW_CHUNK_DELAY_MS, diff --git a/librechat.example.yaml b/librechat.example.yaml index e407c300dd9..3f25483564b 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -1390,7 +1390,11 @@ endpoints: # messages: # unattributedAssistantContent: model_output # `model_output` (default) or `inspect` # pii: -# action: audit # `block` (default) or `audit` +# action: audit # `block` (default), `audit`, or `redact` (reserved for transformation-aware ingress; current routes still block) +# # `redact` requires the later encrypted owner-view and canonical-message integration. +# # maxCharacters: 65536 # bounded total inspected text per transformation session +# # maxMatches: 256 # bounded matches per transformation session +# # Custom patterns may set category: email, phone, name, credential, or custom. # fields: [name, text, summary, quote, answer, decision_response, decision_reason, content_part, attachment_reference, assembled_context] # starterPatterns: [sk_prefix, bearer_header, api_key_header] # customPatterns: diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index 84271fad946..43bac3fd1db 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -1742,6 +1742,35 @@ describe('initializeAgent — maxContextTokens', () => { jest.clearAllMocks(); }); + it.each(['us.openai.gpt-6-sol', 'global.openai.gpt-6-astra', 'us.openai.gpt-5.6-terra'])( + 'budgets %s using the Bedrock context window', + async (model) => { + const { agent, req, res, loadTools, db } = createMocks({ + provider: Providers.BEDROCK, + model, + maxOutputTokens: 4096, + useRealTokenLookup: true, + }); + + const result = await initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([Providers.BEDROCK]), + isInitialAgent: true, + }, + db, + ); + + expect(mockGetModelMaxTokens).toHaveBeenCalledWith(model, EModelEndpoint.bedrock, undefined); + expect(result.maxContextTokens).toBe(Math.round((950000 - 4096) * 0.95)); + expect(result.maxContextTokens).toBeGreaterThan(38079); + }, + ); + it('uses user-configured maxContextTokens when provided via model_parameters', async () => { const userValue = 50000; const { agent, req, res, loadTools, db } = createMocks({ diff --git a/packages/api/src/agents/background.spec.ts b/packages/api/src/agents/background.spec.ts index 5418fdeb3f1..2481c853440 100644 --- a/packages/api/src/agents/background.spec.ts +++ b/packages/api/src/agents/background.spec.ts @@ -433,6 +433,12 @@ describe('registerBackgroundTaskTool', () => { expect(automatic.toolDefinitions[0].description).toContain( 'Ordinary tool execution remains process-local', ); + expect(automatic.toolDefinitions[0].description).toContain( + 'A task is outstanding until its result is delivered', + ); + expect(automatic.toolDefinitions[0].description).toContain( + 'Polling or cancelling a finished task retires its pending delivery', + ); }); }); @@ -1369,6 +1375,54 @@ describe('BackgroundTaskRegistryClass', () => { expect(registry.get('u1', 'c1', created.task.id)).toBeUndefined(); }); + it.each(['completed', 'error', 'cancelled', 'blocked'] as const)( + 'preserves pending %s results through both TTLs and releases protection afterward', + async (status) => { + jest.useFakeTimers(); + try { + const registry = new BackgroundTaskRegistryClass(); + const created = registry.create({ + userId: 'u', + conversationId: 'c', + toolCallId: 'call', + toolName: 'execute_code', + harvestStarted: true, + }); + if ('atCapacity' in created) throw new Error('Unexpected capacity rejection'); + if (status === 'cancelled') registry.cancel('u', 'c', created.task.id, 'cancelled'); + else if (status === 'error') registry.fail('u', 'c', created.task.id, 'failed'); + else + registry.complete('u', 'c', created.task.id, { content: 'result', harvestStarted: true }); + registry.markCompletionPersistencePending('u', 'c', created.task.id); + registry.finishHarvest('u', 'c', created.task.id, [{ file_id: 'file' }]); + if (status === 'blocked') registry.blockArtifact('u', 'c', created.task.id, 'blocked'); + registry.claimResult('u', 'c', created.task.id, { kind: 'manual', claimId: 'poll' }); + await jest.advanceTimersByTimeAsync(7 * 60 * 60_000); + expect(registry.get('u', 'c', created.task.id)).toMatchObject({ + completionPersistencePending: true, + }); + registry.markCompletionPersistenceFinished('u', 'c', created.task.id); + expect( + registry.get('u', 'c', created.task.id)?.completionPersistencePending, + ).toBeUndefined(); + if (status === 'blocked') + expect(registry.get('u', 'c', created.task.id)).toMatchObject({ + artifactBlocked: true, + artifact: undefined, + }); + if (status === 'completed') + expect(registry.get('u', 'c', created.task.id)).toMatchObject({ + result: 'result', + resultClaim: { kind: 'manual', claimId: 'poll' }, + }); + await jest.advanceTimersByTimeAsync(61 * 60_000); + expect(registry.get('u', 'c', created.task.id)).toBeUndefined(); + } finally { + jest.useRealTimers(); + } + }, + ); + it('caps concurrent running tasks per conversation', () => { const registry = new BackgroundTaskRegistryClass(); let atCapacity = false; @@ -3057,6 +3111,69 @@ describe('runCheckBackgroundTask (singleton)', () => { expect(claimBackgroundToolResult).toHaveBeenCalledTimes(1); }); + it('counts a finished subagent as outstanding until its result is delivered', async () => { + const store = new InMemorySubagentTaskStore(); + const subagentTasks: HostSubagentTaskConfig = { + store, + scopeId: 'owner:settled-subagent-parent', + completionDelivery: SUBAGENT_COMPLETION_DELIVERY, + }; + const started = store.start({ + scopeId: subagentTasks.scopeId, + idempotencyKey: 'parent-run:parent-agent:call-settled', + parentRunId: 'parent-run', + parentAgentId: 'parent-agent', + parentToolCallId: 'call-settled', + input: 'Research this.', + subagentKind: 'agent', + subagentType: 'researcher', + run: async () => ({ content: 'research done' }), + }); + if (!started.accepted) { + throw new Error('Expected subagent task to start.'); + } + for ( + let i = 0; + i < 50 && store.get(subagentTasks.scopeId, started.task.taskId)?.status === 'running'; + i++ + ) { + await new Promise((resolve) => setImmediate(resolve)); + } + + const listWithWakeups = async (subagentWakeups: string[]) => + JSON.parse( + await runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'settled-subagent-parent', + agentId: 'agent_parent', + args: {}, + subagentTasks, + pendingCompletions: { + list: jest.fn(async () => ({ completions: [], dead: [], complete: true })), + listSubagentWakeups: jest.fn(async () => ({ + taskIds: subagentWakeups, + complete: true, + })), + discard: jest.fn(async () => 'not_pending' as const), + settleClaimed: jest.fn(async () => false), + }, + }), + ); + + /** No durable wake-up (admitted poll-only): nothing will arrive, so nothing is pending. */ + const pollOnly = await listWithWakeups([]); + expect(pollOnly.tasks[0].delivery).toBeUndefined(); + expect(pollOnly.outstanding).toBe(0); + + const listed = await listWithWakeups([started.task.taskId]); + expect(listed.tasks[0]).toEqual( + expect.objectContaining({ status: 'completed', result_available: true, delivery: 'pending' }), + ); + expect(listed.outstanding).toBe(1); + expect(listed.message).toContain('Some finished subagents have not been delivered yet'); + expect(listed.message).not.toContain('cancel'); + }); + it('tells a wakeup-enabled parent to yield on an unchanged running subagent', async () => { const store = new InMemorySubagentTaskStore(); const subagentTasks: HostSubagentTaskConfig = { @@ -3411,3 +3528,451 @@ describe('toolOptionsSchema', () => { expect(parsed).toEqual({ run_in_background: true }); }); }); + +describe('runCheckBackgroundTask delivery semantics', () => { + const pendingControls = ( + overrides: { + list?: () => Promise; + complete?: boolean; + dead?: unknown[]; + subagentWakeups?: string[]; + discard?: () => Promise; + } = {}, + ) => + ({ + list: jest.fn(async () => ({ + completions: await (overrides.list ?? (async () => []))(), + dead: overrides.dead ?? [], + complete: overrides.complete ?? true, + })), + listSubagentWakeups: jest.fn(async () => ({ + taskIds: overrides.subagentWakeups ?? [], + complete: true, + })), + discard: jest.fn(overrides.discard ?? (async () => 'not_pending')), + settleClaimed: jest.fn(async () => true), + }) as never; + + function completedWithWakeup(userId: string, conversationId: string, toolCallId: string) { + const created = backgroundTaskRegistry.create({ + userId, + conversationId, + toolCallId, + toolName: 'bash_tool', + messageId: `${toolCallId}-message`, + }); + if ('atCapacity' in created) { + throw new Error('unexpected capacity'); + } + backgroundTaskRegistry.markCompletionWakeup(userId, conversationId, created.task.id, { + renew: jest.fn(async () => true), + retire: jest.fn(async () => true), + }); + backgroundTaskRegistry.complete(userId, conversationId, created.task.id, { + content: 'finished output', + }); + return created.task.id; + } + + it('counts a finished task as outstanding until its result is delivered', async () => { + const taskId = completedWithWakeup('outstanding-user', 'outstanding-convo', 'outstanding-call'); + + const before = JSON.parse( + await runCheckBackgroundTask({ + userId: 'outstanding-user', + conversationId: 'outstanding-convo', + args: {}, + }), + ); + expect(before.tasks[0]).toEqual( + expect.objectContaining({ + background_task_id: taskId, + status: 'completed', + delivery: 'pending', + }), + ); + expect(before.outstanding).toBe(1); + expect(before.message).toContain('have not been delivered yet'); + + expect( + backgroundTaskRegistry.claimResult('outstanding-user', 'outstanding-convo', taskId, { + kind: 'wakeup', + claimId: 'automatic-delivery', + }), + ).toBe('acquired'); + const after = JSON.parse( + await runCheckBackgroundTask({ + userId: 'outstanding-user', + conversationId: 'outstanding-convo', + args: {}, + }), + ); + expect(after.tasks[0].delivery).toBe('delivered'); + expect(after.outstanding).toBe(0); + expect(after.message).toBeUndefined(); + }); + + it('counts running work as outstanding without claiming undelivered results', async () => { + const created = backgroundTaskRegistry.create({ + userId: 'running-user', + conversationId: 'running-convo', + toolCallId: 'running-call', + toolName: 'bash_tool', + }); + if ('atCapacity' in created) { + throw new Error('unexpected capacity'); + } + + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'running-user', + conversationId: 'running-convo', + args: {}, + }), + ); + expect(listed.tasks[0]).toEqual(expect.objectContaining({ status: 'running' })); + expect(listed.tasks[0].delivery).toBeUndefined(); + expect(listed.outstanding).toBe(1); + expect(listed.message).toBeUndefined(); + }); + + it('lists undelivered results the process-local registry no longer holds', async () => { + const localTaskId = completedWithWakeup('durable-user', 'durable-convo', 'durable-local'); + const pendingCompletions = pendingControls({ + list: async () => [ + { + taskId: localTaskId, + toolName: 'bash_tool', + dispatchedAt: new Date('2026-09-24T12:00:00Z'), + result: { status: 'completed', settledAt: new Date('2026-09-24T12:01:00Z') }, + claimedByWakeup: false, + }, + { + taskId: 'earlier-turn-task', + toolName: 'slow_task', + dispatchedAt: new Date('2026-09-24T11:00:00Z'), + result: { status: 'error', settledAt: new Date('2026-09-24T11:05:00Z') }, + claimedByWakeup: false, + }, + { + taskId: 'other-replica-task', + toolName: 'slow_task', + dispatchedAt: new Date('2026-09-24T11:30:00Z'), + claimedByWakeup: false, + }, + ], + }); + + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'durable-user', + conversationId: 'durable-convo', + args: {}, + pendingCompletions, + }), + ); + + expect( + listed.tasks.map((task: { background_task_id: string }) => task.background_task_id), + ).toEqual([localTaskId, 'earlier-turn-task', 'other-replica-task']); + expect(listed.tasks[1]).toEqual( + expect.objectContaining({ + status: 'error', + delivery: 'pending', + started_at: '2026-09-24T11:00:00.000Z', + settled_at: '2026-09-24T11:05:00.000Z', + }), + ); + expect(listed.tasks[2]).toEqual( + expect.objectContaining({ status: 'running', delivery: 'pending', progress: 0 }), + ); + expect(listed.tasks[1].result).toBeUndefined(); + expect(listed.outstanding).toBe(3); + }); + + it('keeps listing local work when the durable view is unavailable', async () => { + const taskId = completedWithWakeup('degraded-user', 'degraded-convo', 'degraded-call'); + const pendingCompletions = pendingControls({ + list: async () => { + throw new Error('delivery store unavailable'); + }, + }); + + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'degraded-user', + conversationId: 'degraded-convo', + args: {}, + pendingCompletions, + }), + ); + + expect( + listed.tasks.map((task: { background_task_id: string }) => task.background_task_id), + ).toEqual([taskId]); + expect(listed.partial).toBe(true); + expect(listed.warning).toContain('Undelivered results from earlier turns could not be listed'); + }); + + it.each([ + ['discarded', 'cancelled', 'will not arrive as a new turn'], + ['running', 'unavailable', 'cannot be stopped from here'], + ['delivering', 'delivery_scheduled', 'already being delivered'], + ])( + 'reports a %s undelivered completion this process does not hold', + async (outcome, status, message) => { + const pendingCompletions = pendingControls({ discard: async () => outcome }); + + const cancelled = JSON.parse( + await runCheckBackgroundTask({ + userId: 'discard-user', + conversationId: 'discard-convo', + args: { background_task_id: 'earlier-turn-task', action: 'cancel' }, + pendingCompletions, + }), + ); + + expect(cancelled).toEqual( + expect.objectContaining({ status, background_task_id: 'earlier-turn-task' }), + ); + expect(cancelled.message).toContain(message); + }, + ); + + it('reports a local task delivered once the durable store no longer holds its delivery', async () => { + const delivered = completedWithWakeup('reconcile-user', 'reconcile-convo', 'reconcile-done'); + const waiting = completedWithWakeup('reconcile-user', 'reconcile-convo', 'reconcile-waiting'); + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'reconcile-user', + conversationId: 'reconcile-convo', + args: {}, + pendingCompletions: pendingControls({ + list: async () => [ + { + taskId: waiting, + toolName: 'bash_tool', + dispatchedAt: new Date('2026-09-24T12:00:00Z'), + result: { status: 'completed', settledAt: new Date('2026-09-24T12:01:00Z') }, + claimedByWakeup: false, + }, + ], + }), + }), + ); + + const byId = new Map( + listed.tasks.map((task: { background_task_id: string; delivery?: string }) => [ + task.background_task_id, + task.delivery, + ]), + ); + expect(byId.get(delivered)).toBe('delivered'); + expect(byId.get(waiting)).toBe('pending'); + expect(listed.outstanding).toBe(1); + }); + + it('reports a local task whose automatic delivery dead-lettered as failed, not delivered', async () => { + const taskId = completedWithWakeup('dead-user', 'dead-convo', 'dead-call'); + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'dead-user', + conversationId: 'dead-convo', + args: {}, + pendingCompletions: pendingControls({ + dead: [ + { + taskId, + toolName: 'bash_tool', + dispatchedAt: new Date('2026-09-24T12:00:00Z'), + claimedByWakeup: false, + }, + { + taskId: 'restored-dead-task', + toolName: 'slow_task', + dispatchedAt: new Date('2026-09-24T11:00:00Z'), + result: { status: 'completed', settledAt: new Date('2026-09-24T11:01:00Z') }, + claimedByWakeup: false, + }, + ], + }), + }), + ); + + expect(listed.tasks[0]).toEqual( + expect.objectContaining({ background_task_id: taskId, delivery: 'failed' }), + ); + /** A dead letter this process no longer holds is listed from the durable store. */ + expect(listed.tasks[1]).toEqual( + expect.objectContaining({ + background_task_id: 'restored-dead-task', + status: 'completed', + delivery: 'failed', + }), + ); + expect(listed.outstanding).toBe(2); + expect(listed.message).toContain('Automatic delivery failed'); + }); + + it('retires the pending delivery when a local poll claims the durable result', async () => { + const created = backgroundTaskRegistry.create({ + userId: 'local-claim-user', + conversationId: 'local-claim-convo', + toolCallId: 'local-claim-call', + toolName: 'bash_tool', + messageId: 'local-claim-message', + }); + if ('atCapacity' in created) { + throw new Error('unexpected capacity'); + } + const retire = jest.fn(async () => true); + backgroundTaskRegistry.markCompletionWakeup( + 'local-claim-user', + 'local-claim-convo', + created.task.id, + { + renew: jest.fn(async () => true), + retire, + }, + ); + backgroundTaskRegistry.complete('local-claim-user', 'local-claim-convo', created.task.id, { + content: 'finished output', + }); + + const polled = JSON.parse( + await runCheckBackgroundTask({ + userId: 'local-claim-user', + conversationId: 'local-claim-convo', + args: { background_task_id: created.task.id }, + claimBackgroundToolResult: jest.fn(async () => ({ + status: 'acquired' as const, + results: [], + })) as never, + }), + ); + + expect(polled.status).toBe('completed'); + expect(retire).toHaveBeenCalledWith('completion claimed by manual poll', { + onlyIfUnclaimed: true, + }); + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'local-claim-user', + conversationId: 'local-claim-convo', + args: {}, + }), + ); + expect(listed.tasks[0].delivery).toBe('delivered'); + expect(listed.outstanding).toBe(0); + }); + + it('keeps the local view and warns when the durable listing is incomplete', async () => { + const taskId = completedWithWakeup('truncated-user', 'truncated-convo', 'truncated-call'); + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'truncated-user', + conversationId: 'truncated-convo', + args: {}, + pendingCompletions: pendingControls({ complete: false }), + }), + ); + + expect(listed.tasks[0]).toEqual( + expect.objectContaining({ background_task_id: taskId, delivery: 'pending' }), + ); + expect(listed.outstanding).toBe(1); + expect(listed.partial).toBe(true); + expect(listed.warning).toContain('More undelivered results exist'); + }); + + it('lets a finished local task be cancelled without the live-cancellation policy', async () => { + const taskId = completedWithWakeup( + 'settled-cancel-user', + 'settled-cancel-convo', + 'settled-cancel', + ); + + const cancelled = JSON.parse( + await runCheckBackgroundTask({ + userId: 'settled-cancel-user', + conversationId: 'settled-cancel-convo', + args: { background_task_id: taskId, action: 'cancel' }, + }), + ); + + expect(cancelled.status).not.toBe('invalid'); + expect(cancelled).toEqual( + expect.objectContaining({ background_task_id: taskId, status: 'completed' }), + ); + }); + + it('retires the pending delivery of a remote task a manual poll just claimed', async () => { + const pendingCompletions = pendingControls(); + const claimBackgroundToolResult = jest.fn(async () => ({ + status: 'acquired' as const, + results: [ + { + taskId: 'remote-task', + toolName: 'slow_task', + status: 'completed' as const, + output: 'remote result', + settledAt: new Date('2026-09-24T12:00:00Z'), + }, + ], + })); + + const polled = JSON.parse( + await runCheckBackgroundTask({ + userId: 'remote-user', + conversationId: 'remote-convo', + args: { background_task_id: 'remote-task' }, + claimBackgroundToolResult: claimBackgroundToolResult as never, + pendingCompletions, + }), + ); + + expect(polled).toEqual( + expect.objectContaining({ status: 'completed', result: 'remote result' }), + ); + expect( + (pendingCompletions as unknown as { settleClaimed: jest.Mock }).settleClaimed, + ).toHaveBeenCalledWith({ + userId: 'remote-user', + conversationId: 'remote-convo', + taskId: 'remote-task', + }); + }); + + it('reports a failed discard lookup only when nothing else claims the task', async () => { + const cancelled = JSON.parse( + await runCheckBackgroundTask({ + userId: 'discard-user', + conversationId: 'discard-convo', + args: { background_task_id: 'unknown-task', action: 'cancel' }, + pendingCompletions: pendingControls({ + discard: async () => { + throw new Error('delivery store unavailable'); + }, + }), + }), + ); + + expect(cancelled).toEqual(expect.objectContaining({ status: 'unavailable' })); + expect(cancelled.message).toContain('could not be discarded right now'); + }); + + it('falls through to the ordinary lookup when nothing is pending for the task', async () => { + const pendingCompletions = pendingControls({ discard: async () => 'not_pending' }); + + const cancelled = JSON.parse( + await runCheckBackgroundTask({ + userId: 'discard-user', + conversationId: 'discard-convo', + args: { background_task_id: 'unknown-task', action: 'cancel' }, + pendingCompletions, + }), + ); + + expect(cancelled).toEqual(expect.objectContaining({ status: 'not_found' })); + }); +}); diff --git a/packages/api/src/agents/background.ts b/packages/api/src/agents/background.ts index 2c9b744ee47..22cab89ab9d 100644 --- a/packages/api/src/agents/background.ts +++ b/packages/api/src/agents/background.ts @@ -57,8 +57,10 @@ import type { BackgroundToolResultState } from './harvest'; import type { CapabilityToolNames } from './selection'; import { BACKGROUND_TASK_TIMEOUT_MS, + type PendingBackgroundCompletion, type BackgroundToolDeadClaimRecovery, type BackgroundToolWakeupAdmission, + type PendingBackgroundCompletionControls, } from './backgroundCompletion'; import { CREATE_FILE_TOOL_NAME, @@ -358,7 +360,7 @@ Provide a background_task_id to poll one task; omit it to list every background const CHECK_BACKGROUND_TASK_WAKEUP_DESCRIPTION = `Check, control, and retrieve tool or subagent tasks previously dispatched in the background (with run_in_background: true). -Provide a background_task_id to inspect one task; omit it to list every background task in this thread. Background tools and detached subagents use automatic completion delivery: continue independent work or end the turn instead of repeatedly polling an unchanged running task, and the host will resume you when one finishes. Use this tool for explicit status, steer, queue, interrupt, cancel, or cancel_message actions, or as a fallback if automatic delivery is unavailable. Ordinary tool execution remains process-local and does not survive restart; once its result is persisted, completion delivery may continue on another replica. Live subagent controls route across API replicas but do not survive a restart of the process that owns the executor. A completed subagent thread may be continued later through the subagent tool's durable thread id.`; +Provide a background_task_id to inspect one task; omit it to list every background task in this thread. Background tools and detached subagents use automatic completion delivery: continue independent work or end the turn instead of repeatedly polling an unchanged running task, and the host will resume you when one finishes. Use this tool for explicit status, steer, queue, interrupt, cancel, or cancel_message actions, or as a fallback if automatic delivery is unavailable. A task is outstanding until its result is delivered, not merely until it stops running: a finished task whose delivery is "pending" will still arrive as a new turn, so never report it as done or cancelled on the strength of its status alone. Polling or cancelling a finished task retires its pending delivery so it never arrives as a new turn: a poll returns the result now, and cancelling a result this turn can no longer poll discards it. Ordinary tool execution remains process-local and does not survive restart; once its result is persisted, completion delivery may continue on another replica. Live subagent controls route across API replicas but do not survive a restart of the process that owns the executor. A completed subagent thread may be continued later through the subagent tool's durable thread id.`; function checkBackgroundTaskDescription(subagentCompletionWakeups: boolean): string { return subagentCompletionWakeups @@ -834,7 +836,11 @@ export class BackgroundTaskRegistryClass { private sweepBucketTasks(bucket: TaskBucket, now: number): void { for (const [taskId, task] of bucket.tasks) { - if (task.status !== 'running' && now - task.updatedAt > COMPLETED_TASK_TTL_MS) { + if ( + task.status !== 'running' && + task.completionPersistencePending !== true && + now - task.updatedAt > COMPLETED_TASK_TTL_MS + ) { bucket.tasks.delete(taskId); } } @@ -859,7 +865,13 @@ export class BackgroundTaskRegistryClass { } this.lastGlobalSweepAt = now; for (const [bucketKey, bucket] of this.buckets) { - if (now - bucket.lastAccess > IDLE_BUCKET_TTL_MS && bucket.capacityPermits.size === 0) { + if ( + now - bucket.lastAccess > IDLE_BUCKET_TTL_MS && + bucket.capacityPermits.size === 0 && + ![...bucket.tasks.values()].some( + (task) => task.status === 'running' || task.completionPersistencePending === true, + ) + ) { this.buckets.delete(bucketKey); continue; } @@ -1615,7 +1627,14 @@ export class BackgroundTaskRegistryClass { } markCompletionPersistenceFinished(userId: string, conversationId: string, taskId: string): void { - this.update(userId, conversationId, taskId, { completionPersistencePending: undefined }); + const bucket = this.buckets.get(this.key(userId, conversationId)); + const task = bucket?.tasks.get(taskId); + if (bucket == null || task == null) return; + /** Even a policy-blocked task must release retention protection. This only + * clears lifecycle state; the immutable artifact block remains intact. */ + task.completionPersistencePending = undefined; + task.updatedAt = Date.now(); + bucket.lastAccess = task.updatedAt; } markCompletionPersistenceFailed(userId: string, conversationId: string, taskId: string): void { @@ -1840,10 +1859,23 @@ interface SerializedBackgroundTask { result?: string; result_available?: boolean; result_chars?: number; + /** Whether the result has reached the conversation. `pending` results still + * arrive as a new turn unless polled or cancelled first. Absent when the task + * has no automatic delivery, so only a poll ever surfaces its result. */ + delivery?: 'pending' | 'delivered' | 'failed'; note?: string; error?: string; } +const FAILED_DELIVERY_GUIDANCE = + 'Automatic delivery failed for some finished tasks (delivery: "failed"); they will not arrive as a new turn. Poll each to collect its result.'; + +const SUBAGENT_PENDING_DELIVERY_GUIDANCE = + 'Some finished subagents have not been delivered yet (delivery: "pending"); each will arrive as a new turn. Poll one to collect its result now. Do not report them as finished until then.'; + +const PENDING_DELIVERY_GUIDANCE = + 'Some finished tasks have not been delivered yet (delivery: "pending"); each will arrive as a new turn. Poll one to collect its result now, or cancel it so it does not arrive. Do not report these tasks as finished or cancelled until then.'; + /** * Model-facing task timings. The registry keeps epoch milliseconds; everything the * app serializes carries ISO-8601 (`toISOString`), so the poll payload does too. @@ -1900,6 +1932,69 @@ function taskNote(task: BackgroundTask): Pick return {}; } +function taskDelivery(task: BackgroundTask): Pick { + if (task.completionWakeup !== true || task.completionPersistenceFailed === true) { + return {}; + } + if (task.resultClaim != null || task.completionWakeupRetired === true) { + return { delivery: 'delivered' }; + } + return { delivery: 'pending' }; +} + +/** A completion known only to the durable delivery store: dispatched in an earlier + * turn, on another replica, or before a restart, and not delivered yet. */ +function serializePendingCompletion( + completion: PendingBackgroundCompletion, +): SerializedBackgroundTask { + const settled = completion.result; + return { + background_task_id: completion.taskId, + tool: completion.toolName, + status: settled?.status ?? 'running', + progress: settled == null ? 0 : 1, + started_at: completion.dispatchedAt.toISOString(), + ...(settled != null && { settled_at: settled.settledAt.toISOString() }), + delivery: 'pending', + note: + settled == null + ? 'Still running outside this turn; its result will arrive as a new turn when it finishes.' + : 'Finished, but its result has not been delivered; it will arrive as a new turn unless you poll or cancel it.', + }; +} + +/** A local task whose durable delivery settled elsewhere (an automatic wake-up + * on any replica) no longer holds a local claim; the complete durable listing is + * the evidence. An incomplete listing proves nothing, so the local view stands. */ +function reconcileDelivery( + task: SerializedBackgroundTask, + durablePendingTaskIds: ReadonlySet | undefined, + deadTaskIds: ReadonlySet, +): SerializedBackgroundTask { + if (task.delivery !== 'pending' || task.status === 'running') { + return task; + } + if (deadTaskIds.has(task.background_task_id)) { + return { ...task, delivery: 'failed' }; + } + if (durablePendingTaskIds == null || durablePendingTaskIds.has(task.background_task_id)) { + return task; + } + return { ...task, delivery: 'delivered' }; +} + +/** A completion whose automatic delivery dead-lettered and that this process no + * longer holds: never delivered, so only a poll can still collect its result. */ +function serializeDeadCompletion( + completion: PendingBackgroundCompletion, +): SerializedBackgroundTask { + return { + ...serializePendingCompletion(completion), + delivery: 'failed', + note: 'Automatic delivery failed; this result will not arrive as a new turn. Poll it to collect the result.', + }; +} + function serializeTask( task: BackgroundTask, { includeResult }: { includeResult: boolean }, @@ -1914,6 +2009,7 @@ function serializeTask( : {}), ...taskTimings(task), ...resultFields(task, includeResult), + ...taskDelivery(task), ...taskNote(task), ...(task.error !== undefined ? { error: task.error } : {}), }; @@ -1953,6 +2049,8 @@ interface SerializedSubagentTask { error?: string; control_id?: string; message?: string; + /** A finished subagent whose result will still resume the parent turn. */ + delivery?: 'pending'; } function serializeSubagentSnapshot( @@ -2113,6 +2211,8 @@ export async function runCheckBackgroundTask(params: { recoverDeadBackgroundToolClaim?: BackgroundToolDeadClaimRecovery; /** Trusted deployment policy. Defaults false for backward compatibility. */ ordinaryToolCancellation?: boolean; + /** Durable view of this conversation's undelivered background completions. */ + pendingCompletions?: PendingBackgroundCompletionControls; }): Promise { const { userId, conversationId } = params; const args = coerceArgsObject(params.args) ?? {}; @@ -2132,7 +2232,10 @@ export async function runCheckBackgroundTask(params: { if (task != null) { if (action !== 'poll') { if (action === 'cancel') { - if (params.ordinaryToolCancellation !== true) { + /** A finished task has no execution to stop, so the live-cancellation + * policy does not apply: cancelling it falls through to the poll path, + * which retires its pending delivery. */ + if (params.ordinaryToolCancellation !== true && task.status === 'running') { return JSON.stringify({ status: 'invalid', background_task_id: taskId, @@ -2251,6 +2354,23 @@ export async function runCheckBackgroundTask(params: { }); } } + if (durableClaim.status === 'acquired' && task.completionWakeupRetired !== true) { + /** The result reaches the agent here, so its automatic delivery is redundant. */ + await backgroundTaskRegistry + .retireCompletionWakeup( + userId, + conversationId, + taskId, + 'completion claimed by manual poll', + { onlyIfUnclaimed: true }, + ) + .catch((error: unknown) => + logger.warn( + `[background] Failed to retire the delivery of manually claimed task ${taskId}:`, + error, + ), + ); + } if (durableClaim.status === 'not_found' || durableClaim.status === 'not_ready') { const localReplay = task.resultClaim?.kind === 'manual' && task.resultClaim.claimId === invocationId; @@ -2399,6 +2519,42 @@ export async function runCheckBackgroundTask(params: { return JSON.stringify(serializeTask(task, { includeResult: true })); } + /** A failed lookup must not mask a subagent the controls below can still reach. */ + let discardFailed = false; + if (action === 'cancel' && params.pendingCompletions != null) { + let outcome: Awaited> = + 'not_pending'; + try { + outcome = await params.pendingCompletions.discard({ userId, conversationId, taskId }); + } catch (error) { + logger.warn(`[background] Failed to discard pending completion ${taskId}:`, error); + discardFailed = true; + } + if (outcome === 'discarded') { + return JSON.stringify({ + status: 'cancelled', + background_task_id: taskId, + message: 'The finished result was discarded and will not arrive as a new turn.', + }); + } + if (outcome === 'running') { + return JSON.stringify({ + status: 'unavailable', + background_task_id: taskId, + delivery: 'pending', + message: + 'This task is still running outside this turn and cannot be stopped from here. Its result will arrive as a new turn when it finishes; cancel it then to discard the result.', + }); + } + if (outcome === 'delivering') { + return JSON.stringify({ + status: 'delivery_scheduled', + background_task_id: taskId, + message: 'This result is already being delivered as a new turn.', + }); + } + } + const subagentTasks = params.subagentTasks; let subagentPollChecked = false; let subagentPollError: unknown; @@ -2480,6 +2636,15 @@ export async function runCheckBackgroundTask(params: { if (durableClaim.status === 'acquired') { const durableTask = durableClaim.results.find((result) => result.taskId === taskId); if (durableTask != null) { + /** The result reaches the agent here, so its automatic delivery is redundant. */ + await params.pendingCompletions + ?.settleClaimed({ userId, conversationId, taskId }) + .catch((error: unknown) => + logger.warn( + `[background] Failed to retire the delivery of manually claimed task ${taskId}:`, + error, + ), + ); return JSON.stringify(serializeDurableTask(durableTask)); } return JSON.stringify({ @@ -2570,6 +2735,14 @@ export async function runCheckBackgroundTask(params: { } } + if (discardFailed) { + return JSON.stringify({ + status: 'unavailable', + background_task_id: taskId, + message: + 'The pending result could not be discarded right now. It may still arrive as a new turn; retry the cancel shortly.', + }); + } return JSON.stringify({ status: 'not_found', background_task_id: taskId, @@ -2586,7 +2759,37 @@ export async function runCheckBackgroundTask(params: { const tasks = backgroundTaskRegistry.list(userId, conversationId); let subagentTasks: SerializedSubagentTask[] = []; - let listWarning: string | undefined; + const listWarnings: string[] = []; + let pendingCompletions: PendingBackgroundCompletion[] = []; + /** Undelivered task ids from the durable store, when the listing was complete: + * a local task absent from it was delivered on this or another replica. */ + let durablePendingTaskIds: ReadonlySet | undefined; + let deadTaskIds: ReadonlySet = new Set(); + let deadCompletions: PendingBackgroundCompletion[] = []; + if (params.pendingCompletions != null) { + try { + const localTaskIds = new Set(tasks.map((task) => task.id)); + const durable = await params.pendingCompletions.list({ userId, conversationId }); + deadTaskIds = new Set(durable.dead.map(({ taskId }) => taskId)); + /** A dead letter this process no longer holds is still recoverable by a poll. */ + deadCompletions = durable.dead.filter((completion) => !localTaskIds.has(completion.taskId)); + pendingCompletions = durable.completions.filter( + (completion) => !localTaskIds.has(completion.taskId), + ); + if (durable.complete) { + durablePendingTaskIds = new Set(durable.completions.map(({ taskId }) => taskId)); + } else { + listWarnings.push( + 'More undelivered results exist than could be listed; some not shown may still arrive as new turns.', + ); + } + } catch (error) { + logger.warn('[background] Failed to list undelivered background completions:', error); + listWarnings.push( + 'Undelivered results from earlier turns could not be listed; some may still arrive as new turns.', + ); + } + } const completionWakeups = agentUsesSubagentCompletionWakeups( params.subagentTasks, params.agentId, @@ -2607,24 +2810,73 @@ export async function runCheckBackgroundTask(params: { subagentTasks = params.subagentTasks.store .list(params.subagentTasks.scopeId) .map((task) => serializeSubagentSnapshot(task)); - listWarning = `Cross-replica subagent tasks could not be listed: ${error.message}`; + listWarnings.push(`Cross-replica subagent tasks could not be listed: ${error.message}`); } else { throw error; } } } + const ordinaryTasks = [ + ...tasks.map((task) => + reconcileDelivery( + serializeTask(task, { includeResult: false }), + durablePendingTaskIds, + deadTaskIds, + ), + ), + ...pendingCompletions.map(serializePendingCompletion), + ...deadCompletions.map(serializeDeadCompletion), + ]; + /** Pending only with durable evidence: whether a subagent's wake-up exists depends on + * the policy when it was admitted, not on this request's configuration. */ + const finishedSubagents = subagentTasks.filter( + (task) => task.status !== 'running' && task.result_claimed !== true, + ); + if (finishedSubagents.length > 0 && params.pendingCompletions != null) { + try { + const wakeups = await params.pendingCompletions.listSubagentWakeups({ + userId, + conversationId, + }); + const waiting = new Set(wakeups.taskIds); + subagentTasks = subagentTasks.map((task) => + task.status !== 'running' && + task.result_claimed !== true && + waiting.has(task.background_task_id) + ? { ...task, delivery: 'pending' as const } + : task, + ); + } catch (error) { + logger.warn('[background] Failed to list undelivered subagent completions:', error); + listWarnings.push( + 'Undelivered subagent results could not be checked; some may still arrive as new turns.', + ); + } + } + /** Work is outstanding until its result reaches the conversation: a finished + * task with a pending delivery is still going to resume the agent. */ + const isOutstanding = (task: { status: string; delivery?: string }): boolean => + task.status === 'running' || task.delivery === 'pending' || task.delivery === 'failed'; + const outstanding = + ordinaryTasks.filter(isOutstanding).length + subagentTasks.filter(isOutstanding).length; + const isFinishedPending = (task: { status: string; delivery?: string }): boolean => + task.status !== 'running' && task.delivery === 'pending'; + const guidance = [ + ...(ordinaryTasks.some(isFinishedPending) ? [PENDING_DELIVERY_GUIDANCE] : []), + ...(ordinaryTasks.some((task) => task.delivery === 'failed') ? [FAILED_DELIVERY_GUIDANCE] : []), + ...(subagentTasks.some(isFinishedPending) ? [SUBAGENT_PENDING_DELIVERY_GUIDANCE] : []), + ...(completionWakeups && subagentTasks.some((task) => task.status === 'running') + ? [SUBAGENT_WAKEUP_GUIDANCE] + : []), + ]; logger.debug( - `[background] check_background_task listed ${tasks.length + subagentTasks.length} task(s)`, + `[background] check_background_task listed ${ordinaryTasks.length + subagentTasks.length} task(s), ${outstanding} outstanding`, ); return JSON.stringify({ - tasks: [ - ...tasks.map((task) => serializeTask(task, { includeResult: false })), - ...subagentTasks, - ], - ...(completionWakeups && subagentTasks.some((task) => task.status === 'running') - ? { message: SUBAGENT_WAKEUP_GUIDANCE } - : {}), - ...(listWarning != null && { partial: true, warning: listWarning }), + tasks: [...ordinaryTasks, ...subagentTasks], + outstanding, + ...(guidance.length > 0 && { message: guidance.join(' ') }), + ...(listWarnings.length > 0 && { partial: true, warning: listWarnings.join(' ') }), }); } diff --git a/packages/api/src/agents/backgroundCompletion.ts b/packages/api/src/agents/backgroundCompletion.ts index 8ef1d74cca5..c6574030963 100644 --- a/packages/api/src/agents/backgroundCompletion.ts +++ b/packages/api/src/agents/backgroundCompletion.ts @@ -25,6 +25,8 @@ export interface BackgroundToolWakeupRetireOptions { onlyIfUnclaimed?: boolean; /** Reconcile only after the delivery is irreversibly dead-lettered. */ onlyIfDead?: boolean; + /** Report success only when this call retired it, not when it was already delivered. */ + requireTransition?: boolean; } /** Process-local handle for the durable delivery admitted before launch. */ @@ -42,6 +44,9 @@ export interface BackgroundToolWakeupAdmission { * Manual polling requires an atomic unclaimed-only transition: once a * resolver owns a lease, its prepared continuation cannot be cancelled. */ retire: (reason: string, options?: BackgroundToolWakeupRetireOptions) => Promise; + /** Best effort: makes a waiting delivery claimable now that a result it can + * consume exists, e.g. a parent-message projection written without a receipt. */ + expedite?: () => void; } /** Durable ownership repair used by a manual poll after an automatic @@ -62,3 +67,54 @@ export interface BackgroundToolDeadClaimRecoveryInput { export type BackgroundToolDeadClaimRecovery = ( input: BackgroundToolDeadClaimRecoveryInput, ) => Promise; + +/** A background tool completion whose result has not reached its conversation yet, + * read from the durable delivery store rather than a process-local registry. */ +export interface PendingBackgroundCompletion { + taskId: string; + toolName: string; + dispatchedAt: Date; + /** The tool's terminal outcome once it settled; absent while it still runs. */ + result?: { status: 'completed' | 'error' | 'cancelled'; settledAt: Date }; + /** An automatic delivery holds the result and is starting its turn. */ + claimedByWakeup: boolean; +} + +/** + * What cancelling an undelivered completion did: `discarded` retired its delivery, + * so the result never arrives; `running` found the tool still executing where this + * process cannot stop it; `delivering` found the result already being delivered; + * `not_pending` found no undelivered completion for the task. + */ +export type BackgroundCompletionDiscardOutcome = + | 'discarded' + | 'running' + | 'delivering' + | 'not_pending'; + +/** Durable view and control of one principal's undelivered background completions. */ +export interface PendingBackgroundCompletionControls { + /** `complete` is false when more undelivered completions exist than were listed. */ + list: (input: { userId: string; conversationId: string }) => Promise<{ + completions: PendingBackgroundCompletion[]; + /** Completions whose automatic delivery dead-lettered; only a poll recovers them. */ + dead: PendingBackgroundCompletion[]; + complete: boolean; + }>; + /** Subagent tasks whose completion wake-up has not been delivered yet. */ + listSubagentWakeups: (input: { + userId: string; + conversationId: string; + }) => Promise<{ taskIds: string[]; complete: boolean }>; + discard: (input: { + userId: string; + conversationId: string; + taskId: string; + }) => Promise; + /** Retires a task's pending delivery after a manual poll claimed its result. */ + settleClaimed: (input: { + userId: string; + conversationId: string; + taskId: string; + }) => Promise; +} diff --git a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts index eef75fc6777..652ba4df253 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts @@ -1,8 +1,10 @@ import { AGENT_TRIGGER_WORKER_CAPABILITY_BACKGROUND_COMPLETION_RECEIPT_V2 } from '@librechat/data-schemas'; import type { AgentTriggerProducerLeaseStatus } from '@librechat/data-schemas'; +import type { CodeApprovalMode } from 'librechat-data-provider'; import type { EnqueueBackgroundToolCompletion } from './backgroundCompletionWakeup'; import { BACKGROUND_TOOL_WAKEUP_INPUT_MAX_CHARS, + createPendingBackgroundCompletions, createBackgroundToolCompletionWakeupHandler, createBackgroundToolCompletionWakeupResolver, createBackgroundToolDeadClaimRecovery, @@ -45,12 +47,15 @@ function envelope(registrationOverrides = {}) { }); } -function resolverMethods() { +function resolverMethods(codeApprovalMode?: CodeApprovalMode) { const releaseBackgroundToolResultClaims = jest.fn(async () => true); return { releaseBackgroundToolResultClaims, methods: { - getConvo: jest.fn(async () => ({ tenantId: 'tenant-1' })), + getConvo: jest.fn(async () => ({ + tenantId: 'tenant-1', + ...(codeApprovalMode != null && { codeApprovalMode }), + })), getMessages: jest.fn(async () => [ { messageId: 'response-1', @@ -108,6 +113,25 @@ describe('background tool completion wakeups', () => { jest.useRealTimers(); }); + it('expedites its own delivery when a result it can consume appears', async () => { + const expedite = jest.fn(); + const notify = createBackgroundToolCompletionWakeupHandler( + async () => ({ deliveryKey: 'delivery-key-1' }), + async () => true, + async () => true, + undefined, + expedite, + ); + + const admission = await notify(registration()); + if (admission === false) { + throw new Error('Expected an admission'); + } + admission.expedite?.(); + + expect(expedite).toHaveBeenCalledWith('delivery-key-1'); + }); + it('pre-registers the exact task on the invoking response branch', async () => { const enqueue = jest.fn< ReturnType, @@ -440,6 +464,72 @@ describe('background tool completion wakeups', () => { expect(retire).not.toHaveBeenCalled(); }); + describe.each(['projection', 'receipt', 'legacy receipt'] as const)( + '%s approval context', + (source) => { + it.each([undefined, 'ask', 'acceptEdits', 'fullAccess'] as const)( + 'inherits the parent mode %s without granting broader access', + async (mode) => { + const { methods } = resolverMethods(mode); + if (source !== 'projection') { + methods.claimBackgroundToolResults.mockResolvedValue({ + status: 'not_ready', + results: [], + }); + } + if (source === 'receipt') { + methods.claimAgentBackgroundToolResults.mockResolvedValue({ + status: 'acquired', + results: [ + { + taskId: 'task-1', + toolCallId: 'call-1', + toolName: 'slow_tool', + status: 'completed', + output: 'done', + }, + ], + } as never); + } + if (source === 'legacy receipt') { + Reflect.deleteProperty(methods, 'claimAgentBackgroundToolResults'); + methods.getAgentBackgroundToolResult.mockResolvedValue({ + status: 'completed', + output: 'done', + settledAt: new Date(NOW), + } as never); + } + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + const prepared = await resolve(await envelope(), { idempotencyKey: 'delivery-1' }); + + expect(prepared?.status).toBe('ready'); + expect(prepared?.status === 'ready' && prepared.codeApprovalMode).toBe(mode); + expect(methods.getConvo).toHaveBeenCalledTimes(1); + expect(methods.getConvo).toHaveBeenCalledWith('user-1', 'conversation-1'); + }, + ); + }, + ); + + it('rereads the parent approval mode on a delivery retry', async () => { + const { methods } = resolverMethods('fullAccess'); + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + const delivery = await envelope(); + const first = await resolve(delivery, { idempotencyKey: 'delivery-1' }); + expect(first?.status === 'ready' && first.codeApprovalMode).toBe('fullAccess'); + + methods.getConvo.mockResolvedValue({ tenantId: 'tenant-1', codeApprovalMode: 'ask' }); + const retry = await resolve(delivery, { idempotencyKey: 'delivery-1' }); + expect(retry?.status === 'ready' && retry.codeApprovalMode).toBe('ask'); + }); + it('claims a bounded sibling batch and continues from the latest branch leaf', async () => { const { methods } = resolverMethods(); const resolve = createBackgroundToolCompletionWakeupResolver({ @@ -686,6 +776,88 @@ describe('background tool completion wakeups', () => { expect(methods.claimBackgroundToolResults).not.toHaveBeenCalled(); }); + it('backs off by waiting age while the invoking generation keeps running', async () => { + const { methods } = resolverMethods(); + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ status: 'running' }), + }); + const deliveryEnvelope = await envelope(); + + jest.setSystemTime(NOW + 120_000); + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'PARENT_NOT_READY', retryAfter: '12', deferWithoutAttempt: true }, + ); + jest.setSystemTime(NOW + 6 * 60 * 60_000); + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'PARENT_NOT_READY', retryAfter: '60' }, + ); + }); + + it('caps the waiting backoff at the configured interval', async () => { + const { methods } = resolverMethods(); + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ status: 'running' }), + getWaitMaxIntervalMs: () => 20_000, + }); + const deliveryEnvelope = await envelope(); + + jest.setSystemTime(NOW + 6 * 60 * 60_000); + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'PARENT_NOT_READY', retryAfter: '20' }, + ); + }); + + it('keeps backing off while the parent is paused for approval', async () => { + const { methods } = resolverMethods(); + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ status: 'requires_action' }), + }); + const deliveryEnvelope = await envelope(); + + jest.setSystemTime(NOW + 60 * 60_000); + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'PARENT_NOT_READY', retryAfter: '60' }, + ); + }); + + it('re-checks within a second once the parent has settled and only persistence remains', async () => { + const { methods } = resolverMethods(); + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ + status: 'complete', + metadata: { terminalPersistencePending: true }, + }), + }); + const deliveryEnvelope = await envelope(); + + jest.setSystemTime(NOW + 10 * 60_000); + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'PARENT_NOT_READY', retryAfter: '1' }, + ); + }); + + it('backs off by waiting age while the tool result is not durable yet', async () => { + const { methods } = resolverMethods(); + methods.claimBackgroundToolResults.mockResolvedValue({ status: 'missing', results: [] }); + const resolve = createBackgroundToolCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + const deliveryEnvelope = await envelope(); + + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'BACKGROUND_TOOL_RESULT_NOT_READY', retryAfter: '5' }, + ); + jest.setSystemTime(NOW + 5 * 60_000); + await expect(resolve(deliveryEnvelope, { idempotencyKey: 'delivery-1' })).rejects.toMatchObject( + { code: 'BACKGROUND_TOOL_RESULT_NOT_READY', retryAfter: '30' }, + ); + }); + it('does not manufacture terminal evidence from wall-clock age', async () => { const { methods } = resolverMethods(); methods.claimBackgroundToolResults.mockResolvedValue({ status: 'missing', results: [] }); @@ -720,3 +892,143 @@ describe('background tool completion wakeups', () => { ); }); }); + +describe('pending background completions', () => { + const dispatchedAt = new Date(NOW - 60_000); + const settled = { status: 'completed' as const, settledAt: new Date(NOW) }; + const row = (overrides = {}) => ({ + deliveryKey: 'delivery-key-1', + taskId: 'task-1', + toolCallId: 'call-1', + toolName: 'slow_tool', + dispatchedAt, + claimedByWakeup: false, + ...overrides, + }); + const listing = (completions: Array>, truncated = false) => + jest.fn(async () => ({ completions, dead: [row({ taskId: 'task-dead' })], truncated })); + const listTaskIds = jest.fn(async () => ({ taskIds: ['child-1'], truncated: false })); + + it('lists the durable view for the owner without delivery internals', async () => { + const list = listing([row({ result: settled })], true); + const pending = createPendingBackgroundCompletions({ list, listTaskIds, retire: jest.fn() }); + + await expect( + pending.list({ userId: 'user-1', conversationId: 'conversation-1' }), + ).resolves.toEqual({ + dead: [ + { + taskId: 'task-dead', + toolName: 'slow_tool', + dispatchedAt, + claimedByWakeup: false, + }, + ], + completions: [ + { + taskId: 'task-1', + toolName: 'slow_tool', + dispatchedAt, + result: settled, + claimedByWakeup: false, + }, + ], + complete: false, + }); + expect(list).toHaveBeenCalledWith({ + user: 'user-1', + conversationId: 'conversation-1', + sourceId: 'background-tool-completion', + }); + }); + + it('discards a settled, unclaimed result by looking up that task and retiring it exactly', async () => { + const retire = jest.fn(async () => true); + const list = listing([row({ result: settled })]); + const pending = createPendingBackgroundCompletions({ list, listTaskIds, retire }); + + await expect( + pending.discard({ userId: 'user-1', conversationId: 'conversation-1', taskId: 'task-1' }), + ).resolves.toBe('discarded'); + expect(list).toHaveBeenCalledWith({ + user: 'user-1', + conversationId: 'conversation-1', + sourceId: 'background-tool-completion', + taskId: 'task-1', + }); + expect(retire).toHaveBeenCalledWith( + 'delivery-key-1', + 'background-tool-completion', + 'background result discarded by its owner', + { onlyIfUnclaimed: true, requireTransition: true }, + ); + }); + + it("retires a manually claimed task's delivery only while no wake-up holds it", async () => { + const retire = jest.fn(async () => true); + const pending = createPendingBackgroundCompletions({ + list: listing([row({ result: settled })]), + listTaskIds, + retire, + }); + + await expect( + pending.settleClaimed({ + userId: 'user-1', + conversationId: 'conversation-1', + taskId: 'task-1', + }), + ).resolves.toBe(true); + expect(retire).toHaveBeenCalledWith( + 'delivery-key-1', + 'background-tool-completion', + 'completion claimed by manual poll', + { onlyIfUnclaimed: true }, + ); + await expect( + createPendingBackgroundCompletions({ list: listing([]), listTaskIds, retire }).settleClaimed({ + userId: 'user-1', + conversationId: 'conversation-1', + taskId: 'task-1', + }), + ).resolves.toBe(false); + }); + + it('lists undelivered subagent wake-ups from their own source', async () => { + const pending = createPendingBackgroundCompletions({ + list: listing([]), + listTaskIds, + retire: jest.fn(), + }); + + await expect( + pending.listSubagentWakeups({ userId: 'user-1', conversationId: 'conversation-1' }), + ).resolves.toEqual({ taskIds: ['child-1'], complete: true }); + expect(listTaskIds).toHaveBeenCalledWith({ + user: 'user-1', + conversationId: 'conversation-1', + sourceId: 'subagent-completion', + }); + }); + + it.each([ + ['not_pending', [], true], + ['running', [row()], true], + ['delivering', [row({ result: settled, claimedByWakeup: true })], true], + ['delivering', [row({ result: settled })], false], + ])('reports %s without discarding what it cannot', async (outcome, rows, retired) => { + const retire = jest.fn(async () => retired); + const pending = createPendingBackgroundCompletions({ + list: listing(rows), + listTaskIds, + retire, + }); + + await expect( + pending.discard({ userId: 'user-1', conversationId: 'conversation-1', taskId: 'task-1' }), + ).resolves.toBe(outcome); + if (outcome !== 'delivering' || rows[0]?.claimedByWakeup === true) { + expect(retire).not.toHaveBeenCalled(); + } + }); +}); diff --git a/packages/api/src/agents/backgroundCompletionWakeup.ts b/packages/api/src/agents/backgroundCompletionWakeup.ts index e107118bffb..63e6f45c89e 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.ts @@ -10,9 +10,11 @@ import type { } from '@librechat/data-schemas'; import type { BackgroundToolDeadClaimRecovery, + PendingBackgroundCompletion, BackgroundToolWakeupAdmission, BackgroundToolWakeupRegistration, BackgroundToolWakeupRetireOptions, + PendingBackgroundCompletionControls, } from './backgroundCompletion'; import type { AgentTriggerContinuePreparation, @@ -21,7 +23,9 @@ import type { import type { AgentContinueTriggerEnvelope } from './triggers/envelope'; import type { AgentTriggerDispatchContext } from './triggers/dispatch'; import type { AgentTriggerEnqueueOptions } from './triggers/delivery'; +import { WAITING_RETRY_CAP_MS, waitingRetryAfter } from './triggers/backoff'; import { BACKGROUND_TOOL_PRODUCER_LEASE_MS } from './backgroundCompletion'; +import { SUBAGENT_COMPLETION_SOURCE } from './subagentCompletionWakeup'; import { createAgentTriggerEnvelope } from './triggers/envelope'; import { AgentTriggerExecutionError } from './triggers/host'; import { truncateMiddle } from '~/utils'; @@ -94,6 +98,8 @@ export interface BackgroundToolCompletionWakeupResolverDeps { methods: WakeupMethods; getGenerationJob: (conversationId: string) => Promise; getResultBatchSize?: () => number | undefined; + /** Longest a waiting delivery re-checks readiness; the backoff default otherwise. */ + getWaitMaxIntervalMs?: () => number | undefined; } function executionError( @@ -152,6 +158,12 @@ function isParentActive(job: GenerationState | null): boolean { ); } +/** A running or approval-paused parent can stay busy for hours; one that has + * settled and is only finishing terminal persistence clears within moments. */ +function isParentWorking(job: GenerationState | null): boolean { + return job?.status === 'running' || job?.status === 'requires_action'; +} + function timestamp(message: Pick): number { const value = message.createdAt; if (value instanceof Date) { @@ -252,9 +264,12 @@ export function createBackgroundToolCompletionWakeupResolver({ methods, getGenerationJob, getResultBatchSize, + getWaitMaxIntervalMs, }: BackgroundToolCompletionWakeupResolverDeps): NonNullable< AgentTriggerExecutionHostDeps['prepareContinue'] > { + const waitingRetry = (receivedAt: number): string => + waitingRetryAfter(receivedAt, Date.now(), getWaitMaxIntervalMs?.() ?? WAITING_RETRY_CAP_MS); return async ( envelope: AgentContinueTriggerEnvelope, context: AgentTriggerDispatchContext, @@ -288,7 +303,7 @@ export function createBackgroundToolCompletionWakeupResolver({ code: 'PARENT_NOT_READY', retryable: true, status: 409, - retryAfter: '1', + retryAfter: isParentWorking(parentJob) ? waitingRetry(envelope.receivedAt) : '1', deferWithoutAttempt: true, }); } @@ -386,6 +401,7 @@ export function createBackgroundToolCompletionWakeupResolver({ return { status: 'ready', parentMessageId, + ...(parent.codeApprovalMode != null && { codeApprovalMode: parent.codeApprovalMode }), input, releaseOnDefiniteFailure: async () => { const released = await methods.releaseBackgroundToolResultClaims({ @@ -455,6 +471,7 @@ export function createBackgroundToolCompletionWakeupResolver({ return { status: 'ready', parentMessageId, + ...(parent.codeApprovalMode != null && { codeApprovalMode: parent.codeApprovalMode }), input: buildWakeupInput(receiptClaim.results), releaseOnDefiniteFailure: async () => { const projectionReleased = await methods.releaseBackgroundToolResultClaims({ @@ -490,6 +507,7 @@ export function createBackgroundToolCompletionWakeupResolver({ return { status: 'ready', parentMessageId, + ...(parent.codeApprovalMode != null && { codeApprovalMode: parent.codeApprovalMode }), input: buildWakeupInput([ { ...registration, status: receipt.status, output: receipt.output }, ]), @@ -523,18 +541,115 @@ export function createBackgroundToolCompletionWakeupResolver({ code: 'BACKGROUND_TOOL_RESULT_NOT_READY', retryable: true, status: 409, - retryAfter: '1', + retryAfter: waitingRetry(envelope.receivedAt), deferWithoutAttempt: true, }); }; } +/** Lists and discards a conversation's undelivered background completions from the + * durable delivery store, which outlives the process-local task registry: a result + * dispatched in an earlier turn, on another replica, or before a restart is still + * going to arrive, and the owner must be able to see and stop that. */ +export function createPendingBackgroundCompletions(deps: { + list: (input: { + user: string; + conversationId: string; + sourceId: string; + taskId?: string; + }) => Promise<{ + completions: Array; + dead: Array; + truncated: boolean; + }>; + listTaskIds: (input: { + user: string; + conversationId: string; + sourceId: string; + }) => Promise<{ taskIds: string[]; truncated: boolean }>; + retire: RetireBackgroundToolCompletion; +}): PendingBackgroundCompletionControls { + const read = (input: { userId: string; conversationId: string; taskId?: string }) => + deps.list({ + user: input.userId, + conversationId: input.conversationId, + sourceId: BACKGROUND_TOOL_COMPLETION_SOURCE, + ...(input.taskId != null && { taskId: input.taskId }), + }); + return { + list: async (input) => { + const { completions, dead, truncated } = await read(input); + const project = ({ + taskId, + toolName, + dispatchedAt, + result, + claimedByWakeup, + }: PendingBackgroundCompletion): PendingBackgroundCompletion => ({ + taskId, + toolName, + dispatchedAt, + ...(result != null && { result }), + claimedByWakeup, + }); + return { + completions: completions.map(project), + dead: dead.map(project), + complete: !truncated, + }; + }, + discard: async (input) => { + const [completion] = (await read(input)).completions; + if (completion == null) { + return 'not_pending'; + } + if (completion.result == null) { + return 'running'; + } + if (completion.claimedByWakeup) { + return 'delivering'; + } + /** Unclaimed-only: once a resolver owns the delivery its continuation can no + * longer be withdrawn, so that race, including one it already finished, + * reports as delivering rather than discarded. */ + const retired = await deps.retire( + completion.deliveryKey, + BACKGROUND_TOOL_COMPLETION_SOURCE, + 'background result discarded by its owner', + { onlyIfUnclaimed: true, requireTransition: true }, + ); + return retired ? 'discarded' : 'delivering'; + }, + listSubagentWakeups: async (input) => { + const { taskIds, truncated } = await deps.listTaskIds({ + user: input.userId, + conversationId: input.conversationId, + sourceId: SUBAGENT_COMPLETION_SOURCE, + }); + return { taskIds, complete: !truncated }; + }, + settleClaimed: async (input) => { + const [completion] = (await read(input)).completions; + if (completion == null) { + return false; + } + return deps.retire( + completion.deliveryKey, + BACKGROUND_TOOL_COMPLETION_SOURCE, + 'completion claimed by manual poll', + { onlyIfUnclaimed: true }, + ); + }, + }; +} + /** Pre-registers the ordered completion delivery before external tool work starts. */ export function createBackgroundToolCompletionWakeupHandler( enqueue: EnqueueBackgroundToolCompletion, retire: RetireBackgroundToolCompletion, renewProducerLease: RenewBackgroundToolCompletionProducerLease, persistResult?: PersistBackgroundToolCompletionResult, + expedite?: (deliveryKey: string) => void, ): ( registration: BackgroundToolWakeupRegistration, ) => Promise { @@ -598,6 +713,7 @@ export function createBackgroundToolCompletionWakeupHandler( options == null ? retire(admitted.deliveryKey, BACKGROUND_TOOL_COMPLETION_SOURCE, reason) : retire(admitted.deliveryKey, BACKGROUND_TOOL_COMPLETION_SOURCE, reason, options), + ...(expedite == null ? {} : { expedite: () => expedite(admitted.deliveryKey) }), }; }; } diff --git a/packages/api/src/agents/checkpoints/deletion.integration.spec.ts b/packages/api/src/agents/checkpoints/deletion.integration.spec.ts index 84d8b59eb9b..5da8e089e3b 100644 --- a/packages/api/src/agents/checkpoints/deletion.integration.spec.ts +++ b/packages/api/src/agents/checkpoints/deletion.integration.spec.ts @@ -242,7 +242,9 @@ test('retains writer obligations, renewed revisions, and lookup failures', async await deletion.remember(['child']); const getJobs = jest.fn().mockResolvedValue(['pending-host']); const reclaim = createCheckpointDeletionReclaimer(getJobs); - expect(await reclaim(25)).toBe(0); + const activity = { found: false }; + expect(await reclaim(25, activity)).toBe(0); + expect(activity.found).toBe(true); getJobs.mockRejectedValueOnce(new Error('job store unavailable')); await expect(reclaim(25)).rejects.toThrow('reclamation failed'); getJobs.mockImplementationOnce(async () => { diff --git a/packages/api/src/agents/checkpoints/deletion.ts b/packages/api/src/agents/checkpoints/deletion.ts index 074341dc2f9..36fcd1a0cc1 100644 --- a/packages/api/src/agents/checkpoints/deletion.ts +++ b/packages/api/src/agents/checkpoints/deletion.ts @@ -225,9 +225,9 @@ export async function openCheckpointDeletion( /** Replay captured identities after topology deletion; never sweep new generations. */ export function createCheckpointDeletionReclaimer( getOwnerJobs: (userId: string, tenantId?: string) => Promise, -): (limit: number) => Promise { +): (limit: number, activity?: { found: boolean }) => Promise { let after: string | undefined; - return async (limit) => { + return async (limit, activity) => { if (!Number.isSafeInteger(limit) || limit <= 0) throw new Error('Invalid reclamation limit'); const db = mongoose.connection.db; if (!db || mongoose.connection.readyState !== 1) @@ -238,6 +238,7 @@ export function createCheckpointDeletionReclaimer( .sort({ _id: 1 }) .limit(limit) .toArray(); + if (activity != null && targets.length > 0) activity.found = true; after = targets.length === limit ? targets[targets.length - 1]._id : undefined; const jobsByOwner = new Map>(); diff --git a/packages/api/src/agents/client.spec.ts b/packages/api/src/agents/client.spec.ts index 97700e4492a..f3a5b56e9d5 100644 --- a/packages/api/src/agents/client.spec.ts +++ b/packages/api/src/agents/client.spec.ts @@ -1,8 +1,11 @@ import { Tokenizer as AiTokenizer } from 'ai-tokenizer'; import { Providers, StandardGraph } from '@librechat/agents'; -import { HumanMessage } from '@librechat/agents/langchain/messages'; +import { HumanMessage, SystemMessage } from '@librechat/agents/langchain/messages'; import { ContentTypes, DEFAULT_MAX_RETAINED_TOOL_COUNT_CHARS } from 'librechat-data-provider'; import type { TMessage } from 'librechat-data-provider'; +import type { LCTool } from '@librechat/agents'; +import type { FormattedMessageWithContent } from './client'; +import type { EncodingName } from '~/utils/tokenizer'; import { collectToolCallIds, countRetainedToolTokens, @@ -10,11 +13,130 @@ import { prependQuotes, prependFileContext, applyAttachmentOnlyText, - type FormattedMessageWithContent, } from './client'; import { ATTACHMENT_ONLY_TEXT } from '~/files/context'; +import Tokenizer from '~/utils/tokenizer'; describe('createCachedTokenCounter', () => { + const encodings: EncodingName[] = ['o200k_base', 'claude']; + + it.each(encodings)('counts SDK-sized slices accurately with %s', async (encoding) => { + const counter = await createCachedTokenCounter(encoding); + const text = 'word '.repeat(3277).slice(0, 16384); + const exact = Tokenizer.countExactTokens(text, encoding)!; + const expected = encoding === 'claude' ? Math.ceil(exact * 1.1) : exact; + expect(counter(new SystemMessage(text))).toBeGreaterThanOrEqual(expected * 0.99); + expect(counter(new SystemMessage(text))).toBeLessThan(expected * 1.02); + }); + + it.each([false, true])( + 'sends fitting instructions with summarization=%s', + async (summarizationEnabled) => { + const graph = new StandardGraph({ + runId: `bounded-instructions-${summarizationEnabled}`, + agents: [ + { + agentId: 'primary', + provider: Providers.OPENAI, + instructions: 'word '.repeat(1024), + maxContextTokens: 4000, + summarizationEnabled, + }, + ], + tokenCounter: await createCachedTokenCounter('o200k_base'), + }); + graph.overrideTestModel(['ok']); + const result = await graph + .createAgentNode('primary') + .invoke( + { messages: [new HumanMessage('Hi')] }, + { configurable: { thread_id: graph.runId }, recursionLimit: 12 }, + ); + expect(result.messages[result.messages.length - 1]?.content).toBe('ok'); + expect(graph.agentContexts.get('primary')?.instructionTokens).toBeLessThan(1100); + }, + ); + + it('sends a fitting 400-tool programmatic prompt with no custom instructions', async () => { + const tools: LCTool[] = Array.from({ length: 400 }, (_, index) => ({ + name: `lookup_${index}`, + description: + 'Search the project documents and return matching results with source references. '.repeat( + 16, + ), + parameters: { + type: 'object', + properties: { + query: { + type: 'string', + description: 'The search query to match against project documents.', + }, + limit: { + type: 'number', + description: 'The maximum number of matching documents to return.', + }, + }, + required: ['query'], + }, + allowed_callers: ['code_execution'], + })); + const graph = new StandardGraph({ + runId: 'bounded-programmatic-instructions', + agents: [ + { + agentId: 'primary', + provider: Providers.ANTHROPIC, + instructions: '', + maxContextTokens: 500000, + toolDefinitions: [ + { + name: 'run_tools_with_code', + description: 'Execute code', + parameters: { type: 'object', properties: { code: { type: 'string' } } }, + }, + ...tools, + ], + }, + ], + tokenCounter: await createCachedTokenCounter('claude'), + }); + graph.overrideTestModel(['ok']); + const result = await graph + .createAgentNode('primary') + .invoke( + { messages: [new HumanMessage('Hi')] }, + { configurable: { thread_id: graph.runId }, recursionLimit: 12 }, + ); + expect(result.messages[result.messages.length - 1]?.content).toBe('ok'); + const context = graph.agentContexts.get('primary'); + expect(context?.systemMessageTokens).toBeGreaterThan(100000); + expect(context?.instructionTokens).toBeLessThan(150000); + }); + + it('still rejects instructions that genuinely exceed the context budget', async () => { + const graph = new StandardGraph({ + runId: 'oversized-instructions', + agents: [ + { + agentId: 'primary', + provider: Providers.OPENAI, + instructions: 'word '.repeat(5000), + maxContextTokens: 4000, + }, + ], + tokenCounter: await createCachedTokenCounter('o200k_base'), + }); + graph.overrideTestModel(['should not be called']); + await expect( + graph + .createAgentNode('primary') + .invoke( + { messages: [new HumanMessage('Hi')] }, + { configurable: { thread_id: graph.runId }, recursionLimit: 12 }, + ), + ).rejects.toThrow('empty_messages'); + }); + it('enables stable-message reuse in the agents runtime', async () => { const getTokenCount = jest.spyOn(AiTokenizer.prototype, 'count'); try { @@ -33,7 +155,7 @@ describe('createCachedTokenCounter', () => { const agentContext = graph.agentContexts.get('primary'); await agentContext?.tokenCalculationPromise; getTokenCount.mockClear(); - const message = new HumanMessage('Stable retained context'); + const message = new HumanMessage('Stable retained context '.repeat(400)); agentContext?.contextPressureTokenCounts?.count(message); const callsAfterFirstCount = getTokenCount.mock.calls.length; diff --git a/packages/api/src/agents/errors.spec.ts b/packages/api/src/agents/errors.spec.ts index 868e60ea773..bc621f2baf3 100644 --- a/packages/api/src/agents/errors.spec.ts +++ b/packages/api/src/agents/errors.spec.ts @@ -7,12 +7,28 @@ import { resolveLangChainError, getUserFacingProviderError, isFatalAgentInitializationError, + getAgentErrorMetadata, AGENT_ATTACHMENT_LIMIT_EXCEEDED, AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE, isStepLimitError, } from './errors'; +import { MCPAuthenticationRejectedError, MCPAuthenticationRefreshError } from '~/mcp/errors'; +import { OboTokenResolutionError } from '~/mcp/oauth/obo'; +import { OpenIDReauthRequiredError } from '~/utils/oidc'; describe('isFatalAgentInitializationError', () => { + it.each([ + new OpenIDReauthRequiredError('Please sign in again'), + new MCPAuthenticationRejectedError('private-mcp', false), + new MCPAuthenticationRejectedError('private-mcp', true), + new MCPAuthenticationRefreshError(new Error('temporarily unavailable')), + new OboTokenResolutionError('session_refresh_failed', 'Please sign in again', false), + new OboTokenResolutionError('session_refresh_failed', 'Retry later', true), + ])('never hides a credential outcome behind optional-tool fallback: %s', (error) => { + expect(isFatalAgentInitializationError(error)).toBe(true); + expect(isFatalAgentInitializationError(error, { allowExpectedMCPFallback: true })).toBe(true); + }); + it('propagates cancellation even when optional MCP fallback is allowed', () => { const abort = new DOMException('Stopped', 'AbortError'); const controller = new AbortController(); @@ -254,3 +270,39 @@ describe('isStepLimitError', () => { expect(isStepLimitError(error)).toBe(false); }); }); + +describe('getAgentErrorMetadata', () => { + it.each([ + [new OpenIDReauthRequiredError('Sign in again'), 401, undefined], + [new MCPAuthenticationRejectedError('private', false), 403, 'MCP_AUTHENTICATION_REJECTED'], + [new MCPAuthenticationRefreshError(), 503, 'MCP_AUTHENTICATION_REFRESH_FAILED'], + ])('preserves typed MCP statusCode and code: %s', (error, status, code) => { + expect(getAgentErrorMetadata(error)).toEqual({ status, ...(code ? { code } : {}) }); + }); + + it.each([false, true])( + 'classifies OBO failure without rewriting the original error: retryable=%s', + (retryable) => { + const error = new OboTokenResolutionError('exchange_failed', 'Exchange failed', retryable); + expect(getAgentErrorMetadata(error)).toEqual({ + status: retryable ? 503 : 403, + code: retryable ? 'MCP_AUTHENTICATION_REFRESH_FAILED' : 'MCP_AUTHENTICATION_REJECTED', + retryable, + }); + expect(error.reason).toBe('exchange_failed'); + expect(error).not.toHaveProperty('statusCode'); + }, + ); + + it('keeps status precedence and code-only provider failures', () => { + expect(getAgentErrorMetadata({ status: 409, statusCode: 401, code: 'RUN_REPLACED' })).toEqual({ + status: 409, + code: 'RUN_REPLACED', + }); + expect(getAgentErrorMetadata({ code: 'ERR_REMOTE' })).toEqual({ code: 'ERR_REMOTE' }); + }); + + it.each([null, 'bad', 399, 600, 401.5, NaN])('rejects invalid outward status: %s', (status) => { + expect(getAgentErrorMetadata({ status })).toEqual({}); + }); +}); diff --git a/packages/api/src/agents/errors.ts b/packages/api/src/agents/errors.ts index c2caff0a7ec..38777e76563 100644 --- a/packages/api/src/agents/errors.ts +++ b/packages/api/src/agents/errors.ts @@ -4,7 +4,8 @@ import { parseLangChainErrorCode, stripLangChainTroubleshootingUrl, } from 'librechat-data-provider'; -import { isOwnedAbortError } from '~/utils/errors'; +import { MCPErrorCodes, isMCPInitializationError } from '~/mcp/errors'; +import { OboTokenResolutionError } from '~/mcp/oauth/obo'; export const AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE = 'AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE'; export const AGENT_ATTACHMENT_LIMIT_EXCEEDED = 'AGENT_ATTACHMENT_LIMIT_EXCEEDED'; @@ -58,7 +59,7 @@ export function isFatalAgentInitializationError( ): boolean { const code = getErrorCode(error); return ( - isOwnedAbortError(error, options.signal) || + isMCPInitializationError(error, options.signal) || FATAL_AGENT_INITIALIZATION_CODES.has(code as string) || (code === AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE && options.allowExpectedMCPFallback !== true) ); @@ -190,3 +191,29 @@ export function isStepLimitError(error: unknown): boolean { } return false; } + +/** Outward metadata shared by UI generation failures and both remote agent APIs. */ +export function getAgentErrorMetadata( + error: unknown, +): { status?: number; code?: string; retryable?: boolean } | undefined { + if (error instanceof OboTokenResolutionError) { + return { + status: error.retryable ? 503 : 403, + code: error.retryable + ? MCPErrorCodes.AUTHENTICATION_REFRESH_FAILED + : MCPErrorCodes.AUTHENTICATION_REJECTED, + retryable: error.retryable, + }; + } + if (!error || typeof error !== 'object') { + return undefined; + } + const candidate = error as { status?: unknown; statusCode?: unknown; code?: unknown }; + const status = candidate.status ?? candidate.statusCode; + return { + ...(typeof status === 'number' && Number.isInteger(status) && status >= 400 && status < 600 + ? { status } + : {}), + ...(typeof candidate.code === 'string' ? { code: candidate.code } : {}), + }; +} diff --git a/packages/api/src/agents/handlers.background.spec.ts b/packages/api/src/agents/handlers.background.spec.ts index 645a473cf9e..6eac3eedb1d 100644 --- a/packages/api/src/agents/handlers.background.spec.ts +++ b/packages/api/src/agents/handlers.background.spec.ts @@ -388,6 +388,44 @@ describe('createToolExecuteHandler — background tool calls', () => { expect(JSON.parse(dispatch.content).message).not.toContain('host will resume you'); }); + it('expedites the delivery when only the parent-message projection was persisted', async () => { + const tool = makeSearchTool({ calls: 0 }); + const expedite = jest.fn(); + const retire = jest.fn(async () => true); + const handler = createToolExecuteHandler({ + loadTools: async () => ({ loadedTools: [tool] }), + backgroundToolCompletion: { + preregister: jest.fn(async () => ({ + renew: jest.fn(async () => true), + persistResult: jest.fn(async () => false), + retire, + expedite, + })), + persist: jest.fn(async () => true), + claim: jest.fn(async () => ({ status: 'acquired' as const, results: [] })), + }, + }); + + await runBatch(handler, { + toolCalls: [ + { + id: 'call-projection-only', + name: tool.name, + args: { q: 'projection', run_in_background: true }, + stepId: 'step-projection-only', + }, + ], + agentId: 'agent_parent_1', + configurable: buildConfig([tool.name]), + metadata: { thread_id: 'exec_convo', run_id: 'response-1' }, + }); + await flushMicrotasks(); + await flushMicrotasks(); + + expect(expedite).toHaveBeenCalledTimes(1); + expect(retire).not.toHaveBeenCalled(); + }); + it('retires a preregistered delivery when terminal persistence fails', async () => { const tool = makeSearchTool({ calls: 0 }); const retire = jest.fn(async () => true); @@ -2583,6 +2621,59 @@ describe('createToolExecuteHandler — backgrounded code execution', () => { expect(poll[0].artifact).toBeUndefined(); }); + it('releases stored files to same-turn polls while the row patch waits for the dispatch turn', async () => { + const state: CodeToolState = { calls: 0 }; + const emitted: unknown[] = []; + const handler = createToolExecuteHandler({ + loadTools: async () => ({ loadedTools: [makeCodeTool(state)] }), + /** Files are stored; the row patch then waits for the long dispatch turn. */ + persistBackgroundCodeResult: (params) => { + params.onFilesPersisted?.([{ file_id: 'f1', toolCallId: params.toolCallId }]); + return new Promise(() => undefined); + }, + emitAttachment: (attachment) => { + emitted.push(attachment); + }, + }); + const configurable = buildConfig(['execute_code']); + + const dispatch = await runBatch(handler, { + toolCalls: [codeCall({ id: 'call_code_long_turn' })], + agentId: 'a', + configurable, + metadata: { thread_id: 'exec_convo_code_long_turn', run_id: 'msg-long-turn' }, + }); + await flushMicrotasks(); + await flushMicrotasks(); + + const poll = (await runBatch(handler, { + toolCalls: [ + { + id: 'call_poll_long_turn', + name: CHECK_BACKGROUND_TASK_NAME, + args: { background_task_id: JSON.parse(dispatch[0].content).background_task_id }, + }, + ], + agentId: 'a', + configurable, + metadata: { thread_id: 'exec_convo_code_long_turn', run_id: 'msg-long-turn' }, + })) as Array<{ content: string; artifact?: unknown }>; + + const polled = JSON.parse(poll[0].content); + expect(polled.status).toBe('completed'); + expect(polled.note).toContain('attached to the tool call'); + expect(emitted[0]).toEqual({ file_id: 'f1', toolCallId: 'call_code_long_turn' }); + expect(poll[0].artifact).toEqual(CODE_ARTIFACT); + /** Still protected from retention eviction while the row patch waits. */ + expect( + backgroundTaskRegistry.get( + 'exec_user', + 'exec_convo_code_long_turn', + JSON.parse(dispatch[0].content).background_task_id, + )?.completionPersistencePending, + ).toBe(true); + }); + it('falls back to poll-turn delivery when the harvest fails (files not lost)', async () => { const state: CodeToolState = { calls: 0 }; const toolEndCalls: Array<{ name?: string; artifact?: unknown }> = []; diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index f991418273d..969e886c1b5 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -35,6 +35,12 @@ import type { import type { CodeEnvRef, CodeWorkspaceOperation, PtcToolCallEvent } from 'librechat-data-provider'; import type { StructuredToolInterface } from '@librechat/agents/langchain/tools'; import type { CodeEnvFile, CodeSessionContext } from '@librechat/agents'; +import type { + BackgroundToolDeadClaimRecovery, + BackgroundToolWakeupAdmission, + BackgroundToolWakeupRegistration, + PendingBackgroundCompletionControls, +} from './backgroundCompletion'; import type { WorkspaceEditResult, WorkspacePreviewEditResult, @@ -43,11 +49,6 @@ import type { WorkspaceSearchResult, WorkspaceWriteResult, } from '~/code/workspace'; -import type { - BackgroundToolDeadClaimRecovery, - BackgroundToolWakeupAdmission, - BackgroundToolWakeupRegistration, -} from './backgroundCompletion'; import type { SkillFileRecord, PrimeSkillFilesResult } from './skillFiles'; import type { ArtifactDeliveryFailure } from '~/files/code'; import type { BackgroundToolResultState } from './harvest'; @@ -315,6 +316,7 @@ export interface ToolExecuteOptions { reapply?: boolean; backgroundTask?: BackgroundToolResultState; resolveBackgroundTask?: () => BackgroundToolResultState; + onFilesPersisted?: (attachments: unknown[]) => void; }) => Promise<{ attachments?: unknown[]; deliveryReady?: boolean } | null>; /** Shared ordinary-tool completion lifecycle. The delivery is registered * before invoke; settlement is persisted onto the original response row. */ @@ -345,6 +347,9 @@ export interface ToolExecuteOptions { allowUnfinished?: boolean; }) => Promise; recoverDeadClaim?: BackgroundToolDeadClaimRecovery; + /** Durable view of undelivered completions, so a status check counts results + * dispatched in earlier turns, on other replicas, or before a restart. */ + pending?: PendingBackgroundCompletionControls; }; /** Emits an `attachment` SSE event on the current request's live stream. */ emitAttachment?: (attachment: unknown) => void; @@ -6130,6 +6135,8 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand 'background tool result was not persisted', 'definite', ); + } else if (deliveryReady && !durableReceiptReady) { + completionAdmission?.expedite?.(); } } catch (persistError) { if (!durableReceiptReady) { @@ -6167,6 +6174,13 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand : { backgroundTask, resolveBackgroundTask }), output: params.output ?? localTask?.result, artifact: params.artifact, + onFilesPersisted: (attachments) => + backgroundTaskRegistry.finishHarvest( + backgroundUserId, + backgroundConversationId, + task.id, + attachments, + ), }); if (persisted == null) { /** Harvest never persisted anything (missing anchor @@ -6255,10 +6269,8 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand artifact?: unknown; status: 'completed' | 'error' | 'cancelled'; }): Promise => { - if (harvestEnabled) { - await persistBackgroundResult(params); - return; - } + /** Held for the whole persist, including a code harvest that waits for + * a long dispatch turn, so retention pressure cannot evict the task. */ backgroundTaskRegistry.markCompletionPersistencePending( backgroundUserId, backgroundConversationId, @@ -6669,6 +6681,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand subagentTasks, claimBackgroundToolResult: backgroundToolCompletion?.claim, recoverDeadBackgroundToolClaim: backgroundToolCompletion?.recoverDeadClaim, + pendingCompletions: backgroundToolCompletion?.pending, ordinaryToolCancellation, }); const taskSnapshot = getBackgroundTaskSnapshot({ diff --git a/packages/api/src/agents/harvest.lifecycle.spec.ts b/packages/api/src/agents/harvest.lifecycle.spec.ts new file mode 100644 index 00000000000..a41569f6440 --- /dev/null +++ b/packages/api/src/agents/harvest.lifecycle.spec.ts @@ -0,0 +1,156 @@ +import type { GenerationSettlementState } from '../stream/interfaces/IJobStore'; +import type { GenerationSettledListener } from '../stream/GenerationJobManager'; +import type { GenerationSettledSource } from '../stream/settled'; +import type { GenerationJobStatus } from '../types/stream'; +import type { ServerRequest } from '~/types'; +import { InMemoryEventTransport } from '../stream/implementations/InMemoryEventTransport'; +import { InMemoryJobStore } from '../stream/implementations/InMemoryJobStore'; +import { GenerationJobManagerClass } from '../stream/GenerationJobManager'; +import { createBackgroundCodeResultHandler } from './harvest'; +import { waitForGenerationSettled } from '../stream/settled'; + +function sourceForRunningTurn() { + const listeners = new Set(); + const job: { + status: GenerationJobStatus; + createdAt: number; + terminalPersistencePending?: boolean; + } = { status: 'running', createdAt: 1 }; + return { + job, + listeners, + getGenerationSettlementState: jest.fn( + async (): Promise => ({ ...job }), + ), + onGenerationSettled: (listener: GenerationSettledListener) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + settle: () => { + job.status = 'complete'; + for (const listener of [...listeners]) + listener({ + streamId: 'conversation', + conversationId: 'conversation', + userId: 'user', + status: 'complete', + }); + }, + }; +} + +function harvest(source: GenerationSettledSource, row: { unfinished: boolean }) { + const update = jest.fn(async () => ({ matched: true, unfinished: row.unfinished })); + const handler = createBackgroundCodeResultHandler({ + req: { user: { id: 'user' } } as ServerRequest, + preflightCodeOutputBatch: async () => [], + processCodeOutput: jest.fn(), + updateToolCallResult: update, + runPreviewFinalize: jest.fn(), + waitForGenerationSettled: (conversationId, options) => + waitForGenerationSettled(source, conversationId, options), + }); + let result: boolean | undefined; + const pending = handler({ + toolName: 'execute_code', + toolCallId: 'call', + stepId: 'step', + messageId: 'message', + conversationId: 'conversation', + output: 'finished code', + backgroundTask: { + taskId: 'task', + toolName: 'execute_code', + status: 'completed', + settledAt: new Date(), + completionWakeup: true, + }, + }).then((value) => { + result = value?.deliveryReady; + }); + return { update, pending, getResult: () => result }; +} + +describe('background code harvest generation lifecycle', () => { + beforeEach(() => { + jest.useFakeTimers(); + }); + afterEach(() => { + jest.clearAllTimers(); + jest.useRealTimers(); + }); + + it('keeps harvesting a live turn after a transient first status-read failure', async () => { + const source = sourceForRunningTurn(); + source.getGenerationSettlementState.mockRejectedValueOnce(new Error('transient Redis outage')); + const row = { unfinished: true }; + const running = harvest(source, row); + await jest.advanceTimersByTimeAsync(1_100_000); + const beforeSettle = running.getResult(); + expect(source.getGenerationSettlementState.mock.calls.length).toBeGreaterThan(1); + row.unfinished = false; + source.settle(); + await jest.advanceTimersByTimeAsync(0); + expect(beforeSettle).toBeUndefined(); + await running.pending; + expect(running.getResult()).toBe(true); + }); + + it('anchors the original result when a remote successor replaces the observed generation', async () => { + const source = sourceForRunningTurn(); + const row = { unfinished: true }; + const running = harvest(source, row); + await jest.advanceTimersByTimeAsync(1_000_000); + expect(running.update.mock.calls.length).toBe(14); + row.unfinished = false; + source.job.createdAt = 2; + source.job.status = 'requires_action'; + await jest.advanceTimersByTimeAsync(120_000); + const beforeSuccessorSettles = running.getResult(); + expect(running.update.mock.calls.length).toBe(15); + source.settle(); + await jest.advanceTimersByTimeAsync(0); + await running.pending; + expect(beforeSuccessorSettles).toBe(true); + }); + + it('does not retire a result while the real manager is awaiting terminal message persistence', async () => { + const manager = new GenerationJobManagerClass(); + const store = new InMemoryJobStore({ ttlAfterComplete: 0 }); + manager.configure({ + jobStore: store, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + manager.initialize(); + let finishSave!: () => void; + const save = new Promise((resolve) => { + finishSave = resolve; + }); + try { + const job = await manager.createJob('conversation', 'user', 'conversation'); + const row = { unfinished: true }; + const running = harvest(manager, row); + await jest.advanceTimersByTimeAsync(1_000_000); + const completing = manager.completeJob('conversation', 'provider error', job.createdAt, { + beforeErrorPublication: () => save, + }); + await jest.advanceTimersByTimeAsync(120_000); + const beforePersistence = running.getResult(); + expect((await store.getJob('conversation'))?.terminalPersistencePending).toBe(true); + row.unfinished = false; + finishSave(); + await completing; + await jest.advanceTimersByTimeAsync(0); + expect(beforePersistence).toBeUndefined(); + await running.pending; + expect(running.getResult()).toBe(true); + } finally { + finishSave(); + await manager.destroy(); + } + }); +}); diff --git a/packages/api/src/agents/harvest.spec.ts b/packages/api/src/agents/harvest.spec.ts index fe5c98ae0f3..8994cfbb552 100644 --- a/packages/api/src/agents/harvest.spec.ts +++ b/packages/api/src/agents/harvest.spec.ts @@ -153,3 +153,218 @@ describe('createBackgroundToolResultHandler claim ownership', () => { ); }); }); + +describe('createBackgroundCodeResultHandler long dispatch turns', () => { + const RETRY_SCHEDULE_MS = 1_000_000; + const backgroundTask = { + taskId: 'task-1', + toolName: 'execute_code', + status: 'completed' as const, + settledAt: new Date('2026-09-24T00:00:00Z'), + completionWakeup: true as const, + }; + + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + function createHandler( + updateToolCallResult: jest.Mock, + waitForGenerationSettled?: (conversationId: string) => Promise, + ) { + return createBackgroundCodeResultHandler({ + req, + preflightCodeOutputBatch: jest + .fn() + .mockResolvedValue([{ file: params.artifact.files[0], sessionId: 'storage-session' }]), + processCodeOutput: jest.fn().mockResolvedValue({ file: { file_id: 'persisted-file' } }), + updateToolCallResult, + runPreviewFinalize: jest.fn(), + ...(waitForGenerationSettled != null && { waitForGenerationSettled }), + }); + } + + it('waits for a dispatch turn that outlives the retry schedule, then anchors', async () => { + let turnRunning = true; + const updateToolCallResult = jest.fn(async () => ({ + matched: true, + unfinished: turnRunning, + })); + let settle: (settled: boolean) => void = () => undefined; + const waitForGenerationSettled = jest.fn( + () => + new Promise((resolve) => { + settle = resolve; + }), + ); + const onFilesPersisted = jest.fn(); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + + const result = handler({ ...params, backgroundTask, onFilesPersisted }); + await jest.advanceTimersByTimeAsync(RETRY_SCHEDULE_MS); + + expect(onFilesPersisted).toHaveBeenCalledWith([ + { file_id: 'persisted-file', stepId: 'step-1' }, + ]); + expect(onFilesPersisted.mock.invocationCallOrder[0]).toBeLessThan( + updateToolCallResult.mock.invocationCallOrder[0], + ); + expect(waitForGenerationSettled).toHaveBeenCalledWith('conversation-1', { + signal: expect.any(AbortSignal), + }); + const attemptsWhileRunning = updateToolCallResult.mock.calls.length; + await jest.advanceTimersByTimeAsync(60 * 60 * 1_000); + expect(updateToolCallResult).toHaveBeenCalledTimes(attemptsWhileRunning); + + turnRunning = false; + settle(true); + await jest.advanceTimersByTimeAsync(0); + + await expect(result).resolves.toEqual({ + attachments: [{ file_id: 'persisted-file', stepId: 'step-1' }], + deliveryReady: true, + }); + expect(updateToolCallResult).toHaveBeenCalledTimes(attemptsWhileRunning + 1); + }); + + it('waits for a running turn that has not saved its row yet', async () => { + let rowSaved = false; + const updateToolCallResult = jest.fn(async () => ({ matched: rowSaved, unfinished: false })); + const waitForGenerationSettled = jest.fn(async () => { + rowSaved = true; + return true; + }); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + + const result = handler({ ...params, backgroundTask }); + await jest.advanceTimersByTimeAsync(RETRY_SCHEDULE_MS); + + await expect(result).resolves.toEqual(expect.objectContaining({ deliveryReady: true })); + expect(waitForGenerationSettled).toHaveBeenCalledTimes(1); + }); + + it('gives the result up when no generation is running and the row never appears', async () => { + const updateToolCallResult = jest.fn(async () => ({ matched: false, unfinished: false })); + const waitForGenerationSettled = jest.fn(async () => false); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + + const result = handler({ ...params, backgroundTask }); + await jest.advanceTimersByTimeAsync(RETRY_SCHEDULE_MS + 60_000); + + await expect(result).resolves.toEqual(expect.objectContaining({ deliveryReady: false })); + expect(updateToolCallResult).toHaveBeenCalledTimes(14 + 5); + }); + + it('anchors a final save that lands well after the generation reads as settled', async () => { + let settledAt: number | undefined; + const updateToolCallResult = jest.fn(async () => ({ + matched: true, + unfinished: settledAt == null || Date.now() - settledAt < 40_000, + })); + const waitForGenerationSettled = jest.fn(async () => { + settledAt = Date.now(); + return true; + }); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + + const result = handler({ ...params, backgroundTask }); + await jest.advanceTimersByTimeAsync(RETRY_SCHEDULE_MS + 60_000); + + await expect(result).resolves.toEqual(expect.objectContaining({ deliveryReady: true })); + }); + + it('gives the result up when the turn ends without saving the tool call', async () => { + const updateToolCallResult = jest.fn(async () => ({ matched: true, unfinished: true })); + const waitForGenerationSettled = jest.fn(async () => true); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + + const result = handler({ ...params, backgroundTask }); + await jest.advanceTimersByTimeAsync(RETRY_SCHEDULE_MS + 60_000); + + await expect(result).resolves.toEqual(expect.objectContaining({ deliveryReady: false })); + expect(waitForGenerationSettled).toHaveBeenCalledTimes(1); + }); + + it('anchors as soon as the dispatch turn settles instead of at the next scheduled retry', async () => { + let turnRunning = true; + const updateToolCallResult = jest.fn(async () => ({ matched: true, unfinished: turnRunning })); + let settle: (settled: boolean) => void = () => undefined; + const waitForGenerationSettled = jest.fn( + () => + new Promise((resolve) => { + settle = resolve; + }), + ); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + let finished = false; + + const result = handler({ ...params, backgroundTask }).then((value) => { + finished = true; + return value; + }); + /** 45 s in, the schedule is inside a 30 s gap (next step near 69 s). */ + await jest.advanceTimersByTimeAsync(45_000); + const attemptsBeforeSettle = updateToolCallResult.mock.calls.length; + turnRunning = false; + settle(true); + await jest.advanceTimersByTimeAsync(0); + + expect(finished).toBe(true); + await expect(result).resolves.toEqual(expect.objectContaining({ deliveryReady: true })); + expect(updateToolCallResult).toHaveBeenCalledTimes(attemptsBeforeSettle + 1); + }); + + it.each([params.conversationId, 'another-conversation'])( + 'pins only the epoch owned by the host stream %s', + async (generationStreamId) => { + const wait = jest.fn(async () => true); + const handler = createBackgroundCodeResultHandler({ + req, + preflightCodeOutputBatch: async () => [], + processCodeOutput: jest.fn(), + runPreviewFinalize: jest.fn(), + updateToolCallResult: async () => ({ matched: true, unfinished: false }), + generationCreatedAt: 1234, + generationStreamId, + waitForGenerationSettled: wait, + }); + await handler({ ...params, backgroundTask }); + expect(wait).toHaveBeenCalledWith(params.conversationId, { + ...(generationStreamId === params.conversationId && { generationCreatedAt: 1234 }), + signal: expect.any(AbortSignal), + }); + }, + ); + + it('stops listening for the turn once the result is anchored', async () => { + const updateToolCallResult = jest.fn(async () => ({ matched: true, unfinished: false })); + let signal: AbortSignal | undefined; + const waitForGenerationSettled = jest.fn( + (_conversationId: string, options?: { signal?: AbortSignal }) => { + signal = options?.signal; + return new Promise(() => undefined); + }, + ); + const handler = createHandler(updateToolCallResult, waitForGenerationSettled); + + await expect(handler({ ...params, backgroundTask })).resolves.toEqual( + expect.objectContaining({ deliveryReady: true }), + ); + expect(signal?.aborted).toBe(true); + }); + + it('keeps the bounded schedule when no settlement signal is wired', async () => { + const updateToolCallResult = jest.fn(async () => ({ matched: true, unfinished: true })); + const handler = createHandler(updateToolCallResult); + + const result = handler({ ...params, backgroundTask }); + await jest.advanceTimersByTimeAsync(RETRY_SCHEDULE_MS); + + await expect(result).resolves.toEqual(expect.objectContaining({ deliveryReady: false })); + expect(updateToolCallResult).toHaveBeenCalledTimes(14); + }); +}); diff --git a/packages/api/src/agents/harvest.ts b/packages/api/src/agents/harvest.ts index 84738597ad5..63cd9188611 100644 --- a/packages/api/src/agents/harvest.ts +++ b/packages/api/src/agents/harvest.ts @@ -1,5 +1,6 @@ import { logger } from '@librechat/data-schemas'; import type { PreparedCodeOutputEntry } from '~/files/code/preflight'; +import type { GenerationSettledWaitOptions } from '../stream/settled'; import type { CodeExecutionContext } from './execution'; import type { ServerRequest } from '~/types'; @@ -12,6 +13,9 @@ import type { ServerRequest } from '~/types'; const BACKGROUND_PATCH_RETRY_DELAYS_MS = [ 250, 500, 1_000, 2_000, 5_000, 10_000, 20_000, 30_000, 60_000, 120_000, 180_000, 240_000, 300_000, ]; +/** Bounded final attempts after the dispatch epoch has finished or disappeared. + * Live terminal persistence is waited for explicitly, not inferred from this delay. */ +const SETTLED_PATCH_RETRY_DELAYS_MS = [1_000, 5_000, 15_000, 30_000]; interface HarvestFileRef { id: string; name: string; @@ -51,6 +55,10 @@ export interface BackgroundToolResultState { } export interface CodeHarvestDeps { + /** Generation identity captured by the host at dispatch, not at harvest time. */ + generationCreatedAt?: number; + /** The stream owning that epoch; nested child conversations must observe their own. */ + generationStreamId?: string | null; req: ServerRequest; /** Data-schemas method: idempotent tool-call part patch + attachment append. */ updateToolCallResult: (params: { @@ -95,6 +103,12 @@ export interface CodeHarvestDeps { fileId: string; previewRevision?: number; }) => void; + /** Waits for the dispatch epoch, including terminal persistence. Read outages + * remain retryable; replacement ends the wait without following a successor. */ + waitForGenerationSettled?: ( + conversationId: string, + options?: GenerationSettledWaitOptions, + ) => Promise; } export interface CodeHarvestParams { @@ -119,31 +133,48 @@ export interface CodeHarvestParams { /** Re-reads local claim ownership on every retry so a same-generation * manual poll cannot be overwritten by a later automatic continuation. */ resolveBackgroundTask?: () => BackgroundToolResultState; + /** Called once generated files are stored, before the row patch, which can + * wait for the dispatch turn to end: a poll in that turn needs them now. */ + onFilesPersisted?: (attachments: unknown[]) => void; } export type CodeHarvestHandler = ( params: CodeHarvestParams, ) => Promise<{ attachments: unknown[]; deliveryReady?: boolean } | null>; -const sleep = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)); +type BackgroundResultRowParams = { + userId: string; + messageId: string; + conversationId: string; + toolCallId: string; + stepId?: string; + agentId?: string; + output?: string; + attachments?: unknown[]; + backgroundTask?: BackgroundToolResultState; + resolveBackgroundTask?: () => BackgroundToolResultState; +}; -async function persistBackgroundToolResultRow( +/** Resolves `true` once `interrupt` fires, or `false` after `ms`; never both. */ +function sleepUnlessInterrupted(ms: number, interrupt?: Promise): Promise { + return new Promise((resolve) => { + const timer = setTimeout(() => resolve(false), ms); + void interrupt?.then(() => { + clearTimeout(timer); + resolve(true); + }); + }); +} + +/** `false` once `interrupt` cut the schedule short, so the caller can retry at once. */ +async function anchorBackgroundToolResultRow( updateToolCallResult: CodeHarvestDeps['updateToolCallResult'], - params: { - userId: string; - messageId: string; - conversationId: string; - toolCallId: string; - stepId?: string; - agentId?: string; - output?: string; - attachments?: unknown[]; - backgroundTask?: BackgroundToolResultState; - resolveBackgroundTask?: () => BackgroundToolResultState; - }, + params: BackgroundResultRowParams, + retryDelaysMs: readonly number[], + interrupt?: Promise, ): Promise { const { resolveBackgroundTask, ...persistedParams } = params; - for (let attempt = 0; attempt <= BACKGROUND_PATCH_RETRY_DELAYS_MS.length; attempt++) { + for (let attempt = 0; attempt <= retryDelaysMs.length; attempt++) { const currentBackgroundTask = resolveBackgroundTask?.() ?? persistedParams.backgroundTask; const result = await updateToolCallResult({ ...persistedParams, @@ -153,14 +184,74 @@ async function persistBackgroundToolResultRow( if (result.matched && !result.unfinished) { return true; } - if (attempt === BACKGROUND_PATCH_RETRY_DELAYS_MS.length) { + if (attempt === retryDelaysMs.length) { break; } - await sleep(BACKGROUND_PATCH_RETRY_DELAYS_MS[attempt]); + if (await sleepUnlessInterrupted(retryDelaysMs[attempt], interrupt)) { + return false; + } } return false; } +/** + * Patches the result onto the dispatch turn's row, which is absent or unfinished + * while that turn streams. The patch is retried on a schedule, but the turn + * settling is what actually makes the row patchable, so a settle observed + * mid-schedule retries at once instead of at the next scheduled step. A turn that + * outlives the schedule is waited for rather than given up. Only a positive + * settle cuts the schedule short: an unknown generation state keeps the schedule. + */ +async function persistBackgroundToolResultRow( + updateToolCallResult: CodeHarvestDeps['updateToolCallResult'], + params: BackgroundResultRowParams, + waitForGenerationSettled?: CodeHarvestDeps['waitForGenerationSettled'], + generationCreatedAt?: number, +): Promise { + if (waitForGenerationSettled == null) { + return anchorBackgroundToolResultRow( + updateToolCallResult, + params, + BACKGROUND_PATCH_RETRY_DELAYS_MS, + ); + } + const listening = new AbortController(); + const settlement = waitForGenerationSettled(params.conversationId, { + signal: listening.signal, + ...(generationCreatedAt == null ? {} : { generationCreatedAt }), + }).catch((error: unknown) => { + logger.warn( + `[background] Failed waiting for the dispatch turn of message ${params.messageId} to settle:`, + error, + ); + return false; + }); + const settledTurn = new Promise((resolve) => { + void settlement.then((settled) => { + if (settled) resolve(); + }); + }); + try { + const anchored = await anchorBackgroundToolResultRow( + updateToolCallResult, + params, + BACKGROUND_PATCH_RETRY_DELAYS_MS, + settledTurn, + ); + if (anchored) { + return true; + } + await settlement; + return anchorBackgroundToolResultRow( + updateToolCallResult, + params, + SETTLED_PATCH_RETRY_DELAYS_MS, + ); + } finally { + listening.abort(); + } +} + /** Persists an ordinary detached tool result without invoking code-artifact processing. */ export function createBackgroundToolResultHandler( deps: Pick, @@ -218,6 +309,9 @@ export function createBackgroundCodeResultHandler(deps: CodeHarvestDeps): CodeHa preflightCodeOutputBatch, processCodeOutput, runPreviewFinalize, + waitForGenerationSettled, + generationCreatedAt, + generationStreamId, } = deps; return async ({ toolCallId, @@ -233,6 +327,7 @@ export function createBackgroundCodeResultHandler(deps: CodeHarvestDeps): CodeHa reapply, backgroundTask, resolveBackgroundTask, + onFilesPersisted, }) => { const userId = req.user?.id; if (!userId || !messageId || !conversationId) { @@ -311,22 +406,30 @@ export function createBackgroundCodeResultHandler(deps: CodeHarvestDeps): CodeHa } } - const deliveryReady = await persistBackgroundToolResultRow(updateToolCallResult, { - userId, - messageId, - conversationId, - toolCallId, - stepId, - agentId, - output, - attachments, - ...(backgroundTask != null ? { backgroundTask } : {}), - ...(resolveBackgroundTask != null ? { resolveBackgroundTask } : {}), - }); + onFilesPersisted?.(attachments); + const deliveryReady = await persistBackgroundToolResultRow( + updateToolCallResult, + { + userId, + messageId, + conversationId, + toolCallId, + stepId, + agentId, + output, + attachments, + ...(backgroundTask != null ? { backgroundTask } : {}), + ...(resolveBackgroundTask != null ? { resolveBackgroundTask } : {}), + }, + waitForGenerationSettled, + generationStreamId == null || generationStreamId === conversationId + ? generationCreatedAt + : undefined, + ); if (!deliveryReady) { logger.warn( `[background] Could not anchor code result onto message ${messageId} (tool call ${toolCallId}); ` + - 'the dispatch turn never persisted. Poll delivery still returns the result.', + 'the dispatch turn ended without saving that tool call. Poll delivery still returns the result.', ); } return { attachments, ...(backgroundTask != null ? { deliveryReady } : {}) }; diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index 7c0533e006d..b1a305fbfa6 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -29,6 +29,8 @@ export * from './backgroundClaims'; export * from './backgroundCompletionWakeup'; export * from './initialize'; export * from './legacy'; +export * from './listing'; +export * from './listingAvatars'; export * from './lazySubagents'; export * from './lazyHistory'; export * from './memory'; diff --git a/packages/api/src/agents/listing.ts b/packages/api/src/agents/listing.ts new file mode 100644 index 00000000000..6a9680399f4 --- /dev/null +++ b/packages/api/src/agents/listing.ts @@ -0,0 +1,52 @@ +import { PermissionBits, ResourceType } from 'librechat-data-provider'; +import type { IUser, SystemCapability } from '@librechat/data-schemas'; +import { hasManageAgentsCapability } from './reads'; + +/** ACL implementations may return either plain IDs or objects with a string representation. */ +type ResourceId = string | { toString(): string }; + +type AgentListAccessDeps = { + hasCapability: (user: IUser, capability: SystemCapability) => Promise; + findAccessibleResources: (params: { + userId: string; + role?: string; + idOnTheSource?: string; + resourceType: ResourceType; + requiredPermissions: PermissionBits; + }) => Promise; +}; + +export async function getAgentListAccess( + user: IUser, + requiredPermissions: PermissionBits, + deps: AgentListAccessDeps, +): Promise<{ accessibleIds: string[] | null; editableIds: string[] | null }> { + if (typeof requiredPermissions !== 'number' || requiredPermissions < 1) { + throw new Error('requiredPermissions must be a positive number'); + } + + const params = { + userId: user.id, + role: user.role, + idOnTheSource: user.idOnTheSource, + resourceType: ResourceType.AGENT, + }; + // The normal startup path must not wait for a capability DB read before issuing ACL reads. + const aclReads = Promise.all([ + deps.findAccessibleResources({ ...params, requiredPermissions }), + (requiredPermissions & PermissionBits.EDIT) === PermissionBits.EDIT + ? null + : deps.findAccessibleResources({ ...params, requiredPermissions: PermissionBits.EDIT }), + ]); + // A manager can proceed even if a speculative ACL read rejects before the capability resolves. + void aclReads.catch(() => undefined); + if (await hasManageAgentsCapability(user, deps)) { + return { accessibleIds: null, editableIds: null }; + } + + const [accessibleIds, editableIds] = await aclReads; + return { + accessibleIds: accessibleIds.map(String), + editableIds: editableIds?.map(String) ?? null, + }; +} diff --git a/packages/api/src/agents/listingAvatars.spec.ts b/packages/api/src/agents/listingAvatars.spec.ts new file mode 100644 index 00000000000..a6735be150c --- /dev/null +++ b/packages/api/src/agents/listingAvatars.spec.ts @@ -0,0 +1,172 @@ +import { FileSources } from 'librechat-data-provider'; +import { + isFullAgentListAvatarCacheEntry, + getAgentListAvatarRefreshKey, + refreshAgentListAvatarsBeforePage, + refreshManagedAgentListPageAvatars, +} from './listingAvatars'; +import { MAX_AVATAR_REFRESH_AGENTS } from './avatars'; + +const visible = { id: 'agent-visible', avatar: { source: FileSources.s3, filepath: 'old.jpg' } }; +const nextPage = { id: 'agent-next', avatar: { source: FileSources.s3, filepath: 'next.jpg' } }; +const local = { id: 'agent-local', avatar: { source: FileSources.local, filepath: 'local.jpg' } }; + +function params() { + return { + accessibleIds: null, + agents: [visible, local], + cachedEntry: null, + refreshS3Url: jest.fn().mockResolvedValue('signed-visible.jpg'), + cacheSet: jest.fn().mockResolvedValue(undefined), + cacheKey: 'user:tenant:avatars', + ttl: 1800, + }; +} + +describe('Agent listing avatar scope', () => { + it('uses the same tenant-qualified key for listing and upload, including legacy users', () => { + expect(getAgentListAvatarRefreshKey({ id: 'alice', tenantId: 'tenant-a' })).toBe( + 'alice:tenant-a:agents_avatar_refresh', + ); + expect(getAgentListAvatarRefreshKey({ id: 'alice', tenantId: 'tenant-b' })).not.toBe( + getAgentListAvatarRefreshKey({ id: 'alice', tenantId: 'tenant-a' }), + ); + expect(getAgentListAvatarRefreshKey({ id: 'alice' })).toBe('alice::agents_avatar_refresh'); + }); + + it('does not load every tenant avatar before a manager list query', async () => { + const refreshAll = jest.fn().mockResolvedValue({ urlCache: { 'agent-visible': 'signed.jpg' } }); + await expect(refreshAgentListAvatarsBeforePage(null, null, refreshAll)).resolves.toBeNull(); + expect(refreshAll).not.toHaveBeenCalled(); + + const previous = { urlCache: { 'agent-visible': 'signed.jpg' } }; + await expect(refreshAgentListAvatarsBeforePage(null, previous, refreshAll)).resolves.toEqual( + previous, + ); + expect(refreshAll).not.toHaveBeenCalled(); + await expect(refreshAgentListAvatarsBeforePage([], null, refreshAll)).resolves.toEqual( + previous, + ); + expect(refreshAll).toHaveBeenCalledTimes(1); + expect(isFullAgentListAvatarCacheEntry(previous)).toBe(true); + expect(isFullAgentListAvatarCacheEntry({ ...previous, scope: 'page' })).toBe(false); + }); + + it('presigns only the manager search page without writing agent documents', async () => { + const options = params(); + const result = await refreshManagedAgentListPageAvatars(options); + + expect(options.refreshS3Url).toHaveBeenCalledTimes(1); + expect(options.refreshS3Url).toHaveBeenCalledWith(visible.avatar); + expect(result?.urlCache).toEqual({ 'agent-visible': 'signed-visible.jpg' }); + expect(options.cacheSet).toHaveBeenCalledWith(options.cacheKey, result, expect.any(Number)); + expect(options.cacheSet.mock.calls[0][2]).toBeLessThanOrEqual(options.ttl); + + await refreshManagedAgentListPageAvatars({ + ...options, + cachedEntry: result, + agents: [visible], + }); + expect(options.refreshS3Url).toHaveBeenCalledTimes(1); + + await refreshManagedAgentListPageAvatars({ + ...options, + cachedEntry: result, + agents: [nextPage], + }); + expect(options.refreshS3Url).toHaveBeenCalledTimes(2); + }); + + it('bounds each manager cache entry while retaining the newest entries', async () => { + const options = params(); + const all = Array.from({ length: MAX_AVATAR_REFRESH_AGENTS + 1 }, (_, index) => ({ + id: `agent-${index}`, + avatar: { source: FileSources.s3, filepath: `old-${index}.jpg` }, + })); + const first = await refreshManagedAgentListPageAvatars({ + ...options, + agents: all.slice(0, MAX_AVATAR_REFRESH_AGENTS), + }); + const next = await refreshManagedAgentListPageAvatars({ + ...options, + cachedEntry: first, + agents: all.slice(MAX_AVATAR_REFRESH_AGENTS), + }); + + expect(Object.keys(next!.urlCache)).toHaveLength(MAX_AVATAR_REFRESH_AGENTS); + expect(next!.urlCache['agent-0']).toBeUndefined(); + expect(next!.urlCache[`agent-${MAX_AVATAR_REFRESH_AGENTS}`]).toBe('signed-visible.jpg'); + expect(options.cacheSet).toHaveBeenCalledTimes(2); + }); + + it('re-signs entries after their original TTL even when later pages updated the cache', async () => { + const options = params(); + const clock = jest.spyOn(Date, 'now').mockReturnValue(1000); + try { + const first = await refreshManagedAgentListPageAvatars(options); + expect(first?.expiresAt).toBe(1000 + options.ttl); + clock.mockReturnValue(1000 + options.ttl - 1); + const second = await refreshManagedAgentListPageAvatars({ + ...options, + cachedEntry: first, + agents: [nextPage], + }); + expect(options.cacheSet.mock.calls[1][2]).toBe(1); + + clock.mockReturnValue(1000 + options.ttl); + const renewed = await refreshManagedAgentListPageAvatars({ + ...options, + cachedEntry: second, + agents: [visible], + }); + expect(renewed?.expiresAt).toBe(1000 + options.ttl * 2); + expect(options.refreshS3Url).toHaveBeenCalledTimes(3); + } finally { + clock.mockRestore(); + } + }); + + it('re-signs when another user replaces a cached avatar', async () => { + const options = params(); + const previous = await refreshManagedAgentListPageAvatars(options); + const replacement = { + id: visible.id, + avatar: { source: FileSources.s3, filepath: 'replacement.jpg' }, + }; + options.refreshS3Url.mockResolvedValueOnce('signed-replacement.jpg'); + + const refreshed = await refreshManagedAgentListPageAvatars({ + ...options, + cachedEntry: previous, + agents: [replacement], + }); + + expect(options.refreshS3Url).toHaveBeenCalledTimes(2); + expect(refreshed?.urlCache[visible.id]).toBe('signed-replacement.jpg'); + expect(refreshed?.avatarPaths?.[visible.id]).toBe('replacement.jpg'); + expect(options.cacheSet).toHaveBeenCalledTimes(2); + }); + + it('leaves an ordinary ACL viewer on the pre-query full-list refresh path', async () => { + const options = params(); + const previous = { urlCache: { 'agent-visible': 'signed.jpg' } }; + const result = await refreshManagedAgentListPageAvatars({ + ...options, + accessibleIds: ['507f1f77bcf86cd799439011'], + cachedEntry: previous, + }); + expect(result).toBe(previous); + expect(options.refreshS3Url).not.toHaveBeenCalled(); + expect(options.cacheSet).not.toHaveBeenCalled(); + }); + + it('allows retry after an S3 failure without failing the list response', async () => { + const options = params(); + options.refreshS3Url.mockRejectedValueOnce(new Error('S3 unavailable')); + const result = await refreshManagedAgentListPageAvatars(options); + expect(result?.urlCache).toEqual({}); + expect(options.cacheSet).not.toHaveBeenCalled(); + await refreshManagedAgentListPageAvatars({ ...options, cachedEntry: result }); + expect(options.cacheSet).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/api/src/agents/listingAvatars.ts b/packages/api/src/agents/listingAvatars.ts new file mode 100644 index 00000000000..8be8e3bbe02 --- /dev/null +++ b/packages/api/src/agents/listingAvatars.ts @@ -0,0 +1,145 @@ +import { logger } from '@librechat/data-schemas'; +import { FileSources } from 'librechat-data-provider'; +import type { AgentAvatar } from 'librechat-data-provider'; +import type { RefreshS3UrlFn } from './avatars'; +import { AVATAR_REFRESH_BATCH_SIZE, MAX_AVATAR_REFRESH_AGENTS } from './avatars'; + +type AvatarRefreshEntry = { + urlCache: Record; + /** Stored pathname corresponding to each signed page URL, used to detect avatar replacements. */ + avatarPaths?: Record; + scope?: 'page'; + expiresAt?: number; +}; +type ListedAgent = { id?: string; avatar?: AgentAvatar }; + +/** The list writer and avatar upload invalidator must always use the same tenant-scoped key. */ +export function getAgentListAvatarRefreshKey(user: { + id: string; + tenantId?: string | null; +}): string { + return `${user.id}:${user.tenantId ?? ''}:agents_avatar_refresh`; +} + +function validRefreshEntry(entry: unknown): entry is AvatarRefreshEntry { + return ( + entry != null && + typeof entry === 'object' && + 'urlCache' in entry && + entry.urlCache != null && + typeof entry.urlCache === 'object' && + !Array.isArray(entry.urlCache) + ); +} + +export function isFullAgentListAvatarCacheEntry(entry: unknown): boolean { + return validRefreshEntry(entry) && entry.scope !== 'page'; +} + +/** Only ACL-scoped lists use the full-set refresh, which must finish before a cursor snapshot. */ +export async function refreshAgentListAvatarsBeforePage( + accessibleIds: string[] | null, + cachedEntry: unknown, + refreshAll: () => Promise, +): Promise { + if (accessibleIds === null) { + return validRefreshEntry(cachedEntry) ? cachedEntry : null; + } + return refreshAll(); +} + +/** + * A manager has access to the whole tenant, but refreshing all of its avatars before each + * search can presign and write hundreds of unrelated agents. Refresh only the visible page. + * Do not persist these URL changes: updating `updatedAt` after taking a cursor snapshot would + * reorder later pages, and legacy users without tenant context cannot safely update by public ID. + */ +export async function refreshManagedAgentListPageAvatars({ + accessibleIds, + agents, + cachedEntry, + refreshS3Url, + cacheSet, + cacheKey, + ttl, +}: { + accessibleIds: string[] | null; + agents: ListedAgent[]; + cachedEntry: AvatarRefreshEntry | null; + refreshS3Url: RefreshS3UrlFn; + cacheSet: (key: string, entry: AvatarRefreshEntry, ttl: number) => Promise; + cacheKey: string; + ttl: number; +}): Promise { + if (accessibleIds !== null || agents.length === 0) { + return cachedEntry; + } + + const now = Date.now(); + const cachedExpiresAt = cachedEntry?.scope === 'page' ? cachedEntry.expiresAt : undefined; + const cachedPageIsFresh = typeof cachedExpiresAt === 'number' && cachedExpiresAt > now; + const urlCache = cachedPageIsFresh ? { ...cachedEntry?.urlCache } : {}; + const avatarPaths = cachedPageIsFresh ? { ...cachedEntry?.avatarPaths } : {}; + const expiresAt = cachedPageIsFresh ? cachedExpiresAt : now + ttl; + // A different user can replace an agent's avatar without clearing this viewer's cache. + // Never overlay a new avatar with a signed URL for an older stored pathname. + for (const agent of agents) { + if ( + agent.id && + Object.prototype.hasOwnProperty.call(urlCache, agent.id) && + (agent.avatar?.source !== FileSources.s3 || + !agent.avatar.filepath || + avatarPaths[agent.id] !== agent.avatar.filepath) + ) { + delete urlCache[agent.id]; + delete avatarPaths[agent.id]; + } + } + const pending = agents.filter( + (agent) => + agent.id && + agent.avatar?.source === FileSources.s3 && + agent.avatar.filepath && + !Object.prototype.hasOwnProperty.call(urlCache, agent.id), + ); + if (pending.length === 0) { + return cachedPageIsFresh ? { urlCache, avatarPaths, scope: 'page', expiresAt } : null; + } + + let changed = false; + for (let index = 0; index < pending.length; index += AVATAR_REFRESH_BATCH_SIZE) { + await Promise.all( + pending.slice(index, index + AVATAR_REFRESH_BATCH_SIZE).map(async (agent) => { + try { + const url = await refreshS3Url(agent.avatar!); + if (url && agent.id) { + urlCache[agent.id] = url; + avatarPaths[agent.id] = agent.avatar!.filepath!; + changed = true; + } + } catch (error) { + logger.warn('[AgentList] Failed to refresh visible avatar: %o', error); + } + }), + ); + } + + if (changed) { + // A manager can visit far more agents than the full-set refresh ever loads. + // Bound one user's Redis entry to the existing avatar-refresh budget. + const cacheIds = Object.keys(urlCache); + for (const id of cacheIds.slice(0, Math.max(0, cacheIds.length - MAX_AVATAR_REFRESH_AGENTS))) { + delete urlCache[id]; + delete avatarPaths[id]; + } + } + const entry: AvatarRefreshEntry = { urlCache, avatarPaths, scope: 'page', expiresAt }; + if (changed) { + try { + await cacheSet(cacheKey, entry, Math.max(1, expiresAt - Date.now())); + } catch (error) { + logger.warn('[AgentList] Failed to cache refreshed avatars: %o', error); + } + } + return entry; +} diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 918f9797821..3f1285207af 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/packages/api/src/agents/memory.spec.ts @@ -9,6 +9,8 @@ import type { Response } from 'express'; import type { ServerRequest } from '~/types'; import { processMemory, + memoryInstructions, + formatMemoryContext, createMemoryProcessor, createMemoryTool, getMemoryAgentId, @@ -159,7 +161,7 @@ describe('Memory attachment generation fencing', () => { }, }); - await runMemory([]); + await runMemory!([]); expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( 'conversation-1', @@ -797,7 +799,7 @@ describe('memory token limit guidance', () => { }, }); - await process([]); + await process!([]); const runCalls = (Run.create as jest.Mock).mock.calls; const runConfig = runCalls[runCalls.length - 1][0]; @@ -869,6 +871,38 @@ describe('buildInlineMemoryTool content filtering', () => { }); }); +describe('buildInlineMemoryTool read failure', () => { + it('refuses a limited write when current usage could not be read', async () => { + const setMemory = jest.fn(); + const getFormattedMemories = jest.fn().mockResolvedValue({ + withKeys: undefined, + withoutKeys: undefined, + totalTokens: 0, + readFailed: true, + }); + const req = { + config: { + endpoints: { [EModelEndpoint.agents]: { capabilities: [AgentCapabilities.memory] } }, + memory: { disabled: false, tokenLimit: 100 }, + }, + user: { id: 'user-1', personalization: { memories: true } }, + } as ServerRequest; + + const memoryTool = await buildInlineMemoryTool({ + toolName: 'set_memory', + req, + agent: { tools: [AgentCapabilities.memory] }, + userId: 'user-1', + memoryMethods: { setMemory, deleteMemory: jest.fn(), getFormattedMemories }, + getRoleByName: jest.fn(), + }); + + expect(memoryTool).toBeNull(); + expect(getFormattedMemories).toHaveBeenCalledTimes(1); + expect(setMemory).not.toHaveBeenCalled(); + }); +}); + describe('agentHasInlineMemoryTools', () => { it('returns false for a nullish agent', () => { expect(agentHasInlineMemoryTools(null)).toBe(false); @@ -893,6 +927,23 @@ describe('agentHasInlineMemoryTools', () => { }); }); +describe('formatMemoryContext', () => { + it('distinguishes unavailable memory from an eligible empty store', () => { + expect(formatMemoryContext(undefined)).toBeUndefined(); + expect(formatMemoryContext('')).toBe(memoryInstructions); + expect(memoryInstructions).toMatch(/persistent memory.*across conversations/i); + expect(memoryInstructions).not.toContain('automatically stores'); + expect(memoryInstructions).not.toContain('No existing memories'); + expect(memoryInstructions).not.toMatch(/librechat/i); + }); + + it('includes existing memories without losing the capability guidance', () => { + expect(formatMemoryContext('name: Danny')).toBe( + `${memoryInstructions}\n\n# Existing memory about the user:\nname: Danny`, + ); + }); +}); + describe('buildInlineMemoryContext', () => { it('loads keyed memories for an initialized inline-memory agent', async () => { const getFormattedMemories = jest.fn().mockResolvedValue({ @@ -919,6 +970,74 @@ describe('buildInlineMemoryContext', () => { }); }); + it('announces persistent memory before an inline agent has saved anything', async () => { + const getFormattedMemories = jest.fn().mockResolvedValue({ + withKeys: '', + withoutKeys: '', + totalTokens: 0, + }); + const context = await buildInlineMemoryContext({ + agent: { id: 'agent_memory', memoryToolsRegistered: true }, + req: {} as never, + userId: 'user-1', + memoryAvailable: true, + getFormattedMemories, + }); + + expect(context).toBe(memoryInstructions); + expect(context).not.toContain('# Existing memory about the user:'); + expect(getFormattedMemories).toHaveBeenCalledWith({ + userId: 'user-1', + agentId: undefined, + }); + }); + + it('does not load or announce memory when permission is denied', async () => { + const getFormattedMemories = jest.fn(); + await expect( + buildInlineMemoryContext({ + agent: { id: 'agent_memory', memoryToolsRegistered: true }, + req: {} as never, + userId: 'user-1', + memoryAvailable: false, + getFormattedMemories, + }), + ).resolves.toBe(''); + expect(getFormattedMemories).not.toHaveBeenCalled(); + }); + + it('keeps failed loads distinct from an empty store', async () => { + const getFormattedMemories = jest.fn().mockRejectedValue(new Error('read failed')); + await expect( + buildInlineMemoryContext({ + agent: { id: 'agent_memory', memoryToolsRegistered: true }, + req: {} as never, + userId: 'user-1', + memoryAvailable: true, + getFormattedMemories, + }), + ).resolves.toBe(''); + }); + + it('suppresses guidance when a real database read is marked as failed', async () => { + const getFormattedMemories = jest.fn().mockResolvedValue({ + withKeys: undefined, + withoutKeys: undefined, + totalTokens: 0, + readFailed: true, + }); + await expect( + buildInlineMemoryContext({ + agent: { id: 'agent_memory', memoryToolsRegistered: true }, + req: {} as never, + userId: 'user-1', + memoryAvailable: true, + getFormattedMemories, + }), + ).resolves.toBe(''); + expect(getFormattedMemories).toHaveBeenCalledTimes(1); + }); + it('does not load memories when inline tools are unavailable', async () => { const getFormattedMemories = jest.fn(); await expect( @@ -934,6 +1053,37 @@ describe('buildInlineMemoryContext', () => { }); }); +describe('createMemoryProcessor read failure', () => { + it('skips extraction and reuses the failed snapshot within the request', async () => { + const req = {}; + const getFormattedMemories = jest.fn().mockResolvedValue({ + withKeys: undefined, + withoutKeys: undefined, + totalTokens: 0, + readFailed: true, + }); + const result = await createMemoryProcessor({ + req, + res: { headersSent: false, write: jest.fn() } as unknown as Response, + userId: 'user-1', + messageId: 'message-1', + conversationId: 'conversation-1', + memoryMethods: { + setMemory: jest.fn(), + deleteMemory: jest.fn(), + getUserMemories: jest.fn(), + getFormattedMemories, + }, + }); + expect(result).toEqual([undefined, undefined]); + await expect( + getRequestMemories({ req, userId: 'user-1', getFormattedMemories }), + ).resolves.toMatchObject({ readFailed: true, withKeys: undefined, withoutKeys: undefined }); + expect(getFormattedMemories).toHaveBeenCalledTimes(1); + expect(Run.create).not.toHaveBeenCalled(); + }); +}); + describe('getRequestMemories caching', () => { it('memoizes per request, then re-fetches after invalidation', async () => { const getFormattedMemories = jest diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index 9e7627d9aa2..97f62332a79 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -74,7 +74,16 @@ function normalizeMemoryLLMConfig(llmConfig?: Partial): SanitizedMemo } export const memoryInstructions = - 'The system automatically stores important user information and can update or delete memories based on user requests, enabling dynamic memory management.'; + 'Persistent memory is available across conversations within the current memory scope. Saved memories, if any, are shown below. No entries shown here does not mean memory is unavailable. Use memory tools only if provided; claim a memory was saved or deleted only after the action is confirmed.'; + +export function formatMemoryContext(memory: string | undefined): string | undefined { + if (memory == null) { + return undefined; + } + return memory + ? `${memoryInstructions}\n\n# Existing memory about the user:\n${memory}` + : memoryInstructions; +} export const SET_MEMORY_TOOL_NAME = 'set_memory'; export const DELETE_MEMORY_TOOL_NAME = 'delete_memory'; @@ -505,7 +514,7 @@ export function agentHasInlineMemoryTools(agent: InlineMemoryAgent): boolean { ); } -/** Builds the existing-memory system context for an inline-memory agent. */ +/** Builds the memory system context for an inline-memory agent. */ export async function buildInlineMemoryContext({ agent, req, @@ -529,9 +538,7 @@ export async function buildInlineMemoryContext({ agentId: getMemoryAgentId(agent), getFormattedMemories, }); - return memories.withKeys - ? `${memoryInstructions}\n\n# Existing memory about the user:\n${memories.withKeys}` - : ''; + return formatMemoryContext(memories.withKeys) ?? ''; } catch (error) { logger.error('[memory] Error loading inline agent memory context', error); return ''; @@ -691,8 +698,12 @@ export async function buildInlineMemoryTool({ agentId: memoryAgentId, getFormattedMemories: memoryMethods.getFormattedMemories, }); - totalTokens = formatted?.totalTokens ?? 0; - tokenCountsByKey = formatted?.tokenCountsByKey; + /** A formatted read failure has no trustworthy usage total. */ + if (formatted.readFailed) { + return null; + } + totalTokens = formatted.totalTokens ?? 0; + tokenCountsByKey = formatted.tokenCountsByKey; } catch (error) { logger.error( '[memory] Failed to load memory token count for set_memory', @@ -1025,6 +1036,7 @@ export async function createMemoryProcessor({ messageId, memoryMethods, conversationId, + req, config = {}, filters, streamId = null, @@ -1039,6 +1051,8 @@ export async function createMemoryProcessor({ /** Agent partition; omit for the shared personal pool */ agentId?: string; memoryMethods: RequiredMemoryMethods; + /** Reuses the request-scoped formatted snapshot for the chat context. */ + req?: object; config?: MemoryConfig; filters?: FiltersConfig; streamId?: string | null; @@ -1046,27 +1060,41 @@ export async function createMemoryProcessor({ user?: IUser; tenantId?: string; }): Promise< - [ - string, - ( - messages: BaseMessage[], - inspectionMessages?: BaseMessage[], - ) => Promise<(TAttachment | null)[] | undefined>, - ] + | [undefined, undefined] + | [ + string, + ( + messages: BaseMessage[], + inspectionMessages?: BaseMessage[], + ) => Promise<(TAttachment | null)[] | undefined>, + ] > { const { validKeys, instructions, llmConfig, tokenLimit } = config; const finalInstructions = instructions || getDefaultInstructions(validKeys, tokenLimit); - const [{ withKeys, withoutKeys, totalTokens, tokenCountsByKey }, memoryEntries] = - await Promise.all([ - memoryMethods.getFormattedMemories({ - userId, - agentId, - }), - hasActivePiiPatterns(filters?.memories?.pii) - ? memoryMethods.getUserMemories({ userId, agentId }) - : Promise.resolve(undefined), - ]); + const [formatted, memoryEntries] = await Promise.all([ + (req + ? getRequestMemories({ + req, + userId, + agentId, + getFormattedMemories: memoryMethods.getFormattedMemories, + }) + : memoryMethods.getFormattedMemories({ userId, agentId }) + ).catch((error) => { + logger.error('[memory] Error loading automatic memory context', getSafeErrorMetadata(error)); + return undefined; + }), + hasActivePiiPatterns(filters?.memories?.pii) + ? memoryMethods.getUserMemories({ userId, agentId }) + : Promise.resolve(undefined), + ]); + /** Without the current memory snapshot we cannot safely seed token limits + * or assert that no previous memories exist. Skip extraction for this turn. */ + if (!formatted || formatted.readFailed) { + return [undefined, undefined]; + } + const { withKeys, withoutKeys, totalTokens, tokenCountsByKey } = formatted; return [ withoutKeys, diff --git a/packages/api/src/agents/queuedTurns.spec.ts b/packages/api/src/agents/queuedTurns.spec.ts index adb41ec79dd..b4157fd8371 100644 --- a/packages/api/src/agents/queuedTurns.spec.ts +++ b/packages/api/src/agents/queuedTurns.spec.ts @@ -843,6 +843,199 @@ describe('Agent queued-turn continuation', () => { }); describe('Agent queued-turn delivery scheduling', () => { + it('can retry initialization after a transient index setup failure', async () => { + const ensureAgentQueuedTurnIndexes = jest + .fn() + .mockRejectedValueOnce(new Error('mongo unavailable')) + .mockResolvedValue(undefined); + const scheduler = createAgentQueuedTurnScheduler({ + methods: { + ensureAgentQueuedTurnIndexes, + findQueuedTurnsNeedingDelivery: jest.fn(async () => []), + claimQueuedTurnsForAdmissionReconciliation: jest.fn(async () => []), + } as unknown as AgentQueuedTurnMethods, + enqueue: jest.fn(), + getGenerationAdmissionEvidence: async () => null, + }); + await expect(scheduler.initialize()).rejects.toThrow('mongo unavailable'); + await expect(scheduler.initialize()).resolves.toBeUndefined(); + expect(ensureAgentQueuedTurnIndexes).toHaveBeenCalledTimes(2); + await scheduler.stop(); + }); + + it('does not back off when discovery repairs reservations but returns no deliveries', async () => { + jest.useFakeTimers(); + const findQueuedTurnsNeedingDelivery = jest.fn( + async (_limit: number, activity: { found: boolean }) => { + activity.found = true; + return []; + }, + ); + const scheduler = createAgentQueuedTurnScheduler({ + methods: { + ensureAgentQueuedTurnIndexes: jest.fn(async () => undefined), + findQueuedTurnsNeedingDelivery, + claimQueuedTurnsForAdmissionReconciliation: jest.fn(async () => []), + } as unknown as AgentQueuedTurnMethods, + enqueue: jest.fn(), + getGenerationAdmissionEvidence: async () => null, + }); + try { + await scheduler.initialize(); + await jest.advanceTimersByTimeAsync(30_000); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(2); + } finally { + await scheduler.stop(); + jest.useRealTimers(); + } + }); + + it('waits for sibling lease discovery after an error before allowing shutdown or another scan', async () => { + jest.useFakeTimers(); + let release!: (turns: AgentQueuedTurnRecord[]) => void; + const findQueuedTurnsNeedingDelivery = jest + .fn() + .mockRejectedValue(new Error('mongo unavailable')); + const claimQueuedTurnsForAdmissionReconciliation = jest.fn( + () => + new Promise((resolve) => { + release = resolve; + }), + ); + const scheduler = createAgentQueuedTurnScheduler({ + methods: { + ensureAgentQueuedTurnIndexes: jest.fn(async () => undefined), + findQueuedTurnsNeedingDelivery, + claimQueuedTurnsForAdmissionReconciliation, + } as unknown as AgentQueuedTurnMethods, + enqueue: jest.fn(), + getGenerationAdmissionEvidence: async () => null, + }); + try { + const starting = scheduler.initialize(); + await jest.advanceTimersByTimeAsync(120_000); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(1); + let stopped = false; + const stopping = scheduler.stop().then(() => { + stopped = true; + }); + await jest.advanceTimersByTimeAsync(0); + expect(stopped).toBe(false); + release([]); + await Promise.all([starting, stopping]); + expect(stopped).toBe(true); + await jest.advanceTimersByTimeAsync(120_000); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(1); + expect(claimQueuedTurnsForAdmissionReconciliation).toHaveBeenCalledTimes(1); + } finally { + await scheduler.stop(); + jest.useRealTimers(); + } + }); + + it('reduces empty Mongo discovery reads without losing bounded fallback polling', async () => { + jest.useFakeTimers(); + try { + const findQueuedTurnsNeedingDelivery = jest.fn(async () => []); + const claimQueuedTurnsForAdmissionReconciliation = jest.fn(async () => []); + const scheduler = createAgentQueuedTurnScheduler({ + methods: { + ensureAgentQueuedTurnIndexes: jest.fn(async () => undefined), + findQueuedTurnsNeedingDelivery, + claimQueuedTurnsForAdmissionReconciliation, + } as unknown as AgentQueuedTurnMethods, + enqueue: jest.fn(), + getGenerationAdmissionEvidence: async () => null, + }); + await scheduler.initialize(); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(60_000); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(2); + await jest.advanceTimersByTimeAsync(120_000); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(3); + expect(claimQueuedTurnsForAdmissionReconciliation).toHaveBeenCalledTimes(3); + await scheduler.stop(); + await jest.advanceTimersByTimeAsync(120_000); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(3); + } finally { + jest.useRealTimers(); + } + }); + + it('does not treat a deferred indeterminate admission as an empty queue', async () => { + jest.useFakeTimers(); + jest.setSystemTime(new Date(NOW)); + try { + const turn = { + ...queuedTurn('queued-turn-deferred', 1), + status: 'dead' as const, + deliveryKey: 'delivery-deferred', + admissionId: 'delivery-deferred', + admissionProtocolVersion: 2 as const, + terminalReceipt: { + outcome: 'dead' as const, + settledAt: new Date(NOW), + failure: { code: 'ADMISSION_INDETERMINATE', message: 'pending evidence' }, + }, + }; + const claimQueuedTurnsForAdmissionReconciliation = jest.fn(async (input) => [ + { + ...turn, + reconciliationClaimId: input.claimId, + reconciliationClaimBy: input.claimBy, + }, + ]); + const deferAgentQueuedTurnAdmissionReconciliation = jest.fn(async () => true); + const scheduler = createAgentQueuedTurnScheduler({ + methods: { + ensureAgentQueuedTurnIndexes: jest.fn(async () => undefined), + findQueuedTurnsNeedingDelivery: jest.fn(async () => []), + claimQueuedTurnsForAdmissionReconciliation, + deferAgentQueuedTurnAdmissionReconciliation, + } as unknown as AgentQueuedTurnMethods, + enqueue: jest.fn(), + getGenerationAdmissionEvidence: async () => null, + }); + await scheduler.initialize(); + expect(deferAgentQueuedTurnAdmissionReconciliation).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(4_999); + expect(claimQueuedTurnsForAdmissionReconciliation).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1); + expect(claimQueuedTurnsForAdmissionReconciliation).toHaveBeenCalledTimes(2); + await scheduler.stop(); + } finally { + jest.useRealTimers(); + } + }); + + it('wakes recovery after direct publication fails outside its own recovery pass', async () => { + jest.useFakeTimers(); + try { + const findQueuedTurnsNeedingDelivery = jest.fn(async () => []); + const scheduler = createAgentQueuedTurnScheduler({ + methods: { + ensureAgentQueuedTurnIndexes: jest.fn(async () => undefined), + findQueuedTurnsNeedingDelivery, + claimQueuedTurnsForAdmissionReconciliation: jest.fn(async () => []), + reserveAgentQueuedTurnDelivery: jest.fn(async () => { + throw new Error('mongo unavailable'); + }), + } as unknown as AgentQueuedTurnMethods, + enqueue: jest.fn(), + getGenerationAdmissionEvidence: async () => null, + }); + await scheduler.initialize(); + await expect(scheduler.schedule(queuedTurn('queued-turn-1', 1))).rejects.toThrow( + 'mongo unavailable', + ); + await jest.advanceTimersByTimeAsync(0); + expect(findQueuedTurnsNeedingDelivery).toHaveBeenCalledTimes(2); + await scheduler.stop(); + } finally { + jest.useRealTimers(); + } + }); + function queuedTurn(id: string, sequence: number): AgentQueuedTurnRecord { const { claimId: _claimId, claimBy: _claimBy, claimUntil: _claimUntil, ...record } = claim(); return { diff --git a/packages/api/src/agents/queuedTurns.ts b/packages/api/src/agents/queuedTurns.ts index 6eeb6864f92..540e981b2d4 100644 --- a/packages/api/src/agents/queuedTurns.ts +++ b/packages/api/src/agents/queuedTurns.ts @@ -23,6 +23,7 @@ import type { AgentTriggerDeliveryFailure } from './triggers/engine'; import { getAgentTriggerIdempotencyKey, parseAgentTriggerEnvelope } from './triggers/envelope'; import { createAgentTriggerEnvelope } from './triggers/envelope'; import { AgentTriggerExecutionError } from './triggers/host'; +import { createIdleRecoveryLoop } from './recovery'; export const AGENT_QUEUED_TURN_SOURCE = 'agent-queued-turn'; const AGENT_QUEUED_TURN_EVENT = 'agent.queued-turn'; @@ -32,6 +33,7 @@ const RECONCILIATION_LEASE_MS = 2 * 60 * 1000; const RECONCILIATION_BACKOFF_BASE_MS = 5_000; const RECONCILIATION_BACKOFF_MAX_MS = 5 * 60 * 1000; const DEFAULT_RECOVERY_INTERVAL_MS = 30_000; +const DEFAULT_RECOVERY_MAX_IDLE_INTERVAL_MS = 2 * 60_000; const DEFAULT_RECOVERY_LIMIT = 100; const MAX_FAILURE_CODE_LENGTH = 128; const MAX_FAILURE_MESSAGE_LENGTH = 2048; @@ -87,7 +89,7 @@ export interface AgentQueuedTurnSchedulerDeps { } export interface AgentQueuedTurnScheduler { - initialize: () => Promise; + initialize: (options?: { maxIdleIntervalMs?: number }) => Promise; stop: () => Promise; schedule: (turn: AgentQueuedTurnRecord) => Promise; recover: () => Promise; @@ -107,7 +109,7 @@ export interface AgentQueuedTurnLifecycle { rawSource: unknown, input: AgentQueuedTurnExecutionAdmission, ) => Promise; - initialize: () => Promise; + initialize: (options?: { maxIdleIntervalMs?: number }) => Promise; stop: () => Promise; schedule: (turn: AgentQueuedTurnRecord) => Promise; cancel: ( @@ -732,10 +734,12 @@ function createAgentQueuedTurnScheduler({ recoveryIntervalMs = DEFAULT_RECOVERY_INTERVAL_MS, recoveryLimit = DEFAULT_RECOVERY_LIMIT, }: AgentQueuedTurnSchedulerDeps): AgentQueuedTurnScheduler { - let timer: NodeJS.Timeout | undefined; - let recovery: Promise | undefined; + let loop: ReturnType | undefined; + let initialization: Promise | undefined; + let stopped = false; + let recovery: Promise<{ repaired: number; idle: boolean }> | undefined; - const schedule = async (turn: AgentQueuedTurnRecord): Promise => { + const publish = async (turn: AgentQueuedTurnRecord): Promise => { const envelope = deliveryEnvelope(turn); const deliveryKey = getAgentTriggerIdempotencyKey(envelope); const reserved = await runAsSystem(() => @@ -776,25 +780,43 @@ function createAgentQueuedTurnScheduler({ return receipt.deliveryKey; }; - const recover = (): Promise => { + const schedule = async (turn: AgentQueuedTurnRecord): Promise => { + try { + return await publish(turn); + } catch (error) { + loop?.wake(); + throw error; + } + }; + + const recoverPass = (): Promise<{ repaired: number; idle: boolean }> => { if (recovery != null) { return recovery; } const task = (async () => { const reconciliationNow = new Date(); const reconciliationClaimId = randomUUID(); - const [turns, quarantined] = await Promise.all([ - runAsSystem(() => methods.findQueuedTurnsNeedingDelivery(recoveryLimit)), + const activity = { found: false }; + const discoveries = await Promise.allSettled([ + runAsSystem(() => methods.findQueuedTurnsNeedingDelivery(recoveryLimit, activity)), runAsSystem(() => - methods.claimQueuedTurnsForAdmissionReconciliation({ - claimId: reconciliationClaimId, - claimBy: PROCESS_CLAIM_OWNER, - now: reconciliationNow, - leaseUntil: new Date(reconciliationNow.getTime() + RECONCILIATION_LEASE_MS), - limit: recoveryLimit, - }), + methods.claimQueuedTurnsForAdmissionReconciliation( + { + claimId: reconciliationClaimId, + claimBy: PROCESS_CLAIM_OWNER, + now: reconciliationNow, + leaseUntil: new Date(reconciliationNow.getTime() + RECONCILIATION_LEASE_MS), + limit: recoveryLimit, + }, + activity, + ), ), ]); + // A rejected discovery must not release the single-flight guard while + // its sibling is still writing reconciliation leases in Mongo. + const [deliveries, reconciliations] = discoveries; + const turns = deliveries.status === 'fulfilled' ? deliveries.value : []; + const quarantined = reconciliations.status === 'fulfilled' ? reconciliations.value : []; let repaired = 0; for (const turn of quarantined) { const deliveryKey = turn.deliveryKey; @@ -805,8 +827,11 @@ function createAgentQueuedTurnScheduler({ ) { continue; } - const defer = () => - runAsSystem(() => + const defer = async () => { + const availableAt = new Date( + Date.now() + reconciliationBackoff(turn.reconciliationAttempts), + ); + const deferred = await runAsSystem(() => methods.deferAgentQueuedTurnAdmissionReconciliation({ user: turn.user, ...(turn.tenantId != null && { tenantId: turn.tenantId }), @@ -815,11 +840,13 @@ function createAgentQueuedTurnScheduler({ deliveryKey, claimId: reconciliationClaimId, claimBy: PROCESS_CLAIM_OWNER, - availableAt: new Date( - Date.now() + reconciliationBackoff(turn.reconciliationAttempts), - ), + availableAt, }), ); + if (deferred) { + loop?.noteEligibleAt(availableAt); + } + }; try { const isIndeterminate = turn.terminalReceipt?.outcome === 'dead' && @@ -894,7 +921,7 @@ function createAgentQueuedTurnScheduler({ } for (const turn of turns) { try { - await schedule(turn); + await publish(turn); repaired += 1; } catch (error) { logger.warn( @@ -904,7 +931,11 @@ function createAgentQueuedTurnScheduler({ ); } } - return repaired; + // Do not strand leases already acquired by the successful discovery just + // because its independent sibling failed. Process them, then report failure. + if (deliveries.status === 'rejected') throw deliveries.reason; + if (reconciliations.status === 'rejected') throw reconciliations.reason; + return { repaired, idle: !activity.found && turns.length === 0 && quarantined.length === 0 }; })(); recovery = task; void task.then( @@ -921,28 +952,48 @@ function createAgentQueuedTurnScheduler({ ); return task; }; + const recover = async (): Promise => (await recoverPass()).repaired; return { schedule, recover, - initialize: async () => { - await runAsSystem(() => methods.ensureAgentQueuedTurnIndexes()); - timer = setInterval(() => { - void recover().catch((error: unknown) => { - logger.warn('[agentQueuedTurns] Delivery recovery pass failed', error); + initialize: (options = {}) => { + if (initialization != null) { + return initialization; + } + if (stopped) { + return Promise.resolve(); + } + const maxIdleIntervalMs = + options.maxIdleIntervalMs ?? + Math.max(recoveryIntervalMs, DEFAULT_RECOVERY_MAX_IDLE_INTERVAL_MS); + const next = (async () => { + loop = createIdleRecoveryLoop({ + intervalMs: recoveryIntervalMs, + maxIdleIntervalMs, + scan: async () => (await recoverPass()).idle, + onError: (error) => + logger.warn('[agentQueuedTurns] Delivery recovery pass failed', error), }); - }, recoveryIntervalMs); - timer.unref?.(); - await recover().catch((error: unknown) => { - logger.warn('[agentQueuedTurns] Initial delivery recovery pass failed', error); + await runAsSystem(() => methods.ensureAgentQueuedTurnIndexes()); + if (!stopped) { + await loop.start(); + } + })(); + initialization = next.catch((error: unknown) => { + if (!stopped) { + initialization = undefined; + loop = undefined; + } + throw error; }); + return initialization; }, stop: async () => { - if (timer != null) { - clearInterval(timer); - timer = undefined; - } - await recovery; + stopped = true; + await loop?.stop(); + await initialization?.catch(() => undefined); + await recovery?.catch(() => undefined); }, }; } diff --git a/packages/api/src/agents/reads.spec.ts b/packages/api/src/agents/reads.spec.ts index 1783103ccf6..cc6e6d4fe7e 100644 --- a/packages/api/src/agents/reads.spec.ts +++ b/packages/api/src/agents/reads.spec.ts @@ -10,6 +10,7 @@ import type { IRole, IUser } from '@librechat/data-schemas'; import type { Request, Response } from 'express'; import type { AgentManagementReadDeps } from './reads'; import { createAgentManagementReadHandlers } from './reads'; +import { getAgentListAccess } from './listing'; jest.mock('@librechat/data-schemas', () => { const actual = jest.requireActual('@librechat/data-schemas'); @@ -70,6 +71,119 @@ function makeDeps(overrides: Partial = {}): AgentManage }; } +describe('Agent list access', () => { + it('rejects invalid permission masks before checking the manager bypass', async () => { + const deps = makeDeps({ hasCapability: jest.fn().mockResolvedValue(true) }); + + await expect(getAgentListAccess(user, 0 as PermissionBits, deps)).rejects.toThrow( + 'requiredPermissions must be a positive number', + ); + expect(deps.hasCapability).not.toHaveBeenCalled(); + expect(deps.findAccessibleResources).not.toHaveBeenCalled(); + }); + + it('does not use speculative ACL results for a manager', async () => { + const deps = makeDeps({ hasCapability: jest.fn().mockResolvedValue(true) }); + + const access = await getAgentListAccess(user, PermissionBits.VIEW, deps); + + expect(deps.hasCapability).toHaveBeenCalledWith(user, SystemCapabilities.MANAGE_AGENTS); + expect(deps.findAccessibleResources).toHaveBeenCalledTimes(2); + expect(access).toEqual({ accessibleIds: null, editableIds: null }); + }); + + it('uses the ACL for ordinary viewers and marks only editable agents', async () => { + const editableId = new Types.ObjectId(); + const deps = makeDeps({ + findAccessibleResources: jest + .fn() + .mockImplementation(({ requiredPermissions }) => + Promise.resolve( + requiredPermissions === PermissionBits.EDIT ? [editableId] : [objectId, editableId], + ), + ), + }); + + const access = await getAgentListAccess(user, PermissionBits.VIEW, deps); + + expect(access).toEqual({ + accessibleIds: [objectId.toString(), editableId.toString()], + editableIds: [editableId.toString()], + }); + expect(deps.findAccessibleResources).toHaveBeenCalledWith({ + userId: user.id, + role: user.role, + idOnTheSource: user.idOnTheSource, + resourceType: ResourceType.AGENT, + requiredPermissions: PermissionBits.VIEW, + }); + expect(deps.findAccessibleResources).toHaveBeenCalledWith( + expect.objectContaining({ requiredPermissions: PermissionBits.EDIT }), + ); + }); + + it('falls back to the ACL rather than bypassing on capability-check failure', async () => { + const deps = makeDeps({ + hasCapability: jest.fn().mockRejectedValue(new Error('lookup failed')), + }); + + const access = await getAgentListAccess(user, PermissionBits.VIEW, deps); + + expect(access).toEqual({ + accessibleIds: [objectId.toString()], + editableIds: [objectId.toString()], + }); + expect(deps.findAccessibleResources).toHaveBeenCalledTimes(2); + }); + + it('starts ACL queries while capability lookup is still pending', async () => { + let resolveCapability!: (value: boolean) => void; + const capability = new Promise((resolve) => { + resolveCapability = resolve; + }); + const deps = makeDeps({ hasCapability: jest.fn().mockReturnValue(capability) }); + + const accessPromise = getAgentListAccess(user, PermissionBits.VIEW, deps); + expect(deps.findAccessibleResources).toHaveBeenCalledTimes(2); + resolveCapability(false); + await expect(accessPromise).resolves.toEqual({ + accessibleIds: [objectId.toString()], + editableIds: [objectId.toString()], + }); + }); + + it('does not require the speculative ACL reads to succeed for a manager', async () => { + const deps = makeDeps({ + hasCapability: jest.fn().mockResolvedValue(true), + findAccessibleResources: jest.fn().mockRejectedValue(new Error('ACL unavailable')), + }); + + await expect(getAgentListAccess(user, PermissionBits.VIEW, deps)).resolves.toEqual({ + accessibleIds: null, + editableIds: null, + }); + }); + + it('propagates ACL lookup errors for a user without management capability', async () => { + const deps = makeDeps({ + findAccessibleResources: jest.fn().mockRejectedValue(new Error('ACL unavailable')), + }); + + await expect(getAgentListAccess(user, PermissionBits.VIEW, deps)).rejects.toThrow( + 'ACL unavailable', + ); + }); + + it('does not repeat the EDIT lookup for an EDIT-scoped list', async () => { + const deps = makeDeps(); + + const access = await getAgentListAccess(user, PermissionBits.EDIT, deps); + + expect(access).toEqual({ accessibleIds: [objectId.toString()], editableIds: null }); + expect(deps.findAccessibleResources).toHaveBeenCalledTimes(1); + }); +}); + describe('Agent Management read handlers', () => { it('lists only ACL-discovered records in the authenticated tenant', async () => { const deps = makeDeps(); diff --git a/packages/api/src/agents/reads.ts b/packages/api/src/agents/reads.ts index 737094ed7c5..c1cfa111c2a 100644 --- a/packages/api/src/agents/reads.ts +++ b/packages/api/src/agents/reads.ts @@ -90,7 +90,10 @@ async function canViewAgent( }); } -async function hasManageAgentsCapability(user: IUser, deps: AgentManagementReadDeps) { +export async function hasManageAgentsCapability( + user: IUser, + deps: Pick, +): Promise { const capability = ResourceCapabilityMap[ResourceType.AGENT]; try { if (capability != null && (await deps.hasCapability(user, capability))) { diff --git a/packages/api/src/agents/recovery.spec.ts b/packages/api/src/agents/recovery.spec.ts new file mode 100644 index 00000000000..506d8c0625e --- /dev/null +++ b/packages/api/src/agents/recovery.spec.ts @@ -0,0 +1,233 @@ +import { createIdleRecoveryLoop } from './recovery'; + +describe('durable agent idle recovery', () => { + beforeEach(() => { + jest.useFakeTimers(); + jest.setSystemTime(new Date('2026-09-01T00:00:00Z')); + }); + + afterEach(() => jest.useRealTimers()); + + it('backs off only after empty passes and wakes without waiting for the idle ceiling', async () => { + const scan = jest.fn(async () => true); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + await loop.start(); + expect(scan).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(59_999); + expect(scan).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1); + expect(scan).toHaveBeenCalledTimes(2); + await jest.advanceTimersByTimeAsync(120_000); + expect(scan).toHaveBeenCalledTimes(3); + loop.wake(); + await jest.advanceTimersByTimeAsync(0); + expect(scan).toHaveBeenCalledTimes(4); + await jest.advanceTimersByTimeAsync(59_999); + expect(scan).toHaveBeenCalledTimes(4); + await loop.stop(); + }); + + it('never lets an in-flight empty result discard a wake or start overlapping scans', async () => { + let release!: (idle: boolean) => void; + const scan = jest.fn( + () => + new Promise((resolve) => { + release = resolve; + }), + ); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + const first = loop.start(); + await jest.advanceTimersByTimeAsync(0); + loop.wake(); + loop.wake(); + expect(scan).toHaveBeenCalledTimes(1); + release(true); + await first; + await jest.advanceTimersByTimeAsync(0); + expect(scan).toHaveBeenCalledTimes(2); + release(true); + await jest.advanceTimersByTimeAsync(0); + await jest.advanceTimersByTimeAsync(30_000); + expect(scan).toHaveBeenCalledTimes(2); + await jest.advanceTimersByTimeAsync(30_000); + expect(scan).toHaveBeenCalledTimes(3); + release(false); + await loop.stop(); + }); + + it('scans at the base cadence after a failed or nonempty pass', async () => { + const onError = jest.fn(); + const scan = jest + .fn() + .mockResolvedValueOnce(true) + .mockRejectedValueOnce(new Error('mongo unavailable')) + .mockResolvedValueOnce(false) + .mockResolvedValue(true); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError, + }); + await loop.start(); + await jest.advanceTimersByTimeAsync(60_000); + expect(onError).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(30_000); + expect(scan).toHaveBeenCalledTimes(3); + await jest.advanceTimersByTimeAsync(30_000); + expect(scan).toHaveBeenCalledTimes(4); + await loop.stop(); + }); + + it('scans at known eligibility times rather than sleeping past deferred work', async () => { + const scan = jest.fn(async () => true); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + await loop.start(); + const now = Date.now(); + loop.noteEligibleAt(new Date(now + 14_000)); + loop.noteEligibleAt(new Date(now + 5_000)); + await jest.advanceTimersByTimeAsync(5_000); + expect(scan).toHaveBeenCalledTimes(2); + await jest.advanceTimersByTimeAsync(9_000); + expect(scan).toHaveBeenCalledTimes(3); + await loop.stop(); + }); + + it('discovers work written by another replica without a process-local wake', async () => { + let hasWork = false; + const scan = jest.fn(async () => !hasWork); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + await loop.start(); + await jest.advanceTimersByTimeAsync(60_000); + expect(scan).toHaveBeenCalledTimes(2); + hasWork = true; + await jest.advanceTimersByTimeAsync(119_999); + expect(scan).toHaveBeenCalledTimes(2); + await jest.advanceTimersByTimeAsync(1); + expect(scan).toHaveBeenCalledTimes(3); + await jest.advanceTimersByTimeAsync(30_000); + expect(scan).toHaveBeenCalledTimes(4); + await loop.stop(); + }); + + it('never postpones the bounded fallback scan when a later deadline is learned', async () => { + const scan = jest.fn(async () => true); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + await loop.start(); + await jest.advanceTimersByTimeAsync(59_000); + loop.noteEligibleAt(new Date(Date.now() + 120_000)); + await jest.advanceTimersByTimeAsync(1_000); + expect(scan).toHaveBeenCalledTimes(2); + await loop.stop(); + }); + + it.each([false, true])( + 'rescans for a deadline that expires in-flight (stop: %s)', + async (stop) => { + let release!: (idle: boolean) => void; + const scan = jest + .fn, []>() + .mockImplementationOnce( + () => + new Promise((resolve) => { + release = resolve; + }), + ) + .mockResolvedValue(true); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + const starting = loop.start(); + await jest.advanceTimersByTimeAsync(0); + loop.noteEligibleAt(new Date(Date.now() + 5_000)); + loop.noteEligibleAt(new Date(Date.now() + 20_000)); + await jest.advanceTimersByTimeAsync(6_000); + expect(scan).toHaveBeenCalledTimes(1); + const stopping = stop ? loop.stop() : undefined; + release(false); + await starting; + await stopping; + await jest.advanceTimersByTimeAsync(0); + expect(scan).toHaveBeenCalledTimes(stop ? 1 : 2); + await jest.advanceTimersByTimeAsync(14_000); + expect(scan).toHaveBeenCalledTimes(stop ? 1 : 3); + await loop.stop(); + expect(jest.getTimerCount()).toBe(0); + }, + ); + + it('does not rescan or reset its idle timer when start is called twice', async () => { + const scan = jest.fn(async () => true); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + await loop.start(); + await jest.advanceTimersByTimeAsync(59_000); + await loop.start(); + expect(scan).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1_000); + expect(scan).toHaveBeenCalledTimes(2); + await loop.stop(); + }); + + it('waits for its active scan on shutdown and never rearms a timer', async () => { + let release!: (value: boolean) => void; + const scan = jest.fn( + () => + new Promise((resolve) => { + release = resolve; + }), + ); + const loop = createIdleRecoveryLoop({ + intervalMs: 30_000, + maxIdleIntervalMs: 120_000, + scan, + onError: jest.fn(), + }); + const started = loop.start(); + await jest.advanceTimersByTimeAsync(0); + let stopped = false; + const stopping = loop.stop().then(() => { + stopped = true; + }); + loop.wake(); + expect(stopped).toBe(false); + release(true); + await Promise.all([started, stopping]); + await jest.advanceTimersByTimeAsync(300_000); + expect(stopped).toBe(true); + expect(scan).toHaveBeenCalledTimes(1); + expect(jest.getTimerCount()).toBe(0); + }); +}); diff --git a/packages/api/src/agents/recovery.ts b/packages/api/src/agents/recovery.ts new file mode 100644 index 00000000000..93d3ca1b5de --- /dev/null +++ b/packages/api/src/agents/recovery.ts @@ -0,0 +1,164 @@ +export interface IdleRecoveryLoopOptions { + intervalMs: number; + maxIdleIntervalMs: number; + /** True only after a completed scan confirmed no eligible work and no failures. */ + scan: () => Promise; + onError: (error: unknown) => void; +} + +/** Mongo remains the authority: wakes are only hints, and every replica still scans + * within the idle ceiling so a crashed producer or a missed local wake is recovered. */ +export function createIdleRecoveryLoop({ + intervalMs, + maxIdleIntervalMs, + scan, + onError, +}: IdleRecoveryLoopOptions) { + if (!Number.isSafeInteger(intervalMs) || intervalMs <= 0) { + throw new TypeError('Recovery interval must be a positive integer'); + } + if (!Number.isSafeInteger(maxIdleIntervalMs) || maxIdleIntervalMs < intervalMs) { + throw new TypeError('Maximum idle recovery interval must be at least the recovery interval'); + } + + let started = false; + let stopped = false; + let idleStreak = 0; + let nextPollAt = 0; + let wakePending = false; + let generation = 0; + let timer: NodeJS.Timeout | undefined; + let active: Promise | undefined; + const deadlines: number[] = []; + const MAX_DEADLINES = 64; + + const schedule = (): void => { + if (!started || stopped || active != null) { + return; + } + if (timer != null) { + clearTimeout(timer); + } + const now = Date.now(); + let delay = Math.max(0, nextPollAt - now); + if (deadlines.length > 0) { + delay = Math.min(delay, Math.max(0, deadlines[0] - now)); + } + timer = setTimeout(() => { + timer = undefined; + void runNow(); + }, delay); + timer.unref?.(); + }; + + const runNow = (): Promise => { + if (stopped || !started) { + return Promise.resolve(); + } + if (active != null) { + return active; + } + if (timer != null) { + clearTimeout(timer); + timer = undefined; + } + // This scan can cover deadlines already due when it starts. Deadlines + // learned or becoming due while it runs must survive for a follow-up scan. + const now = Date.now(); + while (deadlines.length > 0 && deadlines[0] <= now) { + deadlines.shift(); + } + const before = generation; + const current = Promise.resolve() + .then(scan) + .then( + (idle) => { + idleStreak = idle && before === generation ? idleStreak + 1 : 0; + }, + (error: unknown) => { + idleStreak = 0; + onError(error); + }, + ) + .finally(() => { + if (active === current) { + active = undefined; + } + if (stopped) { + return; + } + nextPollAt = + Date.now() + Math.min(intervalMs * 2 ** Math.min(idleStreak, 30), maxIdleIntervalMs); + if (wakePending) { + wakePending = false; + queueMicrotask(() => void runNow()); + } else { + schedule(); + } + }); + active = current; + return current; + }; + + return { + start: (): Promise => { + if (stopped) { + return Promise.resolve(); + } + if (started) { + return active ?? Promise.resolve(); + } + started = true; + return runNow(); + }, + /** Coalesce notifications while a scan is in progress; none can be erased + * by that scan's earlier empty result. */ + wake: (): void => { + if (stopped) { + return; + } + generation += 1; + idleStreak = 0; + if (active != null) { + wakePending = true; + } else if (started) { + void runNow(); + } + }, + noteEligibleAt: (at: Date): void => { + const time = at.getTime(); + if (!Number.isFinite(time) || stopped) { + return; + } + if (time <= Date.now()) { + if (started) { + generation += 1; + idleStreak = 0; + if (active != null) { + wakePending = true; + } else { + void runNow(); + } + } + return; + } + const index = deadlines.findIndex((value) => value >= time); + if (index !== -1 && deadlines[index] === time) { + return; + } + deadlines.splice(index < 0 ? deadlines.length : index, 0, time); + if (deadlines.length > MAX_DEADLINES) { + deadlines.pop(); + } + schedule(); + }, + stop: async (): Promise => { + stopped = true; + if (timer != null) { + clearTimeout(timer); + timer = undefined; + } + await active; + }, + }; +} diff --git a/packages/api/src/agents/subagentCompletionWakeup.spec.ts b/packages/api/src/agents/subagentCompletionWakeup.spec.ts index d2fd1fc95eb..0f096275573 100644 --- a/packages/api/src/agents/subagentCompletionWakeup.spec.ts +++ b/packages/api/src/agents/subagentCompletionWakeup.spec.ts @@ -1,3 +1,4 @@ +import type { CodeApprovalMode } from 'librechat-data-provider'; import type { IMessage } from '@librechat/data-schemas'; import type { AgentContinueTriggerEnvelope } from './triggers/envelope'; import type { SubagentTaskWakeupRegistration } from './subagentThreads'; @@ -59,7 +60,7 @@ describe('createSubagentCompletionWakeupHandler', () => { const enqueue = enqueueMock(); const notify = createSubagentCompletionWakeupHandler(enqueue); - await notify(registration()); + await expect(notify(registration())).resolves.toBe(true); expect(enqueue).toHaveBeenCalledTimes(1); const [envelopeValue, options] = enqueue.mock.calls[0]!; @@ -111,7 +112,7 @@ describe('createSubagentCompletionWakeupHandler', () => { const enqueue = enqueueMock(); const notify = createSubagentCompletionWakeupHandler(enqueue); - await notify(registration({ parentAgentId: undefined })); + await expect(notify(registration({ parentAgentId: undefined }))).resolves.toBe(false); expect(enqueue).not.toHaveBeenCalled(); }); @@ -153,7 +154,7 @@ function wakeupEnvelope(): AgentContinueTriggerEnvelope { return envelope; } -function resolverMethods() { +function resolverMethods(codeApprovalMode?: CodeApprovalMode) { const subagentTask: IMessage['subagentTask'] = { attemptKey: 'attempt-1', parentRunId: 'response-1', @@ -173,7 +174,11 @@ function resolverMethods() { const methods = { getConvo: jest.fn(async (_userId: string, conversationId: string) => conversationId === 'conversation-1' - ? { conversationId, tenantId: 'tenant-1' } + ? { + conversationId, + tenantId: 'tenant-1', + ...(codeApprovalMode != null && { codeApprovalMode }), + } : { conversationId, tenantId: 'tenant-1', @@ -259,6 +264,27 @@ function orchestrationSnapshot( } describe('createSubagentCompletionWakeupResolver', () => { + it.each([undefined, 'ask', 'acceptEdits', 'fullAccess'] as const)( + 'inherits parent approval mode %s rather than child or event permissions', + async (mode) => { + const { methods } = resolverMethods(mode); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + const delivery = wakeupEnvelope(); + if (delivery.event.payload == null || typeof delivery.event.payload !== 'object') { + throw new Error('Expected a completion payload'); + } + delivery.event.payload = { ...delivery.event.payload, codeApprovalMode: 'fullAccess' }; + + const prepared = await resolve(delivery, { idempotencyKey: 'delivery-1' }); + + expect(prepared?.status).toBe('ready'); + expect(prepared?.status === 'ready' && prepared.codeApprovalMode).toBe(mode); + }, + ); + it('defers without claiming while the parent generation is active', async () => { const { methods } = resolverMethods(); const resolve = createSubagentCompletionWakeupResolver({ @@ -1261,6 +1287,59 @@ describe('createSubagentCompletionWakeupResolver', () => { expect(snapshot.value.note).toContain('Do not infer that no other children ran'); }); + it('backs off by waiting age while the child task is still running', async () => { + const { methods } = resolverMethods(); + methods.getMessages.mockImplementation(async (filter: { conversationId: string }) => + filter.conversationId === 'conversation-1' + ? [ + { + messageId: 'response-1', + parentMessageId: 'user-1', + isCreatedByUser: false, + createdAt: new Date(NOW - 30), + }, + ] + : [{ messageId: 'task-1:user', conversationId: 'thread-1', isCreatedByUser: true }], + ); + const resolverAt = (offsetMs: number) => + createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + now: () => NOW + offsetMs, + }); + + await expect( + resolverAt(5_000)(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ code: 'CHILD_NOT_READY', retryAfter: '5' }); + await expect( + resolverAt(90_000)(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ code: 'CHILD_NOT_READY', retryAfter: '9' }); + await expect( + resolverAt(20 * 60_000)(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ code: 'CHILD_NOT_READY', retryAfter: '60' }); + await expect( + createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + now: () => NOW + 20 * 60_000, + getWaitMaxIntervalMs: () => 30_000, + })(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ code: 'CHILD_NOT_READY', retryAfter: '30' }); + }); + + it('backs off by waiting age while the parent generation keeps running', async () => { + const { methods } = resolverMethods(); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ status: 'running' }), + now: () => NOW + 150_000, + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ code: 'PARENT_NOT_READY', retryAfter: '15' }); + }); + it('dead-letters a child whose process disappeared after the task timeout grace', async () => { const { methods } = resolverMethods(); methods.getMessages.mockImplementation(async (filter: { conversationId: string }) => diff --git a/packages/api/src/agents/subagentCompletionWakeup.ts b/packages/api/src/agents/subagentCompletionWakeup.ts index 612e22c0c0a..0b70b205a21 100644 --- a/packages/api/src/agents/subagentCompletionWakeup.ts +++ b/packages/api/src/agents/subagentCompletionWakeup.ts @@ -9,6 +9,7 @@ import type { SubagentTaskWakeupRegistration } from './subagentThreads'; import type { AgentContinueTriggerEnvelope } from './triggers/envelope'; import type { AgentTriggerDispatchContext } from './triggers/dispatch'; import type { AgentTriggerEnqueueOptions } from './triggers/delivery'; +import { WAITING_RETRY_CAP_MS, waitingRetryAfter } from './triggers/backoff'; import { boundedSubagentTaskResult } from './subagentTaskRouting'; import { createAgentTriggerEnvelope } from './triggers/envelope'; import { AgentTriggerExecutionError } from './triggers/host'; @@ -81,6 +82,8 @@ export interface SubagentCompletionWakeupResolverDeps { methods: WakeupMethods; getGenerationJob: (conversationId: string) => Promise; now?: () => number; + /** Longest a waiting delivery re-checks readiness; the backoff default otherwise. */ + getWaitMaxIntervalMs?: () => number | undefined; } function payloadRegistration( @@ -139,6 +142,12 @@ function isParentActive(job: GenerationState | null): boolean { ); } +/** A running or approval-paused parent can stay busy for hours; one that has + * settled and is only finishing terminal persistence clears within moments. */ +function isParentWorking(job: GenerationState | null): boolean { + return job?.status === 'running' || job?.status === 'requires_action'; +} + function sameTenant(actual: string | undefined, expected: string | undefined): boolean { return actual === expected; } @@ -579,9 +588,12 @@ export function createSubagentCompletionWakeupResolver({ methods, getGenerationJob, now = Date.now, + getWaitMaxIntervalMs, }: SubagentCompletionWakeupResolverDeps): NonNullable< AgentTriggerExecutionHostDeps['prepareContinue'] > { + const waitingRetry = (receivedAt: number): string => + waitingRetryAfter(receivedAt, now(), getWaitMaxIntervalMs?.() ?? WAITING_RETRY_CAP_MS); return async ( envelope: AgentContinueTriggerEnvelope, context: AgentTriggerDispatchContext, @@ -616,7 +628,7 @@ export function createSubagentCompletionWakeupResolver({ code: 'PARENT_NOT_READY', retryable: true, status: 409, - retryAfter: '1', + retryAfter: isParentWorking(parentJob) ? waitingRetry(envelope.receivedAt) : '1', deferWithoutAttempt: true, }); } @@ -700,7 +712,7 @@ export function createSubagentCompletionWakeupResolver({ code: 'CHILD_NOT_READY', retryable: true, status: 409, - retryAfter: '1', + retryAfter: waitingRetry(envelope.receivedAt), deferWithoutAttempt: true, }); } @@ -778,6 +790,7 @@ export function createSubagentCompletionWakeupResolver({ return { status: 'ready', parentMessageId, + ...(parent.codeApprovalMode != null && { codeApprovalMode: parent.codeApprovalMode }), input: renderWakeupInput(registration, resultTaskId, claim.message, orchestrationSnapshot), releaseOnDefiniteFailure: async () => { await methods.releaseSubagentTaskResultClaim({ @@ -797,11 +810,11 @@ export function createSubagentCompletionWakeupResolver({ * simply defers until the terminal child message exists. */ export function createSubagentCompletionWakeupHandler( enqueue: EnqueueAgentTrigger, -): (registration: SubagentTaskWakeupRegistration) => Promise { +): (registration: SubagentTaskWakeupRegistration) => Promise { return async (registration) => { const parentAgentId = registration.parentAgentId?.trim(); if (parentAgentId == null || parentAgentId === '' || isEphemeralAgentId(parentAgentId)) { - return; + return false; } const eventId = registration.taskId; const envelope = createAgentTriggerEnvelope({ @@ -837,5 +850,6 @@ export function createSubagentCompletionWakeupHandler( Math.max(Date.now(), registration.createdAt) + WAKEUP_ADMISSION_DELAY_MS, ), }); + return true; }; } diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index 32daa78c447..410261e7113 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -431,6 +431,73 @@ describe('SubagentThreadTaskStore', () => { }); }); + it('announces a settled child only after its terminal message is durable', async () => { + const userId = 'settled-child-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const saveMessage = jest.spyOn(methods, 'saveMessage'); + const terminalSavedAtSettle: boolean[] = []; + const store = new SubagentThreadTaskStore(methods, { + onTaskPrepared: jest.fn(), + onTaskSettled: (settledUserId, settledConversationId, taskIds) => { + expect(taskIds).toHaveLength(1); + expect( + saveMessage.mock.calls.some( + ([, message]) => + (message as { messageId?: string }).messageId === `${taskIds[0]}:assistant`, + ), + ).toBe(true); + expect(settledUserId).toBe(userId); + expect(settledConversationId).toBe(parentConversationId); + terminalSavedAtSettle.push( + saveMessage.mock.calls.some(([, message]) => + String((message as { messageId?: string }).messageId).endsWith(':assistant'), + ), + ); + }, + }); + const config = buildSubagentThreadTaskConfig( + store, + { userId, parentConversationId }, + { completionWakeups: true }, + ); + const started = config.store.start( + taskRequest(config.scopeId, { + parentRunId: 'parent-response-1', + parentAgentId: 'agent_parent_1', + }), + ); + await waitForSettled(store, config.scopeId, started); + saveMessage.mockRestore(); + + expect(terminalSavedAtSettle).toEqual([true]); + }); + + it.each([false, true])( + 'does not announce a child without an admitted delivery (wakeups enabled: %s)', + async (completionWakeups) => { + const userId = 'poll-only-settled-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const onTaskSettled = jest.fn(); + const store = new SubagentThreadTaskStore(methods, { + onTaskPrepared: jest.fn(() => false), + onTaskSettled, + }); + const config = buildSubagentThreadTaskConfig( + store, + { userId, parentConversationId }, + { completionWakeups }, + ); + const started = config.store.start( + taskRequest(config.scopeId, { parentRunId: 'parent-response-1' }), + ); + await waitForSettled(store, config.scopeId, started); + + expect(onTaskSettled).not.toHaveBeenCalled(); + }, + ); + it('keeps subagent completion delivery poll-only when wakeups are disabled', async () => { const userId = 'poll-only-user'; const parentConversationId = randomUUID(); @@ -951,7 +1018,11 @@ describe('SubagentThreadTaskStore', () => { async (_registration: SubagentTaskWakeupRegistration) => undefined, ); const firstWorker = new SubagentThreadTaskStore(methods, { onTaskPrepared: firstWakeup }); - const secondWorker = new SubagentThreadTaskStore(methods, { onTaskPrepared: replayWakeup }); + const onReplaySettled = jest.fn(); + const secondWorker = new SubagentThreadTaskStore(methods, { + onTaskPrepared: replayWakeup, + onTaskSettled: onReplaySettled, + }); const config = buildSubagentThreadTaskConfig( firstWorker, { userId, parentConversationId }, @@ -999,6 +1070,10 @@ describe('SubagentThreadTaskStore', () => { expect(firstRun).toHaveBeenCalledTimes(1); expect(replayRun).not.toHaveBeenCalled(); + expect(onReplaySettled).toHaveBeenCalledTimes(1); + expect(onReplaySettled).toHaveBeenCalledWith(userId, parentConversationId, [ + requireAccepted(first).task.taskId, + ]); expect(firstWakeup).toHaveBeenCalledTimes(1); const firstRegistration = firstWakeup.mock.calls[0]?.[0]; const replayRegistration = replayWakeup.mock.calls[0]?.[0]; @@ -1056,8 +1131,16 @@ describe('SubagentThreadTaskStore', () => { const parentConversationId = randomUUID(); const threadId = randomUUID(); await saveParent(userId, parentConversationId); - const store = new SubagentThreadTaskStore(methods); - const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const onTaskSettled = jest.fn(); + const store = new SubagentThreadTaskStore(methods, { + onTaskPrepared: jest.fn(), + onTaskSettled, + }); + const config = buildSubagentThreadTaskConfig( + store, + { userId, parentConversationId }, + { completionWakeups: true }, + ); await methods.saveConvo( { userId }, { @@ -1095,7 +1178,7 @@ describe('SubagentThreadTaskStore', () => { }, ); const run = jest.fn(taskRequest(config.scopeId).run); - const retry = store.start( + const retry = config.store.start( taskRequest(config.scopeId, { threadId, idempotencyKey: 'abandoned-attempt', @@ -1106,6 +1189,11 @@ describe('SubagentThreadTaskStore', () => { await waitForSettled(store, config.scopeId, retry); expect(run).not.toHaveBeenCalled(); + expect(onTaskSettled).toHaveBeenCalledTimes(1); + expect(onTaskSettled).toHaveBeenCalledWith(userId, parentConversationId, [ + requireAccepted(retry).task.taskId, + 'abandoned', + ]); expect(store.claim(config.scopeId, requireAccepted(retry).task.taskId)).toMatchObject({ status: 'error', error: @@ -1116,6 +1204,32 @@ describe('SubagentThreadTaskStore', () => { '+subagentTask', ); expect(messages.map((message) => message.subagentTask?.status)).toEqual(['running', 'error']); + + const replaySettled = jest.fn(); + const replayWorker = new SubagentThreadTaskStore(methods, { + onTaskPrepared: jest.fn(), + onTaskSettled: replaySettled, + }); + const replayConfig = buildSubagentThreadTaskConfig( + replayWorker, + { userId, parentConversationId }, + { completionWakeups: true }, + ); + const replay = replayConfig.store.start( + taskRequest(replayConfig.scopeId, { + threadId, + idempotencyKey: 'abandoned-attempt', + requestFingerprint: 'same-inputs', + run, + }), + ); + await waitForSettled(replayWorker, replayConfig.scopeId, replay); + expect(replaySettled).toHaveBeenCalledTimes(1); + expect(replaySettled).toHaveBeenCalledWith(userId, parentConversationId, [ + requireAccepted(retry).task.taskId, + 'abandoned', + ]); + expect(run).not.toHaveBeenCalled(); }); it('holds one active lease per child and exposes provisional ownership safely', async () => { diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index 355ff40a0dc..f3bae71a20c 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -140,6 +140,8 @@ interface PreparedThread { content: string; taskId: string; parentRunId: string; + /** Original running task closed by a retry under a new terminal message id. */ + previousTaskId?: string; }; } @@ -262,7 +264,14 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor taskId: string; tenantId?: string; }) => Promise; - onTaskPrepared?: (registration: SubagentTaskWakeupRegistration) => Promise | void; + onTaskPrepared?: ( + registration: SubagentTaskWakeupRegistration, + ) => Promise | boolean | void; + /** Called once a child's terminal message is durable, with the parent conversation its + * completion wake-up resumes. The child generation settles before that write, so the + * wake-up is ready only from here. Task ids include the original abandoned + * attempt when a retry persisted its result under a new id. */ + onTaskSettled?: (userId: string, parentConversationId: string, taskIds: string[]) => void; } export interface SubagentTaskWakeupRegistration { @@ -616,6 +625,9 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { private readonly releaseOwnerAdmission?: (userId: string, token: string) => Promise; private readonly cancelUnroutedTask?: SubagentThreadTaskStoreOptions['cancelUnroutedTask']; private readonly onTaskPrepared?: SubagentThreadTaskStoreOptions['onTaskPrepared']; + private readonly onTaskSettled?: SubagentThreadTaskStoreOptions['onTaskSettled']; + /** Tasks whose completion wake-up was registered; only they announce settlement. */ + private readonly wakeupTaskIds = new Set(); private taskControlTransport?: SubagentTaskControlTransport; private activityStream = new SubagentActivityStream(new InMemoryEventTransport()); @@ -658,6 +670,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { this.releaseOwnerAdmission = options.releaseOwnerAdmission; this.cancelUnroutedTask = options.cancelUnroutedTask; this.onTaskPrepared = options.onTaskPrepared; + this.onTaskSettled = options.onTaskSettled; } /** Receives payload-free authoritative transitions from the SDK task store. */ @@ -1320,6 +1333,13 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { parentRunId: prepared.replay?.parentRunId ?? request.parentRunId, createdAt: prepared.taskCreatedAt, }); + if (prepared.replay != null) { + this.notifyTaskSettled( + scope, + prepared.replay.taskId, + prepared.replay.previousTaskId, + ); + } if (runtime.signal.aborted) { throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); } @@ -1425,6 +1445,9 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { } throw new Error(publicFailureDetail(error)); } finally { + if (prepared != null) { + this.wakeupTaskIds.delete(prepared.replay?.taskId ?? runtime.taskId); + } if (prepared != null && prepared.replay == null) { this.completeActivity( lease, @@ -2902,6 +2925,10 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { replay: { status: terminal.subagentTask.status as 'completed' | 'error' | 'cancelled', taskId: canonicalTaskId, + ...(terminal.parentMessageId?.endsWith(':user') === true && + terminal.parentMessageId !== `${canonicalTaskId}:user` && { + previousTaskId: terminal.parentMessageId.slice(0, -':user'.length), + }), parentRunId: terminal.subagentTask.parentRunId ?? request.parentRunId, content: terminal.text ?? @@ -2956,6 +2983,9 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { status: 'error', content: abandonedMessage, taskId, + ...(abandoned.messageId.endsWith(':user') && { + previousTaskId: abandoned.messageId.slice(0, -':user'.length), + }), parentRunId: request.parentRunId, }, }; @@ -3192,7 +3222,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ) { return; } - await this.onTaskPrepared({ + const admitted = await this.onTaskPrepared({ userId: scope.userId, parentConversationId: scope.parentConversationId, parentMessageId: task.parentRunId, @@ -3203,6 +3233,9 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { subagentType: request.subagentType, createdAt: task.createdAt, }); + if (admitted !== false) { + this.wakeupTaskIds.add(task.taskId); + } } private async persistCancellation( @@ -3375,6 +3408,26 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { } logger.error(`[subagentThreads] Failed to refresh ${outcome} child thread`, error); } + this.notifyTaskSettled(scope, taskId); + } + + private notifyTaskSettled( + scope: SubagentThreadScope, + taskId: string, + previousTaskId?: string, + ): void { + if (!this.wakeupTaskIds.delete(taskId)) { + return; + } + try { + this.onTaskSettled?.( + scope.userId, + scope.parentConversationId, + previousTaskId == null ? [taskId] : [taskId, previousTaskId], + ); + } catch (error) { + logger.warn('[subagentThreads] Settled-task listener failed', error); + } } private aggregateDetachedUsage(detachedUsage: UsageMetadata[]) { diff --git a/packages/api/src/agents/triggers/README.md b/packages/api/src/agents/triggers/README.md index 93abb3f4bb3..5fe25e09aa1 100644 --- a/packages/api/src/agents/triggers/README.md +++ b/packages/api/src/agents/triggers/README.md @@ -47,6 +47,7 @@ await enqueueAgentTrigger( ## Guarantees - Mongo owns queue state, leases, retry history, and dead letters across restarts and replicas. +- Each replica retains bounded Mongo polling even when idle. See **Idle recovery** below for wake-ups, deadlines, and configuration. - A fresh token fences every claim, including reclaims by the same process. - A delivery is at-least-once. Fire, continue, and steer admission reuse the envelope's stable idempotency identity, so ambiguous retries do not duplicate accepted work. @@ -71,6 +72,48 @@ await enqueueAgentTrigger( `getAgentTriggerDeadLetters` and `requeueAgentTrigger` are intentionally trusted in-process operations. Exposing them through an admin API requires a separate authorization and audit layer. +## Idle recovery + +Queued-turn and maintenance scans start at a 30-second cadence and double the next idle wait +up to two minutes after confirmed-empty discovery. Inspected reservations, reconciliation +candidates, outstanding cleanup markers, full legacy-receipt pages, and failed scans do not +count as empty. A failed discovery waits for its sibling Mongo operation to settle before a +new pass or shutdown can complete; successfully acquired work still gets processed. + +Local failed publication and requeue, unfinished terminal finalization, and purge markers wake +maintenance without waiting for an idle timer. A committed root success or dead letter remains +authoritative when its inline cleanup fails. Healthy terminal writes do not wake a full sweep. +Known reconciliation deadlines cap the next wait, including deadlines that expire while a scan +is in flight. Notifications coalesce behind one active scan; shutdown prevents a follow-up. + +Every replica still scans without notifications, so a crashed producer or a missed cross-replica +wake cannot strand durable work. The cap limits **idle sleep**, not end-to-end recovery time: +scan duration, pagination, active leases, and persistence failures can add time. Delivery claims +retain their separate 15-second idle cap and immediate local enqueue/requeue wake-ups. + +Configure `endpoints.agents.eventDriven.idlePolling` in `librechat.yaml`, for either standard or +experimental clustered startup: + +| Setting | Default (ms) | Allowed range (ms) | +| --- | ---: | ---: | +| `queuedTurnMaxIntervalMs` | 120000 | 30000–300000 | +| `maintenanceMaxIntervalMs` | 120000 | 30000–300000 | +| `deliveryMaxIntervalMs` | 15000 | 1000–300000 | +| `completionWaitMaxIntervalMs` | 60000 | 5000–300000 | + +`completionWaitMaxIntervalMs` caps how long a background or subagent completion waits between +readiness checks while its result or parent turn is not ready. It backs off from 5 seconds by a tenth +of its age; a durable result and a settled generation expedite it, so the cap bounds missed signals. +Parent settlement expedites completion deliveries only in that conversation. A child's durable +terminal result expedites only its task, plus the original task when recovering an abandoned +attempt. Terminal replays signal readiness after registering their delivery, without executing the +child again. A held delivery consumes its wake marker in the same fenced write that releases it +for ordering or defers readiness. These writes use classic operators for DocumentDB compatibility; +older workers can ignore the optional marker and fall back to the configured polling interval. + +Setting a recovery cap to `30000` restores its original fixed recovery frequency. No stored-data +migration is needed; optional activity reporting preserves the existing numeric/boolean results. + ## Remote event ingress Authenticated controllers and source adapters can enqueue the same durable envelope through diff --git a/packages/api/src/agents/triggers/backoff.spec.ts b/packages/api/src/agents/triggers/backoff.spec.ts new file mode 100644 index 00000000000..2fbed8f889d --- /dev/null +++ b/packages/api/src/agents/triggers/backoff.spec.ts @@ -0,0 +1,28 @@ +import { WAITING_RETRY_CAP_MS, WAITING_RETRY_FLOOR_MS, waitingRetryAfter } from './backoff'; + +describe('waitingRetryAfter', () => { + const since = Date.parse('2026-09-24T12:00:00Z'); + const floor = String(WAITING_RETRY_FLOOR_MS / 1_000); + + it('re-checks a fresh wait at the engine floor', () => { + expect(waitingRetryAfter(since, since)).toBe(floor); + expect(waitingRetryAfter(since, since + 45_000)).toBe(floor); + }); + + it('waits a tenth of the elapsed wait once that exceeds the floor', () => { + expect(waitingRetryAfter(since, since + 120_000)).toBe('12'); + expect(waitingRetryAfter(since, since + 5 * 60_000)).toBe('30'); + }); + + it('never waits longer than the cap', () => { + expect(waitingRetryAfter(since, since + 6 * 60 * 60_000)).toBe( + String(WAITING_RETRY_CAP_MS / 1_000), + ); + expect(waitingRetryAfter(since, since + 10 * 60_000, 15_000)).toBe('15'); + }); + + it('treats a missing or future start as a fresh wait', () => { + expect(waitingRetryAfter(Number.NaN, since)).toBe(floor); + expect(waitingRetryAfter(since + 60_000, since)).toBe(floor); + }); +}); diff --git a/packages/api/src/agents/triggers/backoff.ts b/packages/api/src/agents/triggers/backoff.ts new file mode 100644 index 00000000000..d5507b5b62f --- /dev/null +++ b/packages/api/src/agents/triggers/backoff.ts @@ -0,0 +1,26 @@ +/** How long a deferred delivery may wait before re-checking a condition nobody has + * announced a change to. Matches the delivery engine's floor for readiness deferrals. */ +export const WAITING_RETRY_FLOOR_MS = 5_000; +export const WAITING_RETRY_CAP_MS = 60_000; +const WAITING_RETRY_AGE_FRACTION = 0.1; + +/** + * Seconds until a waiting delivery re-checks readiness: a tenth of how long it + * has waited, between the engine's five-second floor and the cap. Producers + * expedite the delivery the moment its condition changes (a result becoming + * durable, a generation settling), so this only bounds re-reads while nothing + * has changed. Without it every pending delivery is re-claimed every five + * seconds for as long as its tool runs or its conversation stays busy. + */ +export function waitingRetryAfter( + waitingSinceMs: number, + nowMs: number = Date.now(), + capMs: number = WAITING_RETRY_CAP_MS, +): string { + const waitedMs = Number.isFinite(waitingSinceMs) ? Math.max(0, nowMs - waitingSinceMs) : 0; + const delayMs = Math.min( + capMs, + Math.max(WAITING_RETRY_FLOOR_MS, waitedMs * WAITING_RETRY_AGE_FRACTION), + ); + return String(Math.ceil(delayMs / 1_000)); +} diff --git a/packages/api/src/agents/triggers/delivery.integration.spec.ts b/packages/api/src/agents/triggers/delivery.integration.spec.ts index bc78de4fa16..364edf0be9f 100644 --- a/packages/api/src/agents/triggers/delivery.integration.spec.ts +++ b/packages/api/src/agents/triggers/delivery.integration.spec.ts @@ -4,6 +4,7 @@ import { createMethods, createModels } from '@librechat/data-schemas'; import type { AgentTriggerDeliveryPersistence, AgentTriggerService } from './service'; import type { AgentTriggerFetch } from './host'; import { __resetShutdownStateForTests } from '../../app/shutdown'; +import { prepareAgentTriggerDelivery } from './delivery'; import { createAgentTriggerEnvelope } from './envelope'; import { createAgentTriggerService } from './service'; @@ -83,6 +84,79 @@ async function eventuallySucceeded(deliveryKey: string) { } describe('durable trigger delivery integration', () => { + it('recovers an abandoned publication across two idle workers without a producer wake', async () => { + const fetcher = jest.fn, Parameters>( + async () => + new Response( + JSON.stringify({ + status: 'started', + streamId: 'cross-replica', + conversationId: 'cross-replica', + generationCreatedAt: 25, + }), + { status: 200, headers: { 'content-type': 'application/json' } }, + ), + ); + const bundles = [createMethods(mongoose), createMethods(mongoose)]; + const scans = bundles.map((methods) => + jest.spyOn(methods, 'recoverAgentTriggerLanePublications'), + ); + const workers = bundles.map((methods) => + createAgentTriggerService({ + methods: methods as typeof methods & AgentTriggerDeliveryPersistence, + fetch: fetcher, + mintToken: () => 'trigger-token', + purgeRecoveryIntervalMs: 20, + deliveryOptions: { + concurrency: 1, + tickMs: 5, + maxIdleTickMs: 40, + }, + }), + ); + const until = async (condition: () => boolean | Promise) => { + const deadline = Date.now() + 5_000; + while (!(await condition())) { + if (Date.now() >= deadline) throw new Error('Timed out waiting for cross-replica recovery'); + await new Promise((resolve) => setTimeout(resolve, 10)); + } + }; + try { + for (const worker of workers) + await worker.initialize({ + address: { address: '127.0.0.1', family: 'IPv4', port: 3080 }, + idlePolling: { maintenanceMaxIntervalMs: 80 }, + }); + await until(() => scans.every((scan) => scan.mock.calls.length >= 3)); + const prepared = prepareAgentTriggerDelivery(envelope(), { orderingKey: 'crashed-producer' }); + // A producer died after its staging write. Neither worker receives enqueue() + // or wake(); both must discover and fence publication via their Mongo fallback. + await mongoose.models.AgentTriggerDelivery.create({ + ...prepared, + laneSequence: 0, + status: 'staging', + attempts: 0, + requeueCount: 0, + claimAvailableAt: prepared.availableAt, + stagingRecoveryAt: new Date(), + }); + await until( + async () => + (await bundles[0].getAgentTriggerDelivery(prepared.deliveryKey))?.status === 'succeeded', + ); + expect(fetcher).toHaveBeenCalledTimes(1); + const delivered = await bundles[1].getAgentTriggerDelivery(prepared.deliveryKey); + expect(delivered).toMatchObject({ + status: 'succeeded', + attempts: 1, + history: [{ outcome: 'succeeded' }], + }); + } finally { + await Promise.all(workers.map((worker) => worker.stop())); + scans.forEach((scan) => scan.mockRestore()); + } + }); + it('moves a trusted envelope through Mongo, the lease worker, and host admission', async () => { const fetcher = jest.fn, Parameters>( async () => diff --git a/packages/api/src/agents/triggers/engine.spec.ts b/packages/api/src/agents/triggers/engine.spec.ts index e00a600654b..697834d4f8c 100644 --- a/packages/api/src/agents/triggers/engine.spec.ts +++ b/packages/api/src/agents/triggers/engine.spec.ts @@ -1,3 +1,4 @@ +import { context, ROOT_CONTEXT } from '@opentelemetry/api'; import type { AgentTriggerDeliveryFailure, AgentTriggerDeliveryRecord, @@ -91,6 +92,166 @@ describe('createAgentTriggerDeliveryEngine', () => { }); }); + it.each([ + ['honours a longer requested wait', '30', 30_000], + ['keeps the default re-check for a shorter request', '1', 5_000], + ['keeps the default re-check without a request', undefined, 5_000], + ])('%s for a readiness deferral', async (_label, retryAfter, expectedDelayMs) => { + const store = storeWith(); + const dispatch = jest.fn(async () => { + throw new AgentTriggerExecutionError('The parent generation has not settled yet.', { + mode: 'continue', + certainty: 'definite', + code: 'PARENT_NOT_READY', + retryable: true, + deferWithoutAttempt: true, + ...(retryAfter != null && { retryAfter }), + }); + }); + const engine = createAgentTriggerDeliveryEngine( + { store, dispatch, now: () => START, workerId: 'worker-1' }, + { concurrency: 1 }, + ); + + await engine.runTick(); + + expect(store.defer).toHaveBeenCalledWith( + expect.objectContaining({ + availableAt: new Date(START.getTime() + expectedDelayMs), + }), + ); + expect(store.retry).not.toHaveBeenCalled(); + }); + + it('re-dispatches at once a delivery its deferral found expedited while held', async () => { + const dispatchesWithin = async (deferred: boolean | 'expedited'): Promise => { + jest.useFakeTimers(); + try { + /** Models the store: the deferred row is claimable again only when the + * deferral honored a wake marker and moved it to now. */ + let claimable = true; + const store = storeWith({ + claimNext: jest.fn(async () => { + if (!claimable) { + return null; + } + claimable = false; + return delivery(); + }), + defer: jest.fn(async () => { + claimable = deferred === 'expedited'; + return deferred; + }), + }); + let dispatches = 0; + const dispatch = jest.fn(async () => { + dispatches++; + if (dispatches > 1) { + return successResult(); + } + throw new AgentTriggerExecutionError('The parent generation has not settled yet.', { + mode: 'continue', + certainty: 'definite', + code: 'PARENT_NOT_READY', + retryable: true, + deferWithoutAttempt: true, + retryAfter: '30', + }); + }); + const engine = createAgentTriggerDeliveryEngine( + { store, dispatch, now: () => START, workerId: 'worker-1' }, + { concurrency: 1, tickMs: 60_000, maxIdleTickMs: 60_000 }, + ); + engine.start(); + await jest.advanceTimersByTimeAsync(100); + await engine.stop(); + return dispatches; + } finally { + jest.useRealTimers(); + } + }; + + await expect(dispatchesWithin(true)).resolves.toBe(1); + await expect(dispatchesWithin('expedited')).resolves.toBe(2); + }); + + it('reclaims an ordering release made due by a held wake signal without waiting for its stale deadline', async () => { + jest.useFakeTimers(); + try { + let claimable = true; + const store = storeWith({ + claimNext: jest.fn(async () => { + if (!claimable) return null; + claimable = false; + return delivery(); + }), + findEarlierUnsettled: jest + .fn() + .mockResolvedValueOnce({ availableAt: new Date(START.getTime() + 60_000) }) + .mockResolvedValue(null), + release: jest.fn(async () => { + claimable = true; + return true; + }), + }); + const dispatch = jest.fn(async () => successResult()); + const engine = createAgentTriggerDeliveryEngine( + { store, dispatch, now: () => START }, + { concurrency: 1, tickMs: 60_000, maxIdleTickMs: 60_000 }, + ); + engine.start(); + await jest.advanceTimersByTimeAsync(100); + await engine.stop(); + expect(store.release).toHaveBeenCalledTimes(1); + expect(dispatch).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } + }); + + it('runs each claim pass under the root trace context', async () => { + const withContext = jest.spyOn(context, 'with'); + const store = storeWith({ claimNext: jest.fn(async () => null) }); + const engine = createAgentTriggerDeliveryEngine( + { store, dispatch: jest.fn(async () => successResult()), now: () => START }, + { concurrency: 1 }, + ); + + try { + await engine.runTick(); + expect(withContext).toHaveBeenCalledWith(ROOT_CONTEXT, expect.any(Function)); + expect(store.claimNext).toHaveBeenCalled(); + } finally { + withContext.mockRestore(); + } + }); + + it('arms each poll timer under the root trace context', async () => { + jest.useFakeTimers(); + const withContext = jest.spyOn(context, 'with'); + try { + const store = storeWith({ claimNext: jest.fn(async () => null) }); + const engine = createAgentTriggerDeliveryEngine( + { store, dispatch: jest.fn(async () => successResult()), now: () => START }, + { concurrency: 1, tickMs: 1_000 }, + ); + engine.start(); + await jest.advanceTimersByTimeAsync(0); + + const timersArmed = withContext.mock.results.filter( + (result, index) => + withContext.mock.calls[index][0] === ROOT_CONTEXT && + result.type === 'return' && + typeof (result.value as { unref?: unknown } | undefined)?.unref === 'function', + ); + expect(timersArmed.length).toBeGreaterThan(0); + await engine.stop(); + } finally { + withContext.mockRestore(); + jest.useRealTimers(); + } + }); + it('persists generation identity when a bound continuation starts', async () => { const envelope = createAgentTriggerEnvelope({ mode: 'continue', diff --git a/packages/api/src/agents/triggers/engine.ts b/packages/api/src/agents/triggers/engine.ts index 8fd6592a3b8..e6d6f9e8fce 100644 --- a/packages/api/src/agents/triggers/engine.ts +++ b/packages/api/src/agents/triggers/engine.ts @@ -1,4 +1,5 @@ import { randomUUID } from 'node:crypto'; +import { context, ROOT_CONTEXT } from '@opentelemetry/api'; import { logger, runAsSystem } from '@librechat/data-schemas'; import type { AgentTriggerExecutionResult } from './host'; import { createAgentTriggerBatchEnvelope } from './batch'; @@ -157,7 +158,7 @@ export interface AgentTriggerDeliveryStore { claimToken: string; attempt: number; availableAt: Date; - }) => Promise; + }) => Promise; complete: (input: { id: string; workerId: string; @@ -283,6 +284,29 @@ function normalizeFailure(failure: AgentTriggerDeliveryFailure): AgentTriggerDel }; } +/** The delay an execution error asked for, as seconds or an HTTP date, clamped. */ +function requestedRetryAfterMs(error: unknown, now: Date): number | undefined { + if (!(error instanceof AgentTriggerExecutionError) || error.retryAfter == null) { + return; + } + const seconds = Number(error.retryAfter); + if (Number.isFinite(seconds) && seconds >= 0) { + return Math.min(seconds * 1_000, MAX_RETRY_AFTER_MS); + } + const absolute = Date.parse(error.retryAfter); + if (Number.isFinite(absolute) && absolute > now.getTime()) { + return Math.min(absolute - now.getTime(), MAX_RETRY_AFTER_MS); + } + return; +} + +/** A readiness deferral may ask to wait longer than the default re-check, never + * shorter: waiting producers back off by age, while every existing `retryAfter` + * below the default keeps the cadence it has always had. */ +function readinessDeferMs(error: unknown, now: Date): number { + return Math.max(DEFAULT_DEFER_MS, requestedRetryAfterMs(error, now) ?? 0); +} + function retryAt( error: unknown, attempt: number, @@ -291,15 +315,9 @@ function retryAt( capMs: number, random: () => number, ): Date { - if (error instanceof AgentTriggerExecutionError && error.retryAfter != null) { - const seconds = Number(error.retryAfter); - if (Number.isFinite(seconds) && seconds >= 0) { - return new Date(now.getTime() + Math.min(seconds * 1_000, MAX_RETRY_AFTER_MS)); - } - const absolute = Date.parse(error.retryAfter); - if (Number.isFinite(absolute) && absolute > now.getTime()) { - return new Date(Math.min(absolute, now.getTime() + MAX_RETRY_AFTER_MS)); - } + const requestedMs = requestedRetryAfterMs(error, now); + if (requestedMs != null) { + return new Date(now.getTime() + requestedMs); } const exponent = Math.min(attempt - 1, 30); const delay = Math.min(baseMs * 2 ** exponent, capMs); @@ -483,8 +501,12 @@ export function createAgentTriggerDeliveryEngine( deletionRejected || runtimeNotReady ) { - const delayMs = - error instanceof AgentTriggerDeliveryDeferredError ? error.delayMs : DEFAULT_DEFER_MS; + let delayMs = DEFAULT_DEFER_MS; + if (error instanceof AgentTriggerDeliveryDeferredError) { + delayMs = error.delayMs; + } else if (runtimeNotReady) { + delayMs = readinessDeferMs(error, attemptedAt); + } const availableAt = new Date(attemptedAt.getTime() + delayMs); noteEligibleAt(availableAt); const deferred = await deps.store.defer({ @@ -693,7 +715,12 @@ export function createAgentTriggerDeliveryEngine( if (activeClaim != null) { return activeClaim; } - activeClaim = runAsSystem(runClaimPass) + /** Claim passes are started from timers and from `wake()` calls made inside + * request handlers. Run them under the root context so a pass never joins + * whichever request happened to wake the engine — otherwise every later tick + * inherits that request's trace for the life of the timer chain. */ + activeClaim = context + .with(ROOT_CONTEXT, () => runAsSystem(runClaimPass)) .then((result) => { /** Only a pass that confirmed an empty queue may advance the idle backoff: work * resets it, and a failed claim proves nothing, so it polls on at the base @@ -773,19 +800,23 @@ export function createAgentTriggerDeliveryEngine( if (eligibleDeadlinesMs.length > 0) { delay = Math.max(0, Math.min(delay, eligibleDeadlinesMs[0] - now().getTime())); } - timer = setTimeout(async () => { - if (stopped) { - return; - } - const nowMs = now().getTime(); - while (eligibleDeadlinesMs.length > 0 && eligibleDeadlinesMs[0] <= nowMs) { - eligibleDeadlinesMs.shift(); - } - await claimAvailable().catch((error) => - logger.error('[agent-triggers] delivery claim pass failed:', error), - ); - schedule(); - }, delay); + /** Created under the root context too: `schedule()` runs from `wake()` inside + * request handlers, and each tick reschedules from its own callback. */ + timer = context.with(ROOT_CONTEXT, () => + setTimeout(async () => { + if (stopped) { + return; + } + const nowMs = now().getTime(); + while (eligibleDeadlinesMs.length > 0 && eligibleDeadlinesMs[0] <= nowMs) { + eligibleDeadlinesMs.shift(); + } + await claimAvailable().catch((error) => + logger.error('[agent-triggers] delivery claim pass failed:', error), + ); + schedule(); + }, delay), + ); timer.unref(); }; diff --git a/packages/api/src/agents/triggers/host.spec.ts b/packages/api/src/agents/triggers/host.spec.ts index b5dcf2ac4ed..f7f731205b1 100644 --- a/packages/api/src/agents/triggers/host.spec.ts +++ b/packages/api/src/agents/triggers/host.spec.ts @@ -588,6 +588,37 @@ describe('createAgentTriggerExecutionHost continue adapter', () => { }); }); + it.each([undefined, 'ask', 'acceptEdits', 'fullAccess'] as const)( + 'forwards prepared coding mode %s without trusting the event payload', + async (mode) => { + const envelope = createContinueEnvelope(); + envelope.event.payload = { codeApprovalMode: 'fullAccess' }; + const fetcher = fetchMock(async () => + response({ + streamId: 'conversation-1', + conversationId: 'conversation-1', + status: 'started', + }), + ); + const host = createAgentTriggerExecutionHost( + deps(fetcher, { + prepareContinue: async () => ({ + status: 'ready', + input: 'durable result', + parentMessageId: 'response-1', + codeApprovalMode: mode, + }), + }), + ); + + await host.dispatch(envelope); + + const body = JSON.parse(String(fetcher.mock.calls[0][1]?.body)); + expect(body.codeApprovalMode).toBe(mode); + if (mode === undefined) expect(body).not.toHaveProperty('codeApprovalMode'); + }, + ); + it('carries a prepared queued-turn payload and settles it after admission', async () => { const envelope = createContinueEnvelope(); const admitted = { diff --git a/packages/api/src/agents/triggers/host.ts b/packages/api/src/agents/triggers/host.ts index d62d0adf575..4dad9baf8bb 100644 --- a/packages/api/src/agents/triggers/host.ts +++ b/packages/api/src/agents/triggers/host.ts @@ -1,6 +1,6 @@ import { logger, tenantStorage } from '@librechat/data-schemas'; import { Constants, EModelEndpoint } from 'librechat-data-provider'; -import type { TFile } from 'librechat-data-provider'; +import type { CodeApprovalMode, TFile } from 'librechat-data-provider'; import type { AgentContinueTriggerEnvelope, AgentFireTriggerEnvelope, @@ -58,6 +58,9 @@ export type AgentTriggerContinuePreparation = files?: Partial[]; quotes?: string[]; manualSkills?: string[]; + /** Parent-selected coding preference for a completion turn. Admission + * revalidates it against live policy; event payloads cannot supply it. */ + codeApprovalMode?: CodeApprovalMode; /** Trusted source identity committed by execution enrollment before the * provider-start fence opens. */ admissionSource?: AgentContinuationAdmissionSource; @@ -668,6 +671,9 @@ async function startRun( ...(readyPreparation?.quotes != null && { quotes: readyPreparation.quotes, }), + ...(readyPreparation?.codeApprovalMode != null && { + codeApprovalMode: readyPreparation.codeApprovalMode, + }), ...(readyPreparation?.manualSkills != null && { manualSkills: readyPreparation.manualSkills, }), diff --git a/packages/api/src/agents/triggers/service.delivery.spec.ts b/packages/api/src/agents/triggers/service.delivery.spec.ts index f3b6b4c1034..0b1278141aa 100644 --- a/packages/api/src/agents/triggers/service.delivery.spec.ts +++ b/packages/api/src/agents/triggers/service.delivery.spec.ts @@ -7,9 +7,11 @@ import { SYSTEM_TENANT_ID, } from '@librechat/data-schemas'; import type { AgentTriggerDeliveryPersistence, AgentTriggerStoredRecord } from './service'; +import type { AgentTriggerDeliveryStore } from './engine'; import { AgentTriggerServiceUnavailableError, createAgentTriggerService } from './service'; import { __resetShutdownStateForTests } from '../../app/shutdown'; import { createAgentTriggerEnvelope } from './envelope'; +import * as deliveryEngineModule from './engine'; jest.mock('@librechat/data-schemas', () => { const actual = jest.requireActual('@librechat/data-schemas'); @@ -207,6 +209,155 @@ describe('durable agent trigger service', () => { await service.stop(); }); + describe('waiting completion deliveries', () => { + const address = { address: '127.0.0.1', family: 'IPv4' as const, port: 3080 }; + const flush = () => new Promise((resolve) => setImmediate(resolve)); + const result = { status: 'completed' as const, output: 'done', settledAt: START }; + + it('expedites and claims the delivery whose background result just became durable', async () => { + const methods = deliveryMethods({ + persistAgentBackgroundToolResult: jest.fn(async () => true), + expediteAgentTriggerDeliveries: jest.fn(async () => ({ expedited: 1, held: 0 })), + }); + const service = createAgentTriggerService({ + methods, + deliveryOptions: { concurrency: 1, tickMs: 60_000 }, + }); + await service.initialize({ address }); + jest.mocked(methods.claimNextAgentTriggerDelivery).mockClear(); + + await expect( + service.persistBackgroundToolResult({ + deliveryKey: 'trigger_background', + sourceId: 'background-tool-completion', + result, + }), + ).resolves.toBe(true); + await flush(); + + expect(methods.expediteAgentTriggerDeliveries).toHaveBeenCalledWith({ + deliveryKeys: ['trigger_background'], + sourceIds: ['background-tool-completion', 'subagent-completion'], + now: expect.any(Date), + }); + expect(methods.claimNextAgentTriggerDelivery).toHaveBeenCalled(); + await service.stop(); + }); + + it('targets only named subagent completions when a child settles', async () => { + const methods = deliveryMethods({ + expediteAgentTriggerDeliveries: jest.fn(async () => ({ expedited: 1, held: 0 })), + }); + const service = createAgentTriggerService({ methods }); + await service.initialize({ address }); + service.expediteCompletionWakeups({ + user: 'user-1', + conversationId: 'parent-1', + taskIds: ['child-1'], + }); + await flush(); + expect(methods.expediteAgentTriggerDeliveries).toHaveBeenCalledWith({ + user: 'user-1', + conversationId: 'parent-1', + taskIds: ['child-1'], + sourceIds: ['subagent-completion'], + now: expect.any(Date), + }); + await service.stop(); + }); + + it('does not expedite a result the store refused to persist', async () => { + const methods = deliveryMethods({ + persistAgentBackgroundToolResult: jest.fn(async () => false), + expediteAgentTriggerDeliveries: jest.fn(async () => ({ expedited: 1, held: 0 })), + }); + const service = createAgentTriggerService({ + methods, + deliveryOptions: { concurrency: 1, tickMs: 60_000 }, + }); + await service.initialize({ address }); + + await expect( + service.persistBackgroundToolResult({ + deliveryKey: 'trigger_background', + sourceId: 'background-tool-completion', + result, + }), + ).resolves.toBe(false); + await flush(); + + expect(methods.expediteAgentTriggerDeliveries).not.toHaveBeenCalled(); + await service.stop(); + }); + + it("expedites a principal's completion deliveries when one of its generations settles", async () => { + const methods = deliveryMethods({ + expediteAgentTriggerDeliveries: jest.fn(async () => ({ expedited: 2, held: 0 })), + }); + let settled: ((event: { userId: string; conversationId: string }) => void) | undefined; + const unsubscribe = jest.fn(); + const service = createAgentTriggerService({ + methods, + deliveryOptions: { concurrency: 1, tickMs: 60_000 }, + subscribeGenerationSettled: (listener) => { + settled = listener; + return unsubscribe; + }, + }); + await service.initialize({ address }); + jest.mocked(methods.claimNextAgentTriggerDelivery).mockClear(); + + settled?.({ userId: '507f1f77bcf86cd799439011', conversationId: 'conversation-1' }); + await flush(); + + expect(methods.expediteAgentTriggerDeliveries).toHaveBeenCalledWith({ + user: '507f1f77bcf86cd799439011', + conversationId: 'conversation-1', + sourceIds: ['background-tool-completion', 'subagent-completion'], + now: expect.any(Date), + }); + expect(methods.claimNextAgentTriggerDelivery).toHaveBeenCalled(); + await service.stop(); + expect(unsubscribe).toHaveBeenCalledTimes(1); + }); + + it('runs a claim pass even when nothing moved, since a match may already be due', async () => { + const methods = deliveryMethods({ + expediteAgentTriggerDeliveries: jest.fn(async () => ({ expedited: 0, held: 0 })), + }); + let settled: ((event: { userId: string; conversationId: string }) => void) | undefined; + const service = createAgentTriggerService({ + methods, + deliveryOptions: { concurrency: 1, tickMs: 60_000 }, + subscribeGenerationSettled: (listener) => { + settled = listener; + return () => undefined; + }, + }); + await service.initialize({ address }); + jest.mocked(methods.claimNextAgentTriggerDelivery).mockClear(); + + settled?.({ userId: '507f1f77bcf86cd799439011', conversationId: 'conversation-1' }); + await flush(); + + expect(methods.expediteAgentTriggerDeliveries).toHaveBeenCalledTimes(1); + expect(methods.claimNextAgentTriggerDelivery).toHaveBeenCalled(); + await service.stop(); + }); + }); + + describe('completion wait configuration', () => { + const address = { address: '127.0.0.1', family: 'IPv4' as const, port: 3080 }; + + it('exposes the configured completion wait cap', async () => { + const service = createAgentTriggerService({ methods: deliveryMethods() }); + expect(service.getCompletionWaitMaxIntervalMs()).toBe(60_000); + await service.initialize({ address, idlePolling: { completionWaitMaxIntervalMs: 20_000 } }); + expect(service.getCompletionWaitMaxIntervalMs()).toBe(20_000); + await service.stop(); + }); + }); + it('advertises ordinary completion but not detached-action capability without durable storage', async () => { const methods = deliveryMethods(); const service = createAgentTriggerService({ @@ -321,6 +472,183 @@ describe('durable agent trigger service', () => { expect(deleteAgentTriggerDeliveriesByUser).toHaveBeenCalledWith('507f1f77bcf86cd799439011'); }); + it.each(['complete', 'dead', 'retire'] as const)( + 'wakes for unfinished %s finalization, but not healthy terminal writes', + async (operation) => { + jest.useFakeTimers(); + let service: ReturnType | undefined; + try { + let store!: AgentTriggerDeliveryStore; + const create = deliveryEngineModule.createAgentTriggerDeliveryEngine; + jest + .spyOn(deliveryEngineModule, 'createAgentTriggerDeliveryEngine') + .mockImplementation((deps, options) => { + store = deps.store; + return create(deps, options); + }); + let unfinished = false; + const terminal = jest.fn(async (_input: unknown, recovery?: { required: boolean }) => { + if (recovery != null) recovery.required = unfinished; + return true; + }); + const methods = deliveryMethods({ + completeAgentTriggerDelivery: terminal, + deadLetterAgentTriggerDelivery: terminal, + retireAgentTriggerDelivery: terminal, + }); + service = createAgentTriggerService({ methods, deliveryOptions: { tickMs: 300_000 } }); + await service.initialize({ address: { address: '127.0.0.1', family: 'IPv4', port: 3080 } }); + await jest.advanceTimersByTimeAsync(180_000); + const count = (methods.recoverAgentTriggerBatchReceipts as jest.Mock).mock.calls.length; + const input = { + id: 'row', + workerId: 'worker', + claimToken: 'token', + attempt: 1, + settledAt: new Date(), + }; + const finish = () => { + if (operation === 'complete') { + return store.complete({ + ...input, + result: { mode: 'fire', status: 'started', conversationId: 'conversation-1' }, + }); + } + if (operation === 'dead') { + return store.dead({ + ...input, + error: { + code: 'FAILED', + message: 'failed', + retryable: false, + certainty: 'definite', + attemptedAt: new Date(), + }, + }); + } + return service!.retire('delivery-key', 'source', 'cancelled'); + }; + await expect(finish()).resolves.toBe(true); + await jest.advanceTimersByTimeAsync(0); + expect(methods.recoverAgentTriggerBatchReceipts).toHaveBeenCalledTimes(count); + unfinished = true; + await expect(finish()).resolves.toBe(true); + await jest.advanceTimersByTimeAsync(0); + expect(methods.recoverAgentTriggerBatchReceipts).toHaveBeenCalledTimes(count + 1); + expect(getTenantId()).not.toBe(SYSTEM_TENANT_ID); + } finally { + await service?.stop(); + jest.restoreAllMocks(); + jest.useRealTimers(); + } + }, + ); + + it.each(['enqueue', 'requeue'] as const)( + 'wakes maintenance after an ambiguous %s write failure', + async (operation) => { + jest.useFakeTimers(); + let service: ReturnType | undefined; + try { + const methods = deliveryMethods(); + (methods.enqueueAgentTriggerDelivery as jest.Mock).mockRejectedValue( + new Error('publication interrupted'), + ); + (methods.requeueAgentTriggerDelivery as jest.Mock).mockRejectedValue( + new Error('publication interrupted'), + ); + service = createAgentTriggerService({ methods, deliveryOptions: { tickMs: 300_000 } }); + await service.initialize({ address: { address: '127.0.0.1', family: 'IPv4', port: 3080 } }); + await jest.advanceTimersByTimeAsync(180_000); + const count = (methods.recoverAgentTriggerLanePublications as jest.Mock).mock.calls.length; + await expect( + operation === 'enqueue' ? service.enqueue(envelope()) : service.requeue('row'), + ).rejects.toThrow('publication interrupted'); + await jest.advanceTimersByTimeAsync(0); + expect(methods.recoverAgentTriggerLanePublications).toHaveBeenCalledTimes(count + 1); + } finally { + await service?.stop(); + jest.useRealTimers(); + } + }, + ); + + it('backs off empty maintenance scans and wakes for new cleanup work', async () => { + jest.useFakeTimers(); + try { + const recoverAgentTriggerUserPurges = jest.fn(async () => 0); + const methods = deliveryMethods({ recoverAgentTriggerUserPurges }); + const service = createAgentTriggerService({ + methods, + deliveryOptions: { concurrency: 1, tickMs: 300_000 }, + }); + await service.initialize({ + address: { address: '127.0.0.1', family: 'IPv4', port: 3080 }, + }); + await jest.advanceTimersByTimeAsync(0); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(60_000); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(2); + await jest.advanceTimersByTimeAsync(120_000); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(3); + await service.prepareUserPurge('507f1f77bcf86cd799439011', new Date()); + await jest.advanceTimersByTimeAsync(0); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(4); + await service.stop(); + await jest.advanceTimersByTimeAsync(120_000); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(4); + } finally { + jest.useRealTimers(); + } + }); + + it('keeps the base cadence when a recovery marker was inspected but no repair completed', async () => { + jest.useFakeTimers(); + try { + const recoverAgentTriggerBatchReceipts = jest.fn( + async (_limit?: number, activity?: { found: boolean }) => { + if (activity != null) activity.found = true; + return 0; + }, + ); + const service = createAgentTriggerService({ + methods: deliveryMethods({ recoverAgentTriggerBatchReceipts }), + deliveryOptions: { concurrency: 1, tickMs: 300_000 }, + }); + await service.initialize({ address: { address: '127.0.0.1', family: 'IPv4', port: 3080 } }); + await jest.advanceTimersByTimeAsync(0); + expect(recoverAgentTriggerBatchReceipts).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(30_000); + expect(recoverAgentTriggerBatchReceipts).toHaveBeenCalledTimes(2); + await service.stop(); + } finally { + jest.useRealTimers(); + } + }); + + it('does not mistake an isolated maintenance failure for an empty scan', async () => { + jest.useFakeTimers(); + try { + const recoverAgentTriggerUserPurges = jest + .fn() + .mockRejectedValueOnce(new Error('mongo unavailable')) + .mockResolvedValue(0); + const methods = deliveryMethods({ recoverAgentTriggerUserPurges }); + const service = createAgentTriggerService({ + methods, + deliveryOptions: { concurrency: 1, tickMs: 300_000 }, + }); + await service.initialize({ address: { address: '127.0.0.1', family: 'IPv4', port: 3080 } }); + await jest.advanceTimersByTimeAsync(0); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(30_000); + expect(recoverAgentTriggerUserPurges).toHaveBeenCalledTimes(2); + await service.stop(); + } finally { + jest.useRealTimers(); + } + }); + it('arms and disarms post-commit purge recovery in system context', async () => { const prepareAgentTriggerUserPurge = jest.fn(async () => { expect(getTenantId()).toBe(SYSTEM_TENANT_ID); @@ -385,7 +713,11 @@ describe('durable agent trigger service', () => { address: { address: '127.0.0.1', family: 'IPv4', port: 3080 }, }); - expect(expireLegacyAgentEventActorReceipts).toHaveBeenCalledWith(expect.any(Date), 17); + expect(expireLegacyAgentEventActorReceipts).toHaveBeenCalledWith( + expect.any(Date), + 17, + expect.objectContaining({ found: false }), + ); await service.stop(); }); @@ -405,7 +737,10 @@ describe('durable agent trigger service', () => { deliveryOptions: { concurrency: 1, tickMs: 60_000 }, }); await service.initialize({ address: { address: '127.0.0.1', family: 'IPv4', port: 3080 } }); - expect(reclaimCheckpointDeletions).toHaveBeenCalledWith(17); + expect(reclaimCheckpointDeletions).toHaveBeenCalledWith( + 17, + expect.objectContaining({ found: false }), + ); let stopped = false; const stop = service.stop().then(() => { stopped = true; @@ -545,6 +880,7 @@ describe('durable agent trigger service', () => { reason: 'result unavailable', onlyIfUnclaimed: true, }), + { required: false }, ); await expect( service.retire('trigger_1', 'background-tool-completion', 'dead recovery', { @@ -553,6 +889,7 @@ describe('durable agent trigger service', () => { ).resolves.toBe(true); expect(retireAgentTriggerDelivery).toHaveBeenLastCalledWith( expect.objectContaining({ onlyIfDead: true }), + { required: false }, ); await service.stop(); }); diff --git a/packages/api/src/agents/triggers/service.ts b/packages/api/src/agents/triggers/service.ts index 18f7b32c690..63aa707b22c 100644 --- a/packages/api/src/agents/triggers/service.ts +++ b/packages/api/src/agents/triggers/service.ts @@ -26,24 +26,43 @@ import type { import type { AgentTriggerEnqueueOptions, PreparedAgentTriggerDelivery } from './delivery'; import type { BoundAddress } from '../../app/origin'; import { AgentTriggerDeliveryDeferredError, createAgentTriggerDeliveryEngine } from './engine'; +import { BACKGROUND_TOOL_COMPLETION_SOURCE } from '../backgroundCompletionWakeup'; import { isShutdownInProgress, registerShutdownTask } from '../../app/shutdown'; +import { SUBAGENT_COMPLETION_SOURCE } from '../subagentCompletionWakeup'; import { generateAgentTriggerToken } from '../../crypto/jwt'; import { prepareAgentTriggerDelivery } from './delivery'; import { selfOriginFromAddress } from '../../app/origin'; import { createAgentTriggerExecutionHost } from './host'; import { parseAgentTriggerEnvelope } from './envelope'; +import { createIdleRecoveryLoop } from '../recovery'; +import { WAITING_RETRY_CAP_MS } from './backoff'; + +/** Internal sources whose deliveries wait on a result or on their parent generation. */ +const COMPLETION_WAKEUP_SOURCES = [BACKGROUND_TOOL_COMPLETION_SOURCE, SUBAGENT_COMPLETION_SOURCE]; export const AGENT_TRIGGER_TOKEN_TTL = '60s'; const DEFAULT_USER_DRAIN_TIMEOUT_MS = 35_000; const DEFAULT_USER_DRAIN_POLL_MS = 100; const DEFAULT_PURGE_RECOVERY_INTERVAL_MS = 30_000; +const DEFAULT_PURGE_RECOVERY_MAX_IDLE_INTERVAL_MS = 2 * 60_000; const DEFAULT_PURGE_RECOVERY_LIMIT = 25; export interface AgentTriggerServiceOptions { completionResultBatchSize?: number; address?: BoundAddress | string | null; + idlePolling?: { + queuedTurnMaxIntervalMs?: number; + maintenanceMaxIntervalMs?: number; + deliveryMaxIntervalMs?: number; + completionWaitMaxIntervalMs?: number; + }; } +/** A principal's deliveries resuming one conversation, or exact deliveries. */ +export type AgentTriggerCompletionExpedite = + | { user: string; conversationId: string; taskIds?: string[] } + | { deliveryKeys: string[] }; + export interface AgentTriggerServiceDeps { fetch?: AgentTriggerExecutionHostDeps['fetch']; getTimezone?: AgentTriggerExecutionHostDeps['getTimezone']; @@ -57,9 +76,19 @@ export interface AgentTriggerServiceDeps { userDrainPollMs?: number; purgeRecoveryIntervalMs?: number; purgeRecoveryLimit?: number; - reclaimCheckpointDeletions?: (limit: number) => Promise; + reclaimCheckpointDeletions?: (limit: number, activity?: { found: boolean }) => Promise; supportsDetachedActionCompletion?: () => boolean; settleSourceBeforeDeadLetter?: AgentTriggerDeliveryEngineDeps['settleSourceBeforeDeadLetter']; + /** Subscribes to generations reaching a terminal state; returns an unsubscribe. */ + subscribeGenerationSettled?: ( + listener: (event: AgentTriggerGenerationSettledEvent) => void, + ) => () => void; +} + +/** The part of a settled generation that decides which waiting deliveries it may unblock. */ +export interface AgentTriggerGenerationSettledEvent { + userId: string; + conversationId: string; } export interface AgentTriggerDeliveryReceipt { @@ -111,14 +140,15 @@ export interface AgentTriggerDeliveryPersistence { releaseAgentTriggerDelivery: AgentTriggerDeliveryStore['release']; beginAgentTriggerDeliveryAttempt: AgentTriggerDeliveryStore['beginAttempt']; deferAgentTriggerDeliveryAttempt: AgentTriggerDeliveryStore['defer']; - completeAgentTriggerDelivery: AgentTriggerDeliveryStore['complete']; + completeAgentTriggerDelivery: AgentTriggerDeliveryMethods['completeAgentTriggerDelivery']; retireAgentTriggerDelivery: AgentTriggerDeliveryMethods['retireAgentTriggerDelivery']; renewAgentTriggerDeliveryProducerLease: AgentTriggerDeliveryMethods['renewAgentTriggerDeliveryProducerLease']; persistAgentBackgroundToolResult?: AgentTriggerDeliveryMethods['persistAgentBackgroundToolResult']; + expediteAgentTriggerDeliveries?: AgentTriggerDeliveryMethods['expediteAgentTriggerDeliveries']; getAgentBackgroundToolResultClaim?: AgentTriggerDeliveryMethods['getAgentBackgroundToolResultClaim']; releaseAgentBackgroundToolResultClaims?: AgentTriggerDeliveryMethods['releaseAgentBackgroundToolResultClaims']; retryAgentTriggerDelivery: AgentTriggerDeliveryStore['retry']; - deadLetterAgentTriggerDelivery: AgentTriggerDeliveryStore['dead']; + deadLetterAgentTriggerDelivery: AgentTriggerDeliveryMethods['deadLetterAgentTriggerDelivery']; getAgentTriggerDelivery: (deliveryKey: string) => Promise; getAgentTriggerDeliveryStatus: ( deliveryKey: string, @@ -132,17 +162,30 @@ export interface AgentTriggerDeliveryPersistence { availableAt: Date, ) => Promise; countActiveAgentTriggerDeliveriesByUser: (userId: string, now: Date) => Promise; - recoverAgentTriggerLanePublications: (limit?: number) => Promise; - recoverAgentTriggerBatchReceipts: (limit?: number) => Promise; - reclaimInactiveAgentTriggerLanes: (limit?: number) => Promise; + recoverAgentTriggerLanePublications: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; + recoverAgentTriggerBatchReceipts: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; + reclaimInactiveAgentTriggerLanes: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; prepareAgentTriggerUserPurge: ( userId: string, fenceStartedAt: Date, tenantId?: string, ) => Promise; cancelAgentTriggerUserPurge: (userId: string, fenceStartedAt: Date) => Promise; - recoverAgentTriggerUserPurges: (limit?: number) => Promise; - expireLegacyAgentEventActorReceipts?: (now: Date, limit?: number) => Promise; + recoverAgentTriggerUserPurges: (limit?: number, activity?: { found: boolean }) => Promise; + expireLegacyAgentEventActorReceipts?: ( + now: Date, + limit?: number, + activity?: { found: boolean }, + ) => Promise; deleteAgentTriggerDeliveriesByUser: (userId: string) => Promise; } @@ -170,7 +213,7 @@ export interface AgentTriggerService { deliveryKey: string, sourceId: string, reason: string, - options?: { onlyIfUnclaimed?: boolean; onlyIfDead?: boolean }, + options?: { onlyIfUnclaimed?: boolean; onlyIfDead?: boolean; requireTransition?: boolean }, ) => Promise; renewProducerLease: (deliveryKey: string, sourceId: string, leaseUntil: Date) => Promise; persistBackgroundToolResult: (input: { @@ -186,6 +229,10 @@ export interface AgentTriggerService { input: Parameters[0], ) => ReturnType; getBackgroundCompletionResultBatchSize: () => number; + /** Longest a waiting completion delivery re-checks readiness. */ + getCompletionWaitMaxIntervalMs: () => number; + /** Best effort: moves waiting completion deliveries forward after what they wait on changed. */ + expediteCompletionWakeups: (input: AgentTriggerCompletionExpedite) => void; releaseBackgroundToolResultClaims: AgentTriggerDeliveryMethods['releaseAgentBackgroundToolResultClaims']; drainUser: (userId: string) => Promise; prepareUserPurge: (userId: string, fenceStartedAt: Date, tenantId?: string) => Promise; @@ -193,9 +240,29 @@ export interface AgentTriggerService { purgeUser: (userId: string) => Promise; } +/** A producer can commit durable state before an inline finalizer fails. Keep + * its authoritative result intact, but wake maintenance for the remaining marker. + * Rejections are uncertain commits, so they also request recovery. Healthy writes + * do not turn every delivered event into a full maintenance sweep. */ +async function withMaintenanceRecovery( + operation: (recovery: { required: boolean }) => Promise, + wake: () => void, +): Promise { + const recovery = { required: false }; + try { + return await operation(recovery); + } catch (error) { + recovery.required = true; + throw error; + } finally { + if (recovery.required) wake(); + } +} + function createDeliveryStore( methods: AgentTriggerDeliveryPersistence, supportsDetachedActionCompletion: () => boolean, + wakeMaintenance: () => void, ): AgentTriggerDeliveryStore { return { claimNext: (input) => @@ -215,9 +282,17 @@ function createDeliveryStore( release: methods.releaseAgentTriggerDelivery, beginAttempt: methods.beginAgentTriggerDeliveryAttempt, defer: methods.deferAgentTriggerDeliveryAttempt, - complete: methods.completeAgentTriggerDelivery, + complete: (input) => + withMaintenanceRecovery( + (recovery) => methods.completeAgentTriggerDelivery(input, recovery), + wakeMaintenance, + ), retry: methods.retryAgentTriggerDelivery, - dead: methods.deadLetterAgentTriggerDelivery, + dead: (input) => + withMaintenanceRecovery( + (recovery) => methods.deadLetterAgentTriggerDelivery(input, recovery), + wakeMaintenance, + ), }; } @@ -275,10 +350,11 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A } let boundOrigin: string | undefined; let backgroundCompletionResultBatchSize = 8; + let completionWaitMaxIntervalMs = WAITING_RETRY_CAP_MS; let deliveryEngine: AgentTriggerDeliveryEngine | undefined; let initializePromise: Promise | undefined; - let purgeRecoveryPromise: Promise | undefined; - let purgeRecoveryTimer: NodeJS.Timeout | undefined; + let purgeRecoveryPromise: Promise | undefined; + let purgeRecoveryLoop: ReturnType | undefined; let deliveryReady = false; let stopping = false; const isPrincipalActive = deps.isPrincipalActive; @@ -362,9 +438,9 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A } }; - const recoverPurges = (): Promise => { + const recoverPurges = (): Promise => { if (deps.methods == null || stopping) { - return Promise.resolve(); + return Promise.resolve(false); } if (purgeRecoveryPromise != null) { return purgeRecoveryPromise; @@ -378,14 +454,20 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A * half-recovered batch clears a request that a later successful recovery * can no longer re-arm, retaining the lane permanently — so reclamation * still waits for a batch-recovery pass that did not fail. */ - const isolated = (label: string, run: () => Promise): Promise => - run().catch((error) => { + let failed = false; + const activity = { found: false }; + const isolated = async (label: string, run: () => Promise): Promise => { + try { + return await run(); + } catch (error) { + failed = true; logger.error( `[agent-triggers] durable delivery maintenance step failed (${label}):`, error, ); return 0; - }); + } + }; const current = runAsSystem(async () => { const [ purgedUsers, @@ -394,13 +476,16 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A expiredLegacyActorReceipts, retiredCheckpointDeletions, ] = await Promise.all([ - isolated('user purges', () => methods.recoverAgentTriggerUserPurges(purgeRecoveryLimit)), + isolated('user purges', () => + methods.recoverAgentTriggerUserPurges(purgeRecoveryLimit, activity), + ), isolated('lane publications', () => - methods.recoverAgentTriggerLanePublications(purgeRecoveryLimit), + methods.recoverAgentTriggerLanePublications(purgeRecoveryLimit, activity), ), - methods.recoverAgentTriggerBatchReceipts(purgeRecoveryLimit).then( + methods.recoverAgentTriggerBatchReceipts(purgeRecoveryLimit, activity).then( (count) => ({ succeeded: true as const, count }), (error) => { + failed = true; logger.error( '[agent-triggers] durable delivery maintenance step failed (batch receipts):', error, @@ -411,18 +496,22 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A isolated( 'legacy actor receipts', () => - methods.expireLegacyAgentEventActorReceipts?.(new Date(), purgeRecoveryLimit) ?? - Promise.resolve(0), + methods.expireLegacyAgentEventActorReceipts?.( + new Date(), + purgeRecoveryLimit, + activity, + ) ?? Promise.resolve(0), ), isolated( 'checkpoint deletion evidence', - () => deps.reclaimCheckpointDeletions?.(purgeRecoveryLimit) ?? Promise.resolve(0), + () => + deps.reclaimCheckpointDeletions?.(purgeRecoveryLimit, activity) ?? Promise.resolve(0), ), ]); const recoveredBatches = batchRecovery.count; const reclaimedLanes = batchRecovery.succeeded ? await isolated('lane reclamation', () => - methods.reclaimInactiveAgentTriggerLanes(purgeRecoveryLimit), + methods.reclaimInactiveAgentTriggerLanes(purgeRecoveryLimit, activity), ) : 0; if (publishedLanes > 0) { @@ -445,9 +534,20 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A retiredCheckpointDeletions, }); } + return ( + !failed && + !activity.found && + purgedUsers === 0 && + publishedLanes === 0 && + recoveredBatches === 0 && + reclaimedLanes === 0 && + expiredLegacyActorReceipts === 0 && + retiredCheckpointDeletions === 0 + ); }) .catch((error) => { logger.error('[agent-triggers] durable delivery maintenance failed:', error); + return false; }) .finally(() => { if (purgeRecoveryPromise === current) { @@ -459,21 +559,49 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A }; const startPurgeRecovery = (): void => { - if (purgeRecoveryTimer != null) { + void purgeRecoveryLoop?.start(); + }; + + /** Moves waiting completion deliveries forward when what they wait on has + * changed, and marks held ones so their next deferral re-checks at once. The + * claim pass always runs: a matching delivery may already be due here without + * having moved. Best effort: a missed expedite only means the delivery + * re-checks at its backoff instead of immediately. */ + const expediteCompletions = (input: AgentTriggerCompletionExpedite): void => { + const expedite = deps.methods?.expediteAgentTriggerDeliveries; + if (expedite == null || !deliveryReady || stopping) { return; } - void recoverPurges(); - purgeRecoveryTimer = setInterval(() => void recoverPurges(), purgeRecoveryIntervalMs); - purgeRecoveryTimer.unref(); + void runAsSystem(() => + expedite({ + ...('user' in input + ? { + user: input.user, + conversationId: input.conversationId, + ...(input.taskIds != null && { taskIds: input.taskIds }), + } + : { deliveryKeys: input.deliveryKeys }), + sourceIds: + 'user' in input && input.taskIds != null + ? [SUBAGENT_COMPLETION_SOURCE] + : COMPLETION_WAKEUP_SOURCES, + now: new Date(), + }), + ) + .then(() => deliveryEngine?.wake()) + .catch((error) => + logger.warn('[agent-triggers] failed to expedite waiting completion deliveries:', error), + ); }; + let unsubscribeGenerationSettled: (() => void) | undefined; + const stop = async (): Promise => { stopping = true; deliveryReady = false; - if (purgeRecoveryTimer != null) { - clearInterval(purgeRecoveryTimer); - purgeRecoveryTimer = undefined; - } + unsubscribeGenerationSettled?.(); + unsubscribeGenerationSettled = undefined; + await purgeRecoveryLoop?.stop(); await initializePromise?.catch(() => undefined); await deliveryEngine?.stop(); await purgeRecoveryPromise?.catch(() => undefined); @@ -489,6 +617,8 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A return { initialize: (options = {}) => { backgroundCompletionResultBatchSize = options.completionResultBatchSize ?? 8; + completionWaitMaxIntervalMs = + options.idlePolling?.completionWaitMaxIntervalMs ?? WAITING_RETRY_CAP_MS; boundOrigin = selfOriginFromAddress(options.address) ?? boundOrigin; if (deps.methods == null || deliveryReady) { return Promise.resolve(); @@ -506,6 +636,15 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A const methods = deps.methods; initializePromise = runAsSystem(async () => { requireDeliveryOrigin(boundOrigin); + purgeRecoveryLoop = createIdleRecoveryLoop({ + intervalMs: purgeRecoveryIntervalMs, + maxIdleIntervalMs: + options.idlePolling?.maintenanceMaxIntervalMs ?? + Math.max(purgeRecoveryIntervalMs, DEFAULT_PURGE_RECOVERY_MAX_IDLE_INTERVAL_MS), + scan: recoverPurges, + onError: (error) => + logger.error('[agent-triggers] durable delivery maintenance failed:', error), + }); await methods.ensureAgentTriggerDeliveryIndexes(); if (stopping || isShutdownInProgress()) { throw new AgentTriggerServiceUnavailableError( @@ -514,16 +653,28 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A } deliveryEngine = createAgentTriggerDeliveryEngine( { - store: createDeliveryStore(methods, supportsDetachedActionCompletion), + store: createDeliveryStore(methods, supportsDetachedActionCompletion, () => + purgeRecoveryLoop?.wake(), + ), dispatch: dispatchForActivePrincipal, ...(deps.settleSourceBeforeDeadLetter != null && { settleSourceBeforeDeadLetter: deps.settleSourceBeforeDeadLetter, }), }, - deps.deliveryOptions, + { + ...deps.deliveryOptions, + ...(options.idlePolling?.deliveryMaxIntervalMs != null && { + maxIdleTickMs: options.idlePolling.deliveryMaxIntervalMs, + }), + }, ); deliveryReady = true; deliveryEngine.start(); + /** Deliveries waiting on a parent resume that parent's conversation, so a + * settled generation wakes only those, not every waiting task of the user. */ + unsubscribeGenerationSettled ??= deps.subscribeGenerationSettled?.( + ({ userId, conversationId }) => expediteCompletions({ user: userId, conversationId }), + ); startPurgeRecovery(); logger.info('[agent-triggers] durable delivery engine started'); }).finally(() => { @@ -546,7 +697,10 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A }; await requireActivePrincipal(String(prepared.user)); const queued = await runAsSystem(async () => - methods.enqueueAgentTriggerDelivery(durableDelivery), + withMaintenanceRecovery( + () => methods.enqueueAgentTriggerDelivery(durableDelivery), + () => purgeRecoveryLoop?.wake(), + ), ); try { await requireActivePrincipal(String(prepared.user)); @@ -589,7 +743,11 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A runAsSystem(async () => requireMethods().getAgentTriggerDeadLetters(limit)), requeue: (id, availableAt = new Date()) => runAsSystem(async () => { - const revived = await requireMethods().requeueAgentTriggerDelivery(id, availableAt); + const methods = requireMethods(); + const revived = await withMaintenanceRecovery( + () => methods.requeueAgentTriggerDelivery(id, availableAt), + () => purgeRecoveryLoop?.wake(), + ); if (revived != null) { if (availableAt.getTime() > Date.now()) { deliveryEngine?.noteEligibleAt(availableAt); @@ -601,14 +759,23 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A }), retire: (deliveryKey, sourceId, reason, options) => runAsSystem(async () => { - const retired = await requireCleanupMethods().retireAgentTriggerDelivery({ - deliveryKey, - sourceId, - reason, - settledAt: new Date(), - ...(options?.onlyIfUnclaimed === true ? { onlyIfUnclaimed: true } : {}), - ...(options?.onlyIfDead === true ? { onlyIfDead: true } : {}), - }); + const methods = requireCleanupMethods(); + const retired = await withMaintenanceRecovery( + (recovery) => + methods.retireAgentTriggerDelivery( + { + deliveryKey, + sourceId, + reason, + settledAt: new Date(), + ...(options?.onlyIfUnclaimed === true ? { onlyIfUnclaimed: true } : {}), + ...(options?.onlyIfDead === true ? { onlyIfDead: true } : {}), + ...(options?.requireTransition === true ? { requireTransition: true } : {}), + }, + recovery, + ), + () => purgeRecoveryLoop?.wake(), + ); if (retired) { deliveryEngine?.wake(); } @@ -625,7 +792,11 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A persistBackgroundToolResult: (input) => runAsSystem(async () => { const persist = requireMethods().persistAgentBackgroundToolResult; - return persist == null ? false : persist(input); + const persisted = persist == null ? false : await persist(input); + if (persisted) { + expediteCompletions({ deliveryKeys: [input.deliveryKey] }); + } + return persisted; }), getBackgroundToolResultClaim: (input) => runAsSystem(async () => { @@ -633,6 +804,8 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A return getClaim == null ? null : getClaim(input); }), getBackgroundCompletionResultBatchSize: () => backgroundCompletionResultBatchSize, + getCompletionWaitMaxIntervalMs: () => completionWaitMaxIntervalMs, + expediteCompletionWakeups: (input) => expediteCompletions(input), releaseBackgroundToolResultClaims: (input) => runAsSystem(async () => { const release = requireMethods().releaseAgentBackgroundToolResultClaims; @@ -640,9 +813,14 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A }), drainUser, prepareUserPurge: (userId, fenceStartedAt, tenantId) => - runAsSystem(async () => - requireCleanupMethods().prepareAgentTriggerUserPurge(userId, fenceStartedAt, tenantId), - ), + runAsSystem(async () => { + const methods = requireCleanupMethods(); + await withMaintenanceRecovery( + () => methods.prepareAgentTriggerUserPurge(userId, fenceStartedAt, tenantId), + () => purgeRecoveryLoop?.wake(), + ); + purgeRecoveryLoop?.wake(); + }), cancelUserPurge: (userId, fenceStartedAt) => runAsSystem(async () => requireCleanupMethods().cancelAgentTriggerUserPurge(userId, fenceStartedAt), @@ -651,6 +829,13 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A // shutdown has begun. Persistence remains usable even though admissions // and the delivery engine are deliberately no longer ready. purgeUser: (userId) => - runAsSystem(async () => requireCleanupMethods().deleteAgentTriggerDeliveriesByUser(userId)), + runAsSystem(async () => { + const methods = requireCleanupMethods(); + await withMaintenanceRecovery( + () => methods.deleteAgentTriggerDeliveriesByUser(userId), + () => purgeRecoveryLoop?.wake(), + ); + purgeRecoveryLoop?.wake(); + }), }; } diff --git a/packages/api/src/code/bridge.spec.ts b/packages/api/src/code/bridge.spec.ts index 778129d1e4c..faf724de7f3 100644 --- a/packages/api/src/code/bridge.spec.ts +++ b/packages/api/src/code/bridge.spec.ts @@ -319,6 +319,97 @@ describe('getCodeBridgeWorkerStatus', () => { expect(fetchImpl).toHaveBeenCalledTimes(1); }); + test('bypasses cached status without falling back on failure or changing normal polling', async () => { + jest.spyOn(Date, 'now').mockReturnValue(1_000); + const offline = () => + new Response( + JSON.stringify({ + protocolVersion: 1, + workerId: 'personal-vm', + online: false, + ready: false, + }), + ); + const fetchImpl = jest + .fn() + .mockImplementationOnce(async () => offline()) + .mockRejectedValueOnce(new Error('offline')) + .mockImplementation(async () => offline()); + const poll = createCodeBridgeStatusPoller({ fetchImpl, maxConcurrent: 1 }); + const params = { + baseURL: 'https://code.example.com/v1', + token: 'administrator-token', + workerId: 'personal-vm', + }; + await expect(poll(params)).resolves.toEqual({ status: 'offline' }); + await expect(poll({ ...params, bypassCache: true })).rejects.toEqual( + expect.objectContaining({ reason: 'failed' }), + ); + await expect(poll({ ...params, bypassCache: true })).resolves.toEqual({ status: 'offline' }); + expect(fetchImpl).toHaveBeenCalledTimes(3); + await expect(poll(params)).resolves.toEqual({ status: 'offline' }); + expect(fetchImpl).toHaveBeenCalledTimes(3); + }); + + test('never joins an older in-flight poll and shares its upstream concurrency limit', async () => { + const responses: Array<(response: Response) => void> = []; + const fetchImpl = jest.fn(() => new Promise((resolve) => responses.push(resolve))); + const poll = createCodeBridgeStatusPoller({ fetchImpl, maxConcurrent: 2 }); + const params = { + baseURL: 'https://code.example.com/v1', + token: 'administrator-token', + workerId: 'personal-vm', + }; + const previous = poll(params); + const current = poll({ ...params, bypassCache: true }); + expect(current).not.toBe(previous); + expect(poll(params)).toBe(previous); + expect(fetchImpl).toHaveBeenCalledTimes(2); + await expect(poll({ ...params, bypassCache: true })).rejects.toEqual( + expect.objectContaining({ reason: 'busy' }), + ); + await expect(poll({ ...params, workerId: 'second-vm' })).rejects.toEqual( + expect.objectContaining({ reason: 'busy' }), + ); + responses[1]( + new Response( + JSON.stringify({ + protocolVersion: 1, + workerId: 'personal-vm', + online: true, + ready: true, + leaseExpiresInMs: 50_000, + capabilities: { sandboxProfile: 'native-srt', runtimes: ['bash'] }, + }), + ), + ); + await expect(current).resolves.toEqual(expect.objectContaining({ status: 'ready' })); + responses[0]( + new Response( + JSON.stringify({ + protocolVersion: 1, + workerId: 'personal-vm', + online: false, + ready: false, + }), + ), + ); + await expect(previous).resolves.toEqual({ status: 'offline' }); + const next = poll({ ...params, bypassCache: true }); + expect(fetchImpl).toHaveBeenCalledTimes(3); + responses[2]( + new Response( + JSON.stringify({ + protocolVersion: 1, + workerId: 'personal-vm', + online: false, + ready: false, + }), + ), + ); + await next; + }); + test('does not coalesce status requests across credential rotations', async () => { const fetchImpl = jest.fn().mockImplementation((_input, init) => Promise.resolve( diff --git a/packages/api/src/code/bridge.ts b/packages/api/src/code/bridge.ts index 99b732335e6..c3cb7af0e53 100644 --- a/packages/api/src/code/bridge.ts +++ b/packages/api/src/code/bridge.ts @@ -91,6 +91,8 @@ export function createCodeBridgeStatusPoller({ baseURL: string; token: string; workerId: string; + /** Force a new upstream request, bypassing cached results and older in-flight polls. */ + bypassCache?: boolean; }) => Promise { const requests = new Map< string, @@ -98,6 +100,14 @@ export function createCodeBridgeStatusPoller({ >(); let active = 0; return (params) => { + if (params.bypassCache) { + if (active >= maxConcurrent) return Promise.reject(new CodeBridgeStatusError('busy')); + active += 1; + // Mutation validation needs its own observation, but still shares the polling capacity. + return getCodeBridgeWorkerStatus({ ...params, fetchImpl }).finally(() => { + active -= 1; + }); + } const credentialId = createHash('sha256').update(params.token).digest('base64url'); const normalizedBaseURL = params.baseURL.trim().replace(/\/+$/, ''); const key = `${normalizedBaseURL}\u0000${params.workerId}\u0000${credentialId}`; diff --git a/packages/api/src/code/config.spec.ts b/packages/api/src/code/config.spec.ts index 7c556669e34..4b0749f3eb5 100644 --- a/packages/api/src/code/config.spec.ts +++ b/packages/api/src/code/config.spec.ts @@ -5,6 +5,7 @@ import { mergeAccessibleCodeEnvironments, resolveCodeEnvironmentDecisionVersion, resolveCodeEnvironmentMoveVersion, + resolveCodeEnvironmentMoveCapabilities, } from './config'; describe('resolveCodeEnvironmentDecisionVersion', () => { @@ -32,18 +33,25 @@ describe('resolveCodeEnvironmentMoveVersion', () => { it('advertises moves only where the effective policy enables them', () => { expect(resolveCodeEnvironmentMoveVersion(withMoves({ enabled: true }))).toBe(1); + expect(resolveCodeEnvironmentMoveCapabilities(withMoves({ enabled: true }))).toEqual({ + codeEnvironmentMoveVersion: 1, + codeWorkspaceRecoveryVersion: 1, + }); }); it.each([undefined, {}, { enabled: false }])( 'keeps sealed decisions immovable by default: %j', (conversationMoves) => { expect(resolveCodeEnvironmentMoveVersion(withMoves(conversationMoves))).toBeUndefined(); + expect(resolveCodeEnvironmentMoveCapabilities(withMoves(conversationMoves))).toEqual({}); }, ); it('keeps moves off without any stateful code configuration', () => { expect(resolveCodeEnvironmentMoveVersion({} as AppConfig)).toBeUndefined(); expect(resolveCodeEnvironmentMoveVersion(undefined)).toBeUndefined(); + expect(resolveCodeEnvironmentMoveCapabilities({} as AppConfig)).toEqual({}); + expect(resolveCodeEnvironmentMoveCapabilities(undefined)).toEqual({}); }); }); diff --git a/packages/api/src/code/config.ts b/packages/api/src/code/config.ts index 7904fe756d8..3996fea913d 100644 --- a/packages/api/src/code/config.ts +++ b/packages/api/src/code/config.ts @@ -2,8 +2,10 @@ import { logger } from '@librechat/data-schemas'; import { CODE_ENVIRONMENT_DECISION_VERSION, CODE_ENVIRONMENT_MOVE_VERSION, + CODE_WORKSPACE_RECOVERY_VERSION, EModelEndpoint, } from 'librechat-data-provider'; +import type { TStartupConfig } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; import type { AccessibleCodeEnvironmentConfiguration, @@ -44,6 +46,18 @@ export function resolveCodeEnvironmentMoveVersion( : undefined; } +/** Advertises recovery separately so already-open V1 clients retain ordinary environment moves. */ +export function resolveCodeEnvironmentMoveCapabilities( + appConfig?: Pick | null, +): Pick { + const codeEnvironmentMoveVersion = resolveCodeEnvironmentMoveVersion(appConfig); + if (codeEnvironmentMoveVersion == null) return {}; + return { + codeEnvironmentMoveVersion, + codeWorkspaceRecoveryVersion: CODE_WORKSPACE_RECOVERY_VERSION, + }; +} + /** Enables the implicit managed route only after the versioned rollout is complete. */ export function isImplicitStatefulCodeRouteAvailable( configuredVersion?: string, diff --git a/packages/api/src/code/decision.spec.ts b/packages/api/src/code/decision.spec.ts index 42d29df8f8d..01d1798e937 100644 --- a/packages/api/src/code/decision.spec.ts +++ b/packages/api/src/code/decision.spec.ts @@ -190,21 +190,15 @@ describe('resolveConversationCodeEnvironmentMove', () => { ).toEqual({ codeWorkspaces: [vm] }); }); - it('never switches the workspace of an environment the decision already covers', () => { - expect(() => - resolveConversationCodeEnvironmentMove({ - conversation: sealedOn(mac), - from: [mac], - to: [{ environmentId: 'mac', workspaceId: 'canary' }], - }), - ).toThrow(locked); - expect(() => + it('accepts an explicit workspace replacement for the caller to verify against live status', () => { + const replacement = { environmentId: 'mac', workspaceId: 'canary' }; + expect( resolveConversationCodeEnvironmentMove({ - conversation: sealedOn(mac), - from: [mac], - to: [{ environmentId: 'mac', workspaceId: 'canary' }, vm], + conversation: sealedOn(mac, vm), + from: [mac, vm], + to: [replacement, vm], }), - ).toThrow(locked); + ).toEqual({ codeWorkspaces: [replacement, vm] }); }); it('drops an environment the agents stopped using without adding one', () => { diff --git a/packages/api/src/code/decision.ts b/packages/api/src/code/decision.ts index 8c3880b4969..ff1346c9e07 100644 --- a/packages/api/src/code/decision.ts +++ b/packages/api/src/code/decision.ts @@ -124,11 +124,10 @@ export interface ConversationCodeEnvironmentMove { } /** - * Validates an owner's explicit move of a sealed attached decision onto the environments its - * agents now use. A move may drop environments the agents stopped using and add ones they now use, - * but never changes the workspace of an environment the decision already covers and never upgrades - * a conversation that continues without an attached environment. `from` must repeat the persisted selections, so a client acting - * on a stale view of the conversation cannot replace a decision it has not seen. + * Validates the shape and expected decision of an owner's explicit move. The caller must verify + * live registration of every target and, for same-environment replacements, absence of the old + * workspace before persisting. `from` must repeat the persisted selections, so a stale client + * cannot replace a decision it has not seen. A chat without an attached decision stays sealed. */ export function resolveConversationCodeEnvironmentMove({ conversation, @@ -149,19 +148,7 @@ export function resolveConversationCodeEnvironmentMove({ if (!isCodeWorkspaceSelections(to) || to.length === 0) { throw new CodeWorkspaceSelectionError('invalid'); } - const sealed = new Map( - persisted.codeWorkspaces.map(({ environmentId, workspaceId }) => [environmentId, workspaceId]), - ); - let adds = false; - for (const selection of to) { - const sealedWorkspaceId = sealed.get(selection.environmentId); - if (sealedWorkspaceId == null) { - adds = true; - } else if (sealedWorkspaceId !== selection.workspaceId) { - throw new CodeWorkspaceSelectionError('locked'); - } - } - if (!adds && to.length === sealed.size) { + if (sameSelections(to, persisted.codeWorkspaces)) { throw new CodeWorkspaceSelectionError('locked'); } return { codeWorkspaces: canonicalSelections(to) }; diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 268cbca842d..9bc6872003d 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -1640,8 +1640,20 @@ describe('moving a sealed conversation code-environment decision', () => { ); return res; }; + const status = async () => { + const res = response(); + await handlers.status( + { + user: { id: userId, role: 'USER' }, + params: { environmentId: vm.environmentId }, + } as never, + res as never, + ); + return res; + }; return { move, + status, conversations, getConversation, listConversationRuns, @@ -1669,6 +1681,120 @@ describe('moving a sealed conversation code-environment decision', () => { ); }); + test.each(['attached', undefined] as const)( + 'recovers a missing workspace on the same machine with mode %s and keeps the new decision sealed', + async (codeEnvironmentMode) => { + const missing = { ...vm, workspaceId: 'deleted-project' }; + const { move, conversations, fetchImpl } = setup({ + stored: { + conversationId: 'conversation-1', + codeEnvironmentMode, + codeWorkspaces: [missing], + }, + }); + + const res = await move({ from: [missing], to: [vm] }); + + expect(res.statusCode).toBe(200); + expect(conversations.get('conversation-1')).toEqual({ + conversationId: 'conversation-1', + codeEnvironmentMode: 'attached', + codeWorkspaces: [vm], + }); + expect(fetchImpl).toHaveBeenCalledTimes(1); + expect((await move({ from: [missing], to: [vm] })).statusCode).toBe(409); + expect(fetchImpl).toHaveBeenCalledTimes(1); + }, + ); + + test.each([ + { + name: 'old workspace restored', + workspaces: ['deleted-project', 'project-a'], + ready: true, + reason: 'locked', + }, + { name: 'replacement removed', workspaces: ['unrelated'], ready: true, reason: 'missing' }, + { + name: 'worker unavailable', + workspaces: ['project-a'], + ready: false, + reason: 'worker_unavailable', + }, + { name: 'workspace still missing', workspaces: ['project-a'], ready: true, reason: undefined }, + ])( + 'revalidates recovery after a cached status poll: $name', + async ({ workspaces, ready, reason }) => { + jest.spyOn(Date, 'now').mockReturnValue(1_000); + const old = { ...vm, workspaceId: 'deleted-project' }; + const fetchImpl = jest + .fn() + .mockImplementationOnce(async () => workerStatusResponse()) + .mockImplementation(async () => + workerStatusResponse({ workspaces: workspaces.map((id) => ({ id })), ready }), + ); + const { move, status, conversations, replaceDecision } = setup({ + stored: { conversationId: 'conversation-1', codeWorkspaces: [old] }, + fetchImpl, + }); + expect((await status()).statusCode).toBe(200); + expect((await status()).statusCode).toBe(200); + expect(fetchImpl).toHaveBeenCalledTimes(1); + + const res = await move({ from: [old], to: [vm] }); + + expect(fetchImpl).toHaveBeenCalledTimes(2); + expect(res.statusCode).toBe(reason == null ? 200 : 409); + if (reason != null) { + expect(res.body).toEqual(expect.objectContaining({ reason })); + expect(replaceDecision).not.toHaveBeenCalled(); + } + expect(conversations.get('conversation-1')?.codeWorkspaces).toEqual( + reason == null ? [vm] : [old], + ); + }, + ); + + test('refuses replacement if the old workspace is still registered or was restored', async () => { + const old = { ...vm, workspaceId: 'old-project' }; + const { move, conversations, replaceDecision, fetchImpl } = setup({ + stored: { conversationId: 'conversation-1', codeWorkspaces: [old] }, + fetchImpl: jest.fn(async () => + workerStatusResponse({ workspaces: [{ id: old.workspaceId }, { id: vm.workspaceId }] }), + ), + }); + + const res = await move({ from: [old], to: [vm] }); + + expect(res.statusCode).toBe(409); + expect(res.body).toEqual(expect.objectContaining({ reason: 'locked' })); + expect(fetchImpl).toHaveBeenCalledTimes(1); + expect(replaceDecision).not.toHaveBeenCalled(); + expect(conversations.get('conversation-1')?.codeWorkspaces).toEqual([old]); + }); + + test.each([ + { ready: false, workspaces: [{ id: vm.workspaceId }], reason: 'worker_unavailable' }, + { ready: true, workspaces: [{ id: 'unrelated' }], reason: 'missing' }, + { ready: true, workspaces: [], reason: 'worker_unavailable' }, + ])( + 'leaves a missing decision untouched when recovery fails: $reason', + async ({ reason, ...status }) => { + const old = { ...vm, workspaceId: 'deleted-project' }; + const { move, conversations, replaceDecision } = setup({ + stored: { conversationId: 'conversation-1', codeWorkspaces: [old] }, + fetchImpl: jest.fn(async () => workerStatusResponse(status)), + }); + + const res = await move({ from: [old], to: [vm] }); + + expect(res.statusCode).toBe(409); + expect(res.body).toEqual(expect.objectContaining({ reason })); + expect(replaceDecision).not.toHaveBeenCalled(); + expect(conversations.get('conversation-1')?.codeWorkspaces).toEqual([old]); + }, + ); + test('moves a legacy decision that only stored its selections', async () => { const { move, conversations } = setup({ stored: { conversationId: 'conversation-1', codeWorkspaces: [mac] }, @@ -1805,9 +1931,9 @@ describe('moving a sealed conversation code-environment decision', () => { polls: 0, }, { - name: 'a workspace switch inside a sealed environment', + name: 'a replacement on an environment the caller cannot access', body: { from: [mac], to: [{ environmentId: 'mac', workspaceId: 'canary' }] }, - reason: 'locked', + reason: 'invalid', polls: 0, }, ])('rejects $name without polling any worker', async ({ body, reason, polls }) => { diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index ab6d229905a..ec029031054 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -325,6 +325,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) async function assertWorkspaceRegistered( policy: WorkerPolicy, selection: CodeWorkspaceSelection, + previousWorkspaceId?: string, ): Promise { const target = selectWorkerTarget(policy, selection.environmentId); if (target == null) { @@ -340,6 +341,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) baseURL: target.controlPlane.baseURL, token, workerId: target.workerId, + bypassCache: true, }); } catch (error) { if (error instanceof CodeBridgeStatusError) { @@ -356,6 +358,13 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (!current.workspaces.some(({ id }) => id === selection.workspaceId)) { throw new CodeWorkspaceSelectionError('missing'); } + if ( + previousWorkspaceId != null && + previousWorkspaceId !== selection.workspaceId && + current.workspaces.some(({ id }) => id === previousWorkspaceId) + ) { + throw new CodeWorkspaceSelectionError('locked'); + } } /** @@ -421,7 +430,15 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } try { await Promise.all( - move.codeWorkspaces.map((selection) => assertWorkspaceRegistered(policy, selection)), + move.codeWorkspaces.map((selection) => + assertWorkspaceRegistered( + policy, + selection, + conversation.codeWorkspaces?.find( + ({ environmentId }) => environmentId === selection.environmentId, + )?.workspaceId, + ), + ), ); } catch (error) { if (error instanceof CodeWorkspaceSelectionError) { diff --git a/packages/api/src/endpoints/pricing.spec.ts b/packages/api/src/endpoints/pricing.spec.ts index fdff15f84ef..dac94d00b65 100644 --- a/packages/api/src/endpoints/pricing.spec.ts +++ b/packages/api/src/endpoints/pricing.spec.ts @@ -14,6 +14,21 @@ const { getValueKey, getMultiplier, getCacheMultiplier } = createTxMethods(mongo const deps = { getValueKey, getMultiplier, getCacheMultiplier }; describe('buildTokenConfigMap', () => { + it('exposes Bedrock GPT context windows without changing GPT-OSS', () => { + const models = [ + 'us.openai.gpt-6-sol', + 'global.openai.gpt-6-astra', + 'us.openai.gpt-5.6-terra', + 'us.openai.gpt-oss-120b-1:0', + ]; + const map = buildTokenConfigMap({ modelsConfig: { [EModelEndpoint.bedrock]: models } }, deps); + + for (const model of models.slice(0, 3)) { + expect(map[EModelEndpoint.bedrock][model].context).toBe(950000); + } + expect(map[EModelEndpoint.bedrock][models[3]].context).toBe(128000); + }); + it('resolves context windows without pricing by default', () => { const map = buildTokenConfigMap( { diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index 83a0cc9b63d..c84b2d27328 100644 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -26,6 +26,7 @@ export * from './mcp/oauth'; export * from './mcp/auth'; export * from './mcp/zod'; export * from './mcp/errors'; +export * from './mcp/loading'; export * from './mcp/openid'; export * from './mcp/cache'; export * from './mcp/tools'; diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index b5f3b98616b..913a0000295 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -403,6 +403,10 @@ export class MCPConnectionFactory { return { tools: null, connection: null, oauthRequired: false, oauthUrl: null }; } oauthTokens = loaded.value; + if (!oauthTokens && !this.usesObo && isOAuthServer(this.serverConfig)) { + oauthRequired = true; + shouldAttemptAuthenticatedDiscovery = false; + } } let connection: MCPConnection | null = null; diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index 738a6e179e1..016af28aa4e 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -44,6 +44,7 @@ import { processMCPEnv, isPluginSourced } from '~/utils/env'; import { OAuthLifecycleRelay } from './oauth/pending'; import { preProcessGraphTokens } from '~/utils/graph'; import { isOwnedAbortError } from '~/utils/errors'; +import { withMCPRequestSignal } from './signal'; import { formatToolContent } from './parsers'; import { MCPConnection } from './connection'; import { mcpConfig } from './mcpConfig'; @@ -1526,20 +1527,23 @@ Please follow these instructions when using tools from the respective MCP server } const requestTool = () => - connection!.client.request( - { - method: 'tools/call', - params: { - name: toolName, - arguments: toolArguments, + withMCPRequestSignal(options?.signal, (signal) => + connection!.client.request( + { + method: 'tools/call', + params: { + name: toolName, + arguments: toolArguments, + }, + }, + CallToolResultSchema, + { + timeout: connection!.timeout, + resetTimeoutOnProgress: true, + ...options, + signal, }, - }, - CallToolResultSchema, - { - timeout: connection!.timeout, - resetTimeoutOnProgress: true, - ...options, - }, + ), ); const requestedCredentialSetId = connection.getOAuthCredentialSetId?.(); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts index c723985eca0..64b398b4ab4 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts @@ -200,6 +200,95 @@ describe('MCPConnectionFactory OAuth against real SDK Streamable HTTP server', ( jest.clearAllMocks(); }); + it('tries unauthenticated public tool listing only once when an OAuth server has no token', async () => { + const postHeaders: Array = []; + server = await createOAuthMCPServer({ + onResourceRequest: (request) => { + if (request.method === 'POST') postHeaders.push(request.headers.authorization); + }, + }); + const result = await MCPConnectionFactory.discoverTools( + { + serverName: SERVER_NAME, + serverConfig: { type: 'streamable-http', url: server.url, requiresOAuth: true }, + }, + { + useOAuth: true, + user: { id: USER_ID } as IUser, + flowManager: createFlowManager(), + tokenMethods: { + findToken: tokenStore.findToken, + createToken: tokenStore.createToken, + updateToken: tokenStore.updateToken, + deleteTokens: tokenStore.deleteTokens, + }, + }, + ); + + expect(result).toMatchObject({ tools: null, connection: null, oauthRequired: true }); + expect(postHeaders).toEqual([undefined]); + }); + + it('does not cancel finished SDK requests when a shared run signal is aborted', async () => { + let resourcePosts = 0; + server = await createOAuthMCPServer({ + onResourceRequest: (request) => { + if (request.method === 'POST') resourcePosts += 1; + }, + }); + const tokens = await issueTokens(server); + await storeTokens(tokenStore, server, tokens); + const flowManager = createFlowManager(); + const tokenMethods = { + findToken: tokenStore.findToken, + createToken: tokenStore.createToken, + updateToken: tokenStore.updateToken, + deleteTokens: tokenStore.deleteTokens, + }; + const serverConfig = { + type: 'streamable-http' as const, + url: server.url, + requiresOAuth: true, + }; + connection = await MCPConnectionFactory.create( + { serverName: SERVER_NAME, serverConfig }, + { useOAuth: true, user: { id: USER_ID } as IUser, flowManager, tokenMethods }, + ); + const manager = new MCPManager(); + jest.spyOn(manager, 'getConnection').mockResolvedValue(connection); + const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({ + resolveAllowlists: jest.fn().mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + }), + } as unknown as MCPServersRegistry); + const controller = new AbortController(); + try { + for (let index = 0; index < 3; index += 1) { + await expect( + manager.callTool({ + user: { id: USER_ID } as IUser, + serverName: SERVER_NAME, + serverConfig, + toolName: 'echo', + toolArguments: { message: `call ${index}` }, + provider: 'openai', + flowManager, + tokenMethods, + options: { signal: controller.signal }, + }), + ).resolves.toBeDefined(); + } + const completedPosts = resourcePosts; + controller.abort(new Error('Run ended')); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(resourcePosts).toBe(completedPosts); + } finally { + registrySpy.mockRestore(); + } + }); + it('refreshes an expired callback for the active connection and token waiter after publication settles', async () => { server = await createOAuthMCPServer({ issueRefreshTokens: true, diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 652d71b2a5a..af050c348e7 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -4873,6 +4873,42 @@ describe('MCPConnectionFactory', () => { { name: 'tool2', description: 'Second tool', inputSchema: { type: 'object' } }, ]; + it('preserves public tool listing without a redundant OAuth connect when tokens are absent', async () => { + const serverConfig = { + type: 'streamable-http' as const, + url: 'https://mcp.example.com', + requiresOAuth: true, + }; + mockProcessMCPEnv.mockImplementation(({ options }) => options); + mockFlowManager.createFlowWithHandler.mockResolvedValue(null); + mockConnectionInstance.connect.mockResolvedValue(undefined); + mockConnectionInstance.isConnected.mockResolvedValue(true); + mockConnectionInstance.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + tools: mockTools, + complete: true, + }); + + const result = await MCPConnectionFactory.discoverTools( + { serverName: 'public-oauth', serverConfig }, + { + useOAuth: true, + user: mockUser!, + flowManager: mockFlowManager, + tokenMethods: { + findToken: jest.fn(), + createToken: jest.fn(), + updateToken: jest.fn(), + deleteTokens: jest.fn(), + }, + }, + ); + + expect(result.tools).toEqual(mockTools); + expect(result.oauthRequired).toBe(true); + expect(mockMCPConnection).toHaveBeenCalledTimes(1); + expect(mockConnectionInstance.connect).toHaveBeenCalledTimes(1); + }); + it('should discover tools from a successfully connected server', async () => { const basicOptions = { serverName: 'test-server', @@ -6165,7 +6201,7 @@ describe('MCPConnectionFactory', () => { ); expect(result.tools).toEqual(mockTools); - expect(result.oauthRequired).toBe(false); + expect(result.oauthRequired).toBe(true); expect(oauthOptions.oauthStart).not.toHaveBeenCalled(); expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled(); }); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts index 8a23c262588..bcfdf2f890e 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts @@ -379,15 +379,18 @@ describe('MCPConnection.fetchTools pagination', () => { expect(reserve).not.toHaveBeenCalled(); }); - it('hands the caller signal to the SDK so an in-flight page is cancellable', async () => { + it('detaches a completed SDK page request from its caller signal', async () => { const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('a')] }); const conn = createConnectionWithListTools(listTools); - const signal = AbortSignal.timeout(5000); + const controller = new AbortController(); - await conn.fetchToolsSnapshot(Date.now() + 5000, signal); + await conn.fetchToolsSnapshot(Date.now() + 5000, controller.signal); const options = listTools.mock.calls[0][1]!; - expect(options.signal).toBe(signal); + expect(options.signal === controller.signal).toBe(false); + expect(options.signal?.aborted).toBe(false); + controller.abort(); + expect(options.signal?.aborted).toBe(false); }); it('makes no request and no reservation when the signal is already aborted', async () => { diff --git a/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts b/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts index 3956ea2fa9e..628a0b04419 100644 --- a/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts +++ b/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts @@ -81,6 +81,8 @@ export interface OAuthTestServerOptions { refreshGate?: () => Promise | undefined; /** Optional test hook for controlling echo-tool completion. */ echoHandler?: (message: string) => string | Promise; + /** Observes MCP resource requests, including unauthenticated and cancellation POSTs. */ + onResourceRequest?: (request: http.IncomingMessage) => void; } export interface OAuthTokenRequestRecord { @@ -149,6 +151,7 @@ export async function createOAuthMCPServer( requireResourceParameter = false, rejectRefreshTokens = 0, echoHandler, + onResourceRequest, refreshFailure, refreshGate, } = options; @@ -466,6 +469,7 @@ export async function createOAuthMCPServer( } // All other paths require Bearer token auth + onResourceRequest?.(req); const authHeader = req.headers.authorization; if (!authHeader || !authHeader.startsWith('Bearer ')) { writeBearerChallenge(res, 401, 'invalid_token', 'Missing Authorization header'); diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index 24d1ed08b93..818b6c17a15 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -35,6 +35,7 @@ import { reserveMCPToolsChangedRevision } from './toolsChanged'; import { runOutsideTracing } from '~/utils/tracing'; import { mediaTypeEssence } from '~/utils/headers'; import { isAddressAllowed } from '~/auth/domain'; +import { withMCPRequestSignal } from './signal'; import { withTimeout } from '~/utils/promise'; import { isOAuthServer } from './utils'; import { mcpConfig } from './mcpConfig'; @@ -2740,11 +2741,13 @@ export class MCPConnection extends EventEmitter { signal?: AbortSignal, ): Promise { try { - return await this.client.listTools(cursor != null ? { cursor } : undefined, { - timeout: timeoutMs, - maxTotalTimeout: timeoutMs, - signal, - }); + return await withMCPRequestSignal(signal, (requestSignal) => + this.client.listTools(cursor != null ? { cursor } : undefined, { + timeout: timeoutMs, + maxTotalTimeout: timeoutMs, + signal: requestSignal, + }), + ); } catch (error) { this.emitError(error, 'Failed to fetch tools'); throw error; @@ -2793,7 +2796,9 @@ export class MCPConnection extends EventEmitter { try { // Try ping first as it's the lightest check - await this.client.ping({ signal }); + await withMCPRequestSignal(signal, (requestSignal) => + this.client.ping({ signal: requestSignal }), + ); return this.connectionState === 'connected'; } catch (error) { if (probeAborted()) { @@ -2827,13 +2832,19 @@ export class MCPConnection extends EventEmitter { // If we have capabilities, try calling a supported method to verify connection if (capabilities?.tools) { - await this.client.listTools(undefined, { signal }); + await withMCPRequestSignal(signal, (requestSignal) => + this.client.listTools(undefined, { signal: requestSignal }), + ); return this.connectionState === 'connected'; } else if (capabilities?.resources) { - await this.client.listResources(undefined, { signal }); + await withMCPRequestSignal(signal, (requestSignal) => + this.client.listResources(undefined, { signal: requestSignal }), + ); return this.connectionState === 'connected'; } else if (capabilities?.prompts) { - await this.client.listPrompts(undefined, { signal }); + await withMCPRequestSignal(signal, (requestSignal) => + this.client.listPrompts(undefined, { signal: requestSignal }), + ); return this.connectionState === 'connected'; } else { // No capabilities to test, but we're in connected state and initialization succeeded diff --git a/packages/api/src/mcp/loading.spec.ts b/packages/api/src/mcp/loading.spec.ts new file mode 100644 index 00000000000..59c388b6a94 --- /dev/null +++ b/packages/api/src/mcp/loading.spec.ts @@ -0,0 +1,121 @@ +import { OpenIDReauthRequiredError } from '~/utils/oidc'; +import { loadMCPTools } from './loading'; + +const request = (serverName: string, type: 'all' | 'single' = 'all') => ({ + type, + serverName, + toolKey: `${type}:${serverName}`, + config: { serverName }, +}); + +function setup(signal?: AbortSignal) { + return { + userId: 'user-id', + context: { user: { id: 'safe-user-id' }, signal }, + requestedTools: { first: [request('first')] }, + getAvailableTools: jest.fn().mockResolvedValue(undefined), + createTools: jest.fn().mockResolvedValue([{ name: 'bulk' }]), + createTool: jest.fn().mockResolvedValue({ name: 'selected' }), + }; +} + +describe('loadMCPTools', () => { + it('does not schedule any work after an already-cancelled request', async () => { + const controller = new AbortController(); + controller.abort(new Error('stopped')); + const options = setup(controller.signal); + await expect(loadMCPTools(options)).rejects.toBe(controller.signal.reason); + expect(options.createTools).not.toHaveBeenCalled(); + }); + + it.each([false, true])( + 'settles siblings then prefers cancellation: authFailure=%s', + async (authFailure) => { + const controller = new AbortController(); + const options = setup(controller.signal); + const stopped = new Error('stopped'); + let settled = false; + options.createTools.mockImplementationOnce(async () => { + if (authFailure) { + throw new OpenIDReauthRequiredError('Please sign in again'); + } + return []; + }); + options.createTools.mockImplementationOnce( + () => + new Promise((resolve) => + setImmediate(() => { + settled = true; + controller.abort(stopped); + resolve([]); + }), + ), + ); + await expect( + loadMCPTools({ + ...options, + requestedTools: { + first: [request('first')], + second: [request('second')], + }, + }), + ).rejects.toBe(stopped); + expect(settled).toBe(true); + }, + ); + + it('does not create selected tools after cancellation during catalog loading', async () => { + const controller = new AbortController(); + const options = setup(controller.signal); + options.getAvailableTools.mockImplementationOnce(async () => { + controller.abort(); + return {}; + }); + await expect( + loadMCPTools({ + ...options, + requestedTools: { + first: [request('first', 'single')], + }, + }), + ).rejects.toBe(controller.signal.reason); + expect(options.createTool).not.toHaveBeenCalled(); + }); + + it('preserves the first typed error without treating ordinary optional errors as fatal', async () => { + const options = setup(); + const error = new OpenIDReauthRequiredError('Please sign in again'); + options.createTools.mockRejectedValueOnce(new Error('optional tool unavailable')); + await expect(loadMCPTools(options)).resolves.toEqual([]); + options.createTools.mockRejectedValueOnce(error); + await expect(loadMCPTools(options)).rejects.toBe(error); + await expect(loadMCPTools(options)).resolves.toEqual([{ name: 'bulk' }]); + }); + + it('reuses the refreshed catalog for sequential tools without changing parameters', async () => { + const options = setup(); + const catalog = { search: true }; + options.createTool.mockImplementationOnce(async ({ onAvailableTools }) => { + onAvailableTools(catalog); + return { name: 'one' }; + }); + await loadMCPTools({ + ...options, + requestedTools: { + first: [request('first', 'single'), request('first', 'single')], + }, + }); + expect(options.getAvailableTools).toHaveBeenCalledTimes(1); + expect(options.getAvailableTools).toHaveBeenCalledWith('user-id', 'first', { + serverName: 'first', + }); + expect(options.createTool).toHaveBeenLastCalledWith( + expect.objectContaining({ + user: { id: 'safe-user-id' }, + index: 0, + availableTools: catalog, + config: { serverName: 'first' }, + }), + ); + }); +}); diff --git a/packages/api/src/mcp/loading.ts b/packages/api/src/mcp/loading.ts new file mode 100644 index 00000000000..30e4fe88225 --- /dev/null +++ b/packages/api/src/mcp/loading.ts @@ -0,0 +1,131 @@ +import { logger } from '@librechat/data-schemas'; +import { isMCPInitializationError } from './errors'; + +interface MCPToolRequest { + type: 'all' | 'single'; + toolKey: string; + serverName: string; + config: Config; +} + +interface MCPToolParameters { + index: number; + serverName: string; + config: Config; +} + +interface MCPLoadContext { + signal?: AbortSignal; +} + +interface MCPToolLoadOptions { + userId: string; + requestedTools: Record[]>; + context: Context; + availableTools?: Record; + getAvailableTools: ( + userId: string, + server: string, + config: Config, + ) => Promise; + createTools: (params: Context & MCPToolParameters) => Promise; + createTool: ( + params: Context & + MCPToolParameters & { + toolKey: string; + availableTools?: Catalog; + onAvailableTools: (tools: Catalog) => void; + }, + ) => Promise; +} + +/** + * Bulk loads overlap while selected tools reuse each server's catalog in order. + * Settle every started load before cleanup can begin, then prefer request cancellation + * over any saved credential failure. Ordinary optional-tool failures stay soft. + */ +export async function loadMCPTools({ + userId, + requestedTools, + context, + availableTools: catalogs, + getAvailableTools, + createTools, + createTool, +}: MCPToolLoadOptions): Promise { + const { signal } = context; + const loaded: Tool[] = []; + const pending: Promise[] = []; + const failedServers = new Set(); + let failure: { error: unknown } | undefined; + const capture = (error: unknown) => { + if (!failure && isMCPInitializationError(error, signal)) { + failure = { error }; + } + }; + + let index = -1; + for (const [serverName, configs] of Object.entries(requestedTools)) { + index++; + let availableTools = catalogs?.[serverName]; + for (const config of configs) { + if (signal?.aborted || failedServers.has(serverName)) { + break; + } + try { + const params = { ...context, index, serverName: config.serverName, config: config.config }; + if (config.type === 'all' && configs.length === 1) { + pending.push( + createTools(params).catch((error: unknown) => { + capture(error); + logger.error(`Error loading ${serverName} tools:`, error); + return null; + }), + ); + continue; + } + if (!availableTools) { + try { + availableTools = await getAvailableTools(userId, serverName, config.config); + } catch (error) { + capture(error); + logger.error(`Error fetching available tools for MCP server ${serverName}:`, error); + } + } + if (signal?.aborted) { + break; + } + const tool = + config.type === 'all' + ? await createTools(params) + : await createTool({ + ...params, + availableTools, + toolKey: config.toolKey, + onAvailableTools: (tools) => { + availableTools = tools; + }, + }); + if (Array.isArray(tool)) { + loaded.push(...tool); + } else if (tool) { + loaded.push(tool); + } else { + failedServers.add(serverName); + logger.warn( + `MCP tool creation failed for "${config.toolKey}", server may be unavailable or unauthenticated.`, + ); + } + } catch (error) { + capture(error); + logger.error(`Error loading MCP tool for server ${serverName}:`, error); + } + } + } + loaded.push(...(await Promise.all(pending)).flatMap((tools) => tools ?? [])); + signal?.throwIfAborted(); + if (failure) { + throw failure.error; + } + return loaded; +} diff --git a/packages/api/src/mcp/signal.spec.ts b/packages/api/src/mcp/signal.spec.ts new file mode 100644 index 00000000000..59a9e6c86a4 --- /dev/null +++ b/packages/api/src/mcp/signal.spec.ts @@ -0,0 +1,49 @@ +import { withMCPRequestSignal } from './signal'; + +describe('withMCPRequestSignal', () => { + it('relays cancellation to an in-flight request and detaches on rejection', async () => { + const parent = new AbortController(); + const reason = new Error('Stop'); + const remove = jest.spyOn(parent.signal, 'removeEventListener'); + const request = jest.fn( + (signal?: AbortSignal) => + new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { once: true }); + }), + ); + const pending = withMCPRequestSignal(parent.signal, request); + + expect(request).toHaveBeenCalledTimes(1); + expect(request.mock.calls[0][0]).not.toBe(parent.signal); + parent.abort(reason); + await expect(pending).rejects.toBe(reason); + expect(remove).toHaveBeenCalledWith('abort', expect.any(Function)); + }); + + it('does not abort a completed request when its parent is aborted later', async () => { + const parent = new AbortController(); + const request = jest.fn(async (signal?: AbortSignal) => signal); + const child = await withMCPRequestSignal(parent.signal, request); + + expect(child).toBeDefined(); + parent.abort(); + expect(child?.aborted).toBe(false); + }); + + it('does not start a request whose parent is already aborted', async () => { + const parent = new AbortController(); + const reason = new Error('Stop'); + parent.abort(reason); + const request = jest.fn(async () => undefined); + + await expect(withMCPRequestSignal(parent.signal, request)).rejects.toBe(reason); + expect(request).not.toHaveBeenCalled(); + }); + + it('passes through requests without a parent signal', async () => { + const request = jest.fn(async (signal?: AbortSignal) => signal); + + await expect(withMCPRequestSignal(undefined, request)).resolves.toBeUndefined(); + expect(request).toHaveBeenCalledWith(undefined); + }); +}); diff --git a/packages/api/src/mcp/signal.ts b/packages/api/src/mcp/signal.ts new file mode 100644 index 00000000000..db42118e829 --- /dev/null +++ b/packages/api/src/mcp/signal.ts @@ -0,0 +1,23 @@ +/** The MCP SDK retains its abort listener after a response; never give it a long-lived run signal. */ +export async function withMCPRequestSignal( + parent: AbortSignal | undefined, + request: (signal: AbortSignal | undefined) => Promise, +): Promise { + if (!parent) { + return request(undefined); + } + + parent.throwIfAborted(); + const controller = new AbortController(); + const onAbort = () => controller.abort(parent.reason); + parent.addEventListener('abort', onAbort, { once: true }); + try { + if (parent.aborted) { + onAbort(); + parent.throwIfAborted(); + } + return await request(controller.signal); + } finally { + parent.removeEventListener('abort', onAbort); + } +} diff --git a/packages/api/src/protection/detectors/pattern.spec.ts b/packages/api/src/protection/detectors/pattern.spec.ts index 992eb6e593a..cc84a1a6adf 100644 --- a/packages/api/src/protection/detectors/pattern.spec.ts +++ b/packages/api/src/protection/detectors/pattern.spec.ts @@ -72,6 +72,32 @@ describe('pattern content inspector', () => { ).toMatchObject({ ruleId: 'bearer_header' }); }); + it.each([ + { config: { starterPatterns: ['bearer_header'] }, text: 'Bearer contract-token' }, + { + config: { + starterPatterns: [], + customPatterns: [{ id: 'custom', label: 'Custom', regex: 'ORG-[A-Z]+' }], + }, + text: 'ORG-SECRET', + }, + ])('enforces a zero remaining allowance when locating $text', ({ config, text }) => { + const inspector = createPatternContentInspector(config, { linearTime: true }); + + expect(inspector.locate('clean text', 0)).toEqual([]); + expect(inspector.locate('', 0)).toEqual([]); + expect(() => inspector.locate(text, 0)).toThrow(PatternConfigurationError); + expect(inspector.locate(text, 1)).toHaveLength(1); + }); + + it.each([-1, 0.5, Number.NaN, Number.POSITIVE_INFINITY, Number.MAX_SAFE_INTEGER + 1])( + 'rejects invalid remaining match allowance %s before inspecting', + (allowance) => { + const inspector = createPatternContentInspector({}, { linearTime: true }); + expect(() => inspector.locate('clean text', allowance)).toThrow(PatternConfigurationError); + }, + ); + it('memoizes a compiled set by config identity and memory limit', () => { const compile = jest.spyOn(RE2Set.prototype, 'compile'); const config: MessageFilterPiiConfig = { diff --git a/packages/api/src/protection/detectors/pattern.ts b/packages/api/src/protection/detectors/pattern.ts index 373028a3a2e..6456b931be5 100644 --- a/packages/api/src/protection/detectors/pattern.ts +++ b/packages/api/src/protection/detectors/pattern.ts @@ -1,4 +1,4 @@ -import { RE2Set } from 're2js'; +import { RE2JS, RE2Set } from 're2js'; import { logger } from '@librechat/data-schemas'; import { MAX_PII_CUSTOM_REGEX_CHARACTERS, @@ -9,23 +9,25 @@ import { MAX_PII_PATTERNS_PER_SOURCE, getPiiRegexProgramSize, } from 'librechat-data-provider'; -import type { MessageFilterPiiConfig, FilterPiiCustomPatternConfig } from 'librechat-data-provider'; +import type { + FilterPiiCategory, + MessageFilterPiiConfig, + FilterPiiCustomPatternConfig, +} from 'librechat-data-provider'; import type { ProtectionFinding, TextContentFragment } from '../types'; -interface TestablePattern { - test(input: string): boolean; -} - interface CompiledPattern { readonly id: string; readonly label: string; - readonly pattern: TestablePattern; + readonly pattern: RegExp; + readonly category: FilterPiiCategory; } interface PreparedCustomPattern { readonly id: string; readonly label: string; readonly regex: string; + readonly category: FilterPiiCategory; } interface SnapshotPatternContentInspectorConfig { @@ -67,10 +69,18 @@ export interface PatternContentInspectorPreflightCost { readonly regexes: readonly string[]; } +export interface PatternTextMatch { + readonly start: number; + readonly end: number; + readonly category: FilterPiiCategory; +} + export interface PatternContentInspector { readonly active: boolean; inspectFragment(fragment: TextContentFragment): ProtectionFinding | null; inspect(fragments: Iterable): ProtectionFinding | null; + /** Internal match locations, never safe to return as public finding metadata. */ + locate(text: string, maxMatches: number): readonly PatternTextMatch[]; } export interface PatternContentInspectorConfig { @@ -86,9 +96,24 @@ export interface PatternContentInspectorOptions { } const STARTER_PATTERNS: readonly CompiledPattern[] = [ - { id: 'sk_prefix', label: 'sk- prefix token', pattern: /\b(sk-)[a-zA-Z0-9_-]+/ }, - { id: 'bearer_header', label: 'Bearer token', pattern: /\b(Bearer )[^\s"']+/i }, - { id: 'api_key_header', label: 'api-key header', pattern: /\b(api-key:?\s+)[^\s"']+/i }, + { + id: 'sk_prefix', + label: 'sk- prefix token', + pattern: /\b(sk-)[a-zA-Z0-9_-]+/, + category: 'credential', + }, + { + id: 'bearer_header', + label: 'Bearer token', + pattern: /\b(Bearer )[^\s"']+/i, + category: 'credential', + }, + { + id: 'api_key_header', + label: 'api-key header', + pattern: /\b(api-key:?\s+)[^\s"']+/i, + category: 'credential', + }, ]; const STARTER_BY_ID = new Map(STARTER_PATTERNS.map((pattern) => [pattern.id, pattern])); @@ -201,7 +226,25 @@ function readCustomPattern(candidate: unknown, index: number): FilterPiiCustomPa if (typeof regex !== 'string' || regex.length === 0 || regex.length > MAX_PII_PATTERN_LENGTH) { throw configurationError(`customPatterns[${index}].regex is invalid`); } - return { id, label, regex }; + let category: unknown; + try { + category = (candidate as FilterPiiCustomPatternConfig).category; + } catch { + throw configurationError(`customPatterns[${index}] could not be read safely`); + } + if ( + category != null && + (typeof category !== 'string' || + !['email', 'phone', 'name', 'credential', 'custom'].includes(category)) + ) { + throw configurationError(`customPatterns[${index}].category is invalid`); + } + return { + id, + label, + regex, + ...(category == null ? {} : { category: category as FilterPiiCategory }), + }; } function snapshotConfig( @@ -240,7 +283,7 @@ function snapshotConfig( throw configurationError('customPatterns could not be read safely'); } const pattern = readCustomPattern(patternCandidate, index); - custom.push(pattern); + custom.push({ ...pattern, category: pattern.category ?? 'custom' }); regexes.push(pattern.regex); regexCharacters += pattern.regex.length; if (regexCharacters > MAX_PII_CUSTOM_REGEX_CHARACTERS) { @@ -357,10 +400,17 @@ function createSequentialInspector(patterns: readonly CompiledPattern[]): Patter function createInspector( active: boolean, inspectFragment: (fragment: TextContentFragment) => ProtectionFinding | null, + locate?: (text: string, maxMatches: number) => readonly PatternTextMatch[], ): PatternContentInspector { return { active, inspectFragment, + locate(text, maxMatches) { + if (locate == null) { + throw configurationError('redaction requires a linear-time pattern inspector'); + } + return locate(text, maxMatches); + }, inspect(fragments) { for (const fragment of fragments) { const finding = inspectFragment(fragment); @@ -416,7 +466,44 @@ function createLinearInspector( return pattern == null ? null : findingFor(pattern, fragment); }; - return createInspector(prepared.cost.active, inspectFragment); + const compiled = new Map(); + const locate = (text: string, maxMatches: number): readonly PatternTextMatch[] => { + if (!Number.isSafeInteger(maxMatches) || maxMatches < 0) { + throw configurationError('redaction match limit must be a non-negative safe integer'); + } + const matches: PatternTextMatch[] = []; + const selected = customSet == null ? [] : customSet.match(text); + const candidates: Array = prepared.starter.map( + (pattern) => [ + `${pattern.pattern.ignoreCase ? '(?i)' : ''}${pattern.pattern.source}`, + pattern.category, + ], + ); + for (const index of selected) { + const pattern = prepared.custom[index]; + if (pattern != null) { + candidates.push([pattern.regex, pattern.category]); + } + } + for (const [expression, category] of candidates) { + let regex = compiled.get(expression); + if (regex == null) { + regex = RE2JS.compile(expression); + compiled.set(expression, regex); + } + const matcher = regex.matcher(text); + while (matcher.find()) { + const start = matcher.start(); + const end = matcher.end(); + if (end <= start || matches.length >= maxMatches) { + throw configurationError('redaction match limit or empty match encountered'); + } + matches.push({ start, end, category }); + } + } + return matches; + }; + return createInspector(prepared.cost.active, inspectFragment, locate); } function cacheLinearInspector( @@ -480,6 +567,7 @@ export function createPatternContentInspector( id: pattern.id, label: pattern.label, pattern: new RegExp(pattern.regex), + category: pattern.category, })); const inspector = createSequentialInspector([...prepared.starter, ...custom]); if (options.cacheResult !== false) { diff --git a/packages/api/src/protection/index.ts b/packages/api/src/protection/index.ts index 18e325633aa..264f49b381b 100644 --- a/packages/api/src/protection/index.ts +++ b/packages/api/src/protection/index.ts @@ -1,5 +1,6 @@ export * from './types'; export * from './runtime'; +export * from './transform'; export * from './title'; export * from './legacy'; export * from './provenance'; diff --git a/packages/api/src/protection/transform.spec.ts b/packages/api/src/protection/transform.spec.ts new file mode 100644 index 00000000000..88022c29816 --- /dev/null +++ b/packages/api/src/protection/transform.spec.ts @@ -0,0 +1,227 @@ +import type { FiltersConfig } from 'librechat-data-provider'; +import type { TextContentFragment } from './types'; +import { createPiiTextTransformer, PiiTransformationError } from './transform'; +import { inspectContent } from './runtime'; + +jest.mock('@librechat/data-schemas', () => ({ + logger: { warn: jest.fn(), error: jest.fn(), info: jest.fn(), debug: jest.fn() }, +})); + +const pii: NonNullable['pii']> = { + action: 'redact' as const, + starterPatterns: [], + fields: ['text'], + customPatterns: [ + { + id: 'email', + label: 'Email address', + regex: '(?i)[a-z]+@[a-z]+\\.[a-z]+', + category: 'email' as const, + }, + ], +}; + +function fragment(text: string, extras: Partial = {}): TextContentFragment { + return { + id: 'chat.text', + path: '/text', + text, + source: 'message', + field: 'text', + format: 'plain', + treatment: 'replaceable', + provenance: 'user', + ...extras, + } as TextContentFragment; +} + +describe('opt-in PII text transformation', () => { + it('leaves the existing inspection path blocking until a caller integrates transformation', () => { + const filters: FiltersConfig = { messages: { pii } }; + expect(inspectContent([fragment('Alice@Example.com')], { filters })).toMatchObject({ + detectorId: 'pii-pattern', + ruleId: 'email', + }); + }); + + it('uses typed placeholders and returns no matched values in metadata', () => { + const session = createPiiTextTransformer(pii).createSession(); + const first = session.transform(fragment('😀 Alice@Example.com, Alice@Example.com')); + const second = session.transform(fragment('Next: Bob@Example.com and Alice@Example.com')); + + expect(first).toEqual({ + version: 1, + content: '😀 [EMAIL_1], [EMAIL_1]', + replacements: 2, + categories: [{ category: 'EMAIL', count: 2 }], + }); + expect(second.content).toBe('Next: [EMAIL_2] and [EMAIL_1]'); + expect(JSON.stringify(first.categories)).not.toContain('Alice'); + expect(JSON.stringify(second.categories)).not.toContain('Bob'); + const newTurn = createPiiTextTransformer(pii).createSession(); + expect(newTurn.transform(fragment('Bob@Example.com')).content).toBe('[EMAIL_1]'); + }); + + it('never aliases a user-supplied placeholder to another private value', () => { + const session = createPiiTextTransformer(pii).createSession(); + expect(session.transform(fragment('[EMAIL_1] Alice@Example.com')).content).toBe( + '[EMAIL_1] [EMAIL_2]', + ); + expect(() => session.transform(fragment('A forged [EMAIL_2] marker'))).toThrow( + PiiTransformationError, + ); + }); + + it('reserves literal markers across selected fragments without sharing session state', () => { + const transformer = createPiiTextTransformer(pii); + const session = transformer.createSession(); + const independent = transformer.createSession(); + + expect(session.transform(fragment('[EMAIL_1] [EMAIL_3]'))).toEqual({ + version: 1, + content: '[EMAIL_1] [EMAIL_3]', + replacements: 0, + categories: [], + }); + expect(session.transform(fragment('Alice@Example.com')).content).toBe('[EMAIL_2]'); + expect(independent.transform(fragment('Alice@Example.com')).content).toBe('[EMAIL_1]'); + expect(session.transform(fragment('Bob@Example.com Alice@Example.com')).content).toBe( + '[EMAIL_4] [EMAIL_2]', + ); + expect(session.transform(fragment('[EMAIL_1]'))).toMatchObject({ replacements: 0 }); + expect(() => session.transform(fragment('[EMAIL_4]'))).toThrow(PiiTransformationError); + }); + + it('still charges selected literal markers against the character budget', () => { + const session = createPiiTextTransformer({ ...pii, maxCharacters: 18 }).createSession(); + expect(session.transform(fragment('[EMAIL_1]')).replacements).toBe(0); + expect(session.transform(fragment('[EMAIL_2]')).replacements).toBe(0); + expect(() => session.transform(fragment('[EMAIL_3]'))).toThrow(PiiTransformationError); + }); + + it('does not charge an excluded field against the selected character budget', () => { + const text = 'Alice@Example.com'; + const session = createPiiTextTransformer({ + ...pii, + maxCharacters: text.length, + }).createSession(); + const excluded = fragment(text.repeat(100), { field: 'summary' }); + + expect(session.transform(excluded)).toEqual({ + version: 1, + content: excluded.text, + replacements: 0, + categories: [], + }); + expect(session.transform(fragment(text)).content).toBe('[EMAIL_1]'); + }); + + it('neither validates nor reserves excluded markers before or after selected fragments', () => { + const session = createPiiTextTransformer(pii).createSession(); + const excluded = fragment('[EMAIL_1] [EMAIL_2]', { field: 'summary' }); + + expect(session.transform(excluded).content).toBe(excluded.text); + expect(session.transform(fragment('Alice@Example.com')).content).toBe('[EMAIL_1]'); + expect(session.transform(excluded).content).toBe(excluded.text); + expect(session.transform(fragment('Bob@Example.com')).content).toBe('[EMAIL_2]'); + }); + + it('accepts clean and empty fragments at the match limit but rejects one more match', () => { + const session = createPiiTextTransformer({ ...pii, maxMatches: 1 }).createSession(); + + expect(session.transform(fragment('Alice@Example.com')).replacements).toBe(1); + for (const text of ['No private details here.', '']) { + expect(session.transform(fragment(text))).toEqual({ + version: 1, + content: text, + replacements: 0, + categories: [], + }); + } + expect(() => session.transform(fragment('Alice@Example.com'))).toThrow(PiiTransformationError); + }); + + it('redacts selected built-in credential headers without leaking the value', () => { + const session = createPiiTextTransformer({ + action: 'redact', + starterPatterns: ['bearer_header'], + }).createSession(); + const result = session.transform(fragment('Authorization: bEaReR contract-token')); + expect(result.content).toBe('Authorization: [CREDENTIAL_1]'); + expect(JSON.stringify(result.categories)).not.toContain('contract-token'); + }); + + it('merges overlapping patterns so partial sensitive values never survive', () => { + const config = { + action: 'redact' as const, + starterPatterns: [], + customPatterns: [ + { id: 'early', label: 'early', regex: 'ORG-SECRET', category: 'credential' as const }, + { id: 'later', label: 'later', regex: 'SECRET-DATA', category: 'custom' as const }, + ], + }; + const result = createPiiTextTransformer(config) + .createSession() + .transform(fragment('prefix ORG-SECRET-DATA suffix')); + expect(result.content).toBe('prefix [CREDENTIAL_1] suffix'); + expect(result.replacements).toBe(1); + expect(result.content).not.toContain('SECRET'); + }); + + it('respects configured fields without attempting to replace other fields', () => { + const session = createPiiTextTransformer(pii).createSession(); + expect(session.transform(fragment('Alice@Example.com', { field: 'summary' })).content).toBe( + 'Alice@Example.com', + ); + expect(session.transform(fragment('Alice@Example.com')).content).toBe('[EMAIL_1]'); + }); + + it('fails closed for non-replaceable matches and structured or URI content', () => { + for (const extras of [ + { treatment: 'inspect_only' as const }, + { format: 'json' as const }, + { format: 'uri' as const }, + ]) { + expect(() => + createPiiTextTransformer(pii) + .createSession() + .transform(fragment('Alice@Example.com', extras)), + ).toThrow(PiiTransformationError); + } + }); + + it('bounds characters and matches across all fragments in one session', () => { + const config = { + ...pii, + maxCharacters: 50, + maxMatches: 2, + }; + const session = createPiiTextTransformer(config).createSession(); + expect(session.transform(fragment('A@b.com A@b.com')).replacements).toBe(2); + expect(() => session.transform(fragment('C@d.com'))).toThrow(PiiTransformationError); + expect(() => + createPiiTextTransformer({ ...config, maxCharacters: 4 }) + .createSession() + .transform(fragment('Alice@Example.com')), + ).toThrow(PiiTransformationError); + }); + + it('rejects empty regex matches, invalid limits and missing redact action without echoing input', () => { + expect(() => createPiiTextTransformer({ ...pii, maxMatches: -1 })).toThrow(); + expect(() => createPiiTextTransformer({ ...pii, maxMatches: 0 })).toThrow(); + expect(() => createPiiTextTransformer({ ...pii, action: 'block' })).toThrow(); + const session = createPiiTextTransformer({ + action: 'redact', + starterPatterns: [], + customPatterns: [{ id: 'empty', label: 'empty', regex: 'a*' }], + }).createSession(); + let failure: Error | undefined; + try { + session.transform(fragment('secret value')); + } catch (error) { + failure = error as Error; + } + expect(failure).toBeInstanceOf(PiiTransformationError); + expect(failure?.message).not.toContain('secret value'); + }); +}); diff --git a/packages/api/src/protection/transform.ts b/packages/api/src/protection/transform.ts new file mode 100644 index 00000000000..ddfd9695ae0 --- /dev/null +++ b/packages/api/src/protection/transform.ts @@ -0,0 +1,184 @@ +import { + DEFAULT_PII_REDACTION_MAX_CHARACTERS, + DEFAULT_PII_REDACTION_MAX_MATCHES, + MAX_PII_PATTERNS_PER_SOURCE, +} from 'librechat-data-provider'; +import type { FilterPiiAction, FilterPiiCategory } from 'librechat-data-provider'; +import type { PatternContentInspectorConfig, PatternTextMatch } from './detectors/pattern'; +import type { TextContentFragment } from './types'; +import { createPatternContentInspector, PatternConfigurationError } from './detectors/pattern'; + +export type PiiPlaceholderType = Uppercase; + +export interface PiiRedactionConfig extends PatternContentInspectorConfig { + readonly action?: FilterPiiAction; + readonly fields?: readonly string[]; + readonly maxCharacters?: number; + readonly maxMatches?: number; +} + +export interface PiiRedactionMetadata { + readonly category: PiiPlaceholderType; + readonly count: number; +} + +export interface PiiTransformationResult { + readonly version: 1; + readonly content: string; + readonly replacements: number; + readonly categories: readonly PiiRedactionMetadata[]; +} + +export class PiiTransformationError extends Error { + constructor(reason: 'limit' | 'unreplaceable' | 'inspection') { + super(`PII transformation failed: ${reason}`); + this.name = 'PiiTransformationError'; + } +} + +export interface PiiTransformationSession { + /** Applies this rule to selected text fields; the caller must still enforce all other policies. */ + transform(fragment: TextContentFragment): PiiTransformationResult; +} + +function readLimit(value: number | undefined, fallback: number, maximum: number): number { + if (value == null) { + return fallback; + } + if (!Number.isSafeInteger(value) || value <= 0 || value > maximum) { + throw new PatternConfigurationError('Invalid PII transformation limit'); + } + return value; +} + +function mergedMatches(matches: readonly PatternTextMatch[]): PatternTextMatch[] { + const sorted = [...matches].sort( + (left, right) => + left.start - right.start || + right.end - left.end || + left.category.localeCompare(right.category), + ); + const merged: PatternTextMatch[] = []; + for (const match of sorted) { + const previous = merged[merged.length - 1]; + if (previous == null || match.start >= previous.end) { + merged.push(match); + } else if (match.end > previous.end) { + merged[merged.length - 1] = { ...previous, end: match.end }; + } + } + return merged; +} + +/** + * Explicit, opt-in transformation. Legacy inspectors still reject `redact` findings + * until the authenticated submission and storage pipeline consumes this API. + * A session owns only ephemeral, bounded mappings; never log or persist it. + */ +export function createPiiTextTransformer(config: PiiRedactionConfig): { + createSession(): PiiTransformationSession; +} { + if (config.action !== 'redact') { + throw new PatternConfigurationError('PII transformation requires the redact action'); + } + const maxCharacters = readLimit( + config.maxCharacters, + DEFAULT_PII_REDACTION_MAX_CHARACTERS, + 262_144, + ); + const maxMatches = readLimit(config.maxMatches, DEFAULT_PII_REDACTION_MAX_MATCHES, 4_096); + let fields: ReadonlySet | undefined; + if (config.fields != null) { + try { + if (!Array.isArray(config.fields) || config.fields.length > MAX_PII_PATTERNS_PER_SOURCE) { + throw new PatternConfigurationError('Invalid PII transformation fields'); + } + fields = new Set(config.fields); + if (fields.size === 0 || [...fields].some((field) => typeof field !== 'string')) { + throw new PatternConfigurationError('Invalid PII transformation fields'); + } + } catch { + throw new PatternConfigurationError('Invalid PII transformation fields'); + } + } + const inspector = createPatternContentInspector(config, { linearTime: true }); + + return { + createSession() { + const placeholders = new Map(); + const issued = new Set(); + const reserved = new Set(); + const nextByCategory = new Map(); + let remainingCharacters = maxCharacters; + let remainingMatches = maxMatches; + return { + transform(fragment) { + let matches: readonly PatternTextMatch[]; + try { + if (fields != null && !fields.has(fragment.field)) { + return { version: 1, content: fragment.text, replacements: 0, categories: [] }; + } + if (typeof fragment.text !== 'string' || fragment.text.length > remainingCharacters) { + throw new PiiTransformationError('limit'); + } + remainingCharacters -= fragment.text.length; + for (const marker of fragment.text.matchAll( + /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_[1-9]\d*\]/g, + )) { + reserved.add(marker[0]); + if (issued.has(marker[0])) { + throw new PiiTransformationError('inspection'); + } + } + matches = inspector.locate(fragment.text, remainingMatches); + } catch (error) { + if (error instanceof PiiTransformationError) { + throw error; + } + throw new PiiTransformationError('inspection'); + } + if (matches.length === 0) { + return { version: 1, content: fragment.text, replacements: 0, categories: [] }; + } + if ( + fragment.treatment !== 'replaceable' || + !['plain', 'markdown'].includes(fragment.format) + ) { + throw new PiiTransformationError('unreplaceable'); + } + const merged = mergedMatches(matches); + remainingMatches -= matches.length; + const categories = new Map(); + const parts: string[] = []; + let offset = 0; + for (const { start, end, category } of merged) { + const type = category.toUpperCase() as PiiPlaceholderType; + const value = fragment.text.slice(start, end); + const key = JSON.stringify([type, value]); + let placeholder = placeholders.get(key); + if (placeholder == null) { + let next = nextByCategory.get(type) ?? 0; + do { + next++; + placeholder = `[${type}_${next}]`; + } while (reserved.has(placeholder) || issued.has(placeholder)); + nextByCategory.set(type, next); + placeholders.set(key, placeholder); + issued.add(placeholder); + } + categories.set(type, (categories.get(type) ?? 0) + 1); + parts.push(fragment.text.slice(offset, start), placeholder); + offset = end; + } + parts.push(fragment.text.slice(offset)); + return { + version: 1, + content: parts.join(''), + replacements: merged.length, + categories: [...categories].map(([category, count]) => ({ category, count })), + }; + }, + }; + }, + }; +} diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 991ee1b882c..a54ee838beb 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -20,6 +20,7 @@ import type { import type { StandardGraph } from '@librechat/agents'; import type { SerializableJobData, + GenerationSettlementState, CreatedJobData, IEventTransport, UsageMetadata, @@ -568,10 +569,22 @@ export interface CreateGenerationJobOptions { * receiving this claim, then pass the same object to {@link finishTerminalJob} * from a `finally` block. */ +/** A generation reached a terminal state and released its runtime. */ +export interface GenerationSettledEvent { + streamId: string; + conversationId: string; + userId: string; + status: TerminalJobClaim['status']; +} + +export type GenerationSettledListener = (event: GenerationSettledEvent) => void; + export interface TerminalJobClaim { readonly streamId: string; readonly createdAt: number; readonly conversationId?: string; + /** The generation's owner, so settlement can be announced to that principal's waiters. */ + readonly userId?: string; readonly status: 'complete' | 'error' | 'aborted'; readonly error?: string; /** The winner must durably publish either its normal FINAL or a @@ -786,6 +799,7 @@ class GenerationJobManagerClass { /** Makes terminal cleanup idempotent while keeping claims opaque to callers. */ private terminalFinishPromises = new WeakMap>(); + private generationSettledListeners = new Set(); /** Exact local runtime observed when a claim won; never clean a later runtime. */ private terminalClaimRuntimes = new WeakMap(); @@ -1078,6 +1092,7 @@ class GenerationJobManagerClass { createdAt, ...(job?.createdAt === createdAt && job.conversationId != null && { conversationId: job.conversationId }), + ...(job?.createdAt === createdAt && job.userId != null && { userId: job.userId }), status: 'error' as const, error, drainedSteers: Object.freeze([...drainedSteers]), @@ -2256,7 +2271,7 @@ class GenerationJobManagerClass { streamId: string, job: Pick< SerializableJobData, - 'createdAt' | 'conversationId' | 'providerExecutionId' | 'agentEventDeliveryKey' + 'createdAt' | 'conversationId' | 'providerExecutionId' | 'agentEventDeliveryKey' | 'userId' >, message: string, ): Promise { @@ -2288,6 +2303,13 @@ class GenerationJobManagerClass { expectCreatedAt: job.createdAt, }) ) { + /** A direct terminal transition builds no claim, so it announces itself. */ + this.notifyGenerationSettled({ + streamId, + conversationId: job.conversationId, + userId: job.userId, + status: 'error', + }); return true; } } catch (error) { @@ -2298,6 +2320,15 @@ class GenerationJobManagerClass { try { const current = await this.jobStore.getJob(streamId); + if (current?.createdAt === job.createdAt && current.status === 'error') { + this.notifyGenerationSettled({ + streamId, + conversationId: current.conversationId, + userId: current.userId, + status: 'error', + }); + return true; + } if ( current == null || current.createdAt !== job.createdAt || @@ -2396,7 +2427,7 @@ class GenerationJobManagerClass { (options.recoveredSteerPayload != null && !isRecoveredSteerPayload(options.recoveredSteerPayload)) ) { - throw new RecoveredSteerPayloadMismatchError(); + throw new RecoveredSteerPayloadMismatchError('invalid_payload'); } // Capture the active epoch before the store atomically replaces it. A // subscriber attached to that predecessor filters events by generation, @@ -3742,9 +3773,21 @@ class GenerationJobManagerClass { await this.jobStore.releaseIdempotencyKey(legacyKey, expectedClaim); } - /** - * Get job status. - */ + /** Observes identity and final-save ownership without attaching a runtime or + * promoting a slow terminal writer to stale-owner recovery. */ + async getGenerationSettlementState( + streamId: string, + ): Promise { + const job = await this.jobStore.getJob(streamId); + if (job == null) return undefined; + return { + createdAt: job.createdAt, + status: job.status, + terminalPersistencePending: job.terminalPersistencePending, + }; + } + + /** Get job status. */ async getJobStatus(streamId: string): Promise { const jobData = await this.jobStore.getJob(streamId); return jobData?.status as t.GenerationJobStatus | undefined; @@ -3901,6 +3944,7 @@ class GenerationJobManagerClass { ...(jobData.conversationId != null && { conversationId: jobData.conversationId, }), + ...(jobData.userId != null && { userId: jobData.userId }), status, ...(terminalError != null && { error: terminalError }), ...(options.persistencePending === true && { @@ -4071,6 +4115,40 @@ class GenerationJobManagerClass { * claim is idempotent, and every local mutation is pinned to the runtime * object and generation epoch captured when the CAS won. */ + /** + * Calls `listener` after each generation owned by this process reaches a + * terminal state and its runtime is released — by completion, error, or + * abort. Listeners run synchronously and must not throw; failures are logged + * and never affect terminal cleanup. Returns an unsubscribe function. + */ + onGenerationSettled(listener: GenerationSettledListener): () => void { + this.generationSettledListeners.add(listener); + return () => { + this.generationSettledListeners.delete(listener); + }; + } + + private notifyGenerationSettled( + target: Pick, + ): void { + if (target.userId == null || this.generationSettledListeners.size === 0) { + return; + } + const event: GenerationSettledEvent = { + streamId: target.streamId, + conversationId: target.conversationId ?? target.streamId, + userId: target.userId, + status: target.status, + }; + for (const listener of this.generationSettledListeners) { + try { + listener(event); + } catch (listenerError) { + logger.error('[GenerationJobManager] Generation settled listener failed', listenerError); + } + } + } + finishTerminalJob(claim: TerminalJobClaim): Promise { const inFlight = this.terminalFinishPromises.get(claim); if (inFlight) { @@ -4260,6 +4338,7 @@ class GenerationJobManagerClass { metricStatus = 'error'; } recordGenerationJob(this.storeLabel, metricStatus); + this.notifyGenerationSettled(claim); } if (cleanupError != null) { @@ -4616,6 +4695,7 @@ class GenerationJobManagerClass { ...(jobData.conversationId != null && { conversationId: jobData.conversationId, }), + ...(jobData.userId != null && { userId: jobData.userId }), status: 'aborted', persistencePending: true, drainedSteers: Object.freeze([...drainedSteers]), @@ -8865,6 +8945,17 @@ class GenerationJobManagerClass { await this.runApprovalExpiredHandler(streamId, expiredJob); await this.notifyApprovalExpiredRuntime(streamId, expiredJob.createdAt, observedRuntime); + /** Expiry is a direct `requires_action -> aborted` transition that never builds a + * terminal claim, so it announces settlement itself. */ + this.notifyGenerationSettled({ + streamId, + conversationId: expiredJob.conversationId, + userId: expiredJob.userId, + status: 'aborted', + }); + /** Terminal now; releasing ownership keeps the sweep's relay branch from + * announcing this generation a second time. */ + this.releaseJobOwnership(streamId, expiredJob.createdAt); return true; } @@ -9055,7 +9146,17 @@ class GenerationJobManagerClass { await this.runApprovalExpiredHandler(streamId, job); } await this.notifyApprovalExpiredRuntime(streamId, job.createdAt, runtime); - changed = this.releaseJobOwnership(streamId, job.createdAt) || changed; + const released = this.releaseJobOwnership(streamId, job.createdAt); + if (released) { + /** The store won the expiry CAS, so no local claim announced it. */ + this.notifyGenerationSettled({ + streamId, + conversationId: job.conversationId, + userId: job.userId, + status: 'aborted', + }); + } + changed = released || changed; continue; } if ( diff --git a/packages/api/src/stream/SteerRecovery.ts b/packages/api/src/stream/SteerRecovery.ts index 9f624dd88a3..675612525b0 100644 --- a/packages/api/src/stream/SteerRecovery.ts +++ b/packages/api/src/stream/SteerRecovery.ts @@ -1,5 +1,6 @@ +import { logger } from '@librechat/data-schemas'; import type { TFile, TPendingSteer } from 'librechat-data-provider'; -import { getReferencedQuotes } from '~/utils'; +import { getReferencedQuotes } from '../utils/quotes'; /** Immutable user-visible payload a parked steer recovery is allowed to submit. */ export interface RecoveredSteerPayload { @@ -14,16 +15,41 @@ export interface RecoveredSteerPayload { quotes: string[]; } -/** A recovery-shaped request did not reproduce the parked source exactly. */ +const recoveryFailureMessages = { + source_missing: 'Recovered steer source is no longer available', + protocol_mismatch: 'Recovered steer source requires protocol v2', + owner_mismatch: 'Recovered steer source belongs to a different owner', + invalid_payload: 'Recovered steer request payload is invalid', + payload_mismatch: 'Recovered steer payload does not match its parked source', +} as const; + +type RecoveryFailureReason = keyof typeof recoveryFailureMessages; + +/** Retain the wire code for old clients; the reason distinguishes permanent failures. */ export class RecoveredSteerPayloadMismatchError extends Error { readonly code = 'RECOVERY_PAYLOAD_MISMATCH'; + readonly reason: RecoveryFailureReason; - constructor() { - super('Recovered steer payload does not match its parked source'); + constructor(reason: RecoveryFailureReason = 'payload_mismatch') { + super(recoveryFailureMessages[reason]); + this.reason = reason; this.name = 'RecoveredSteerPayloadMismatchError'; } } +export function getSteerRecoveryFailure( + error: RecoveredSteerPayloadMismatchError, + context: { conversationId?: string; streamId: string; recoveredSteerId?: string }, +): { code: 'RECOVERY_PAYLOAD_MISMATCH'; reason: RecoveryFailureReason; error: string } { + const reason = error.reason ?? 'payload_mismatch'; + logger.warn('[SteerRecovery] Recovery rejected', { ...context, reason }); + return { + code: error.code, + reason, + error: 'The queued message could not be recovered. Review it before sending again.', + }; +} + /** Canonical attachment identity: every entry must name a file; ordering, * duplicates, and mutable display metadata do not affect identity. */ export function canonicalRecoveryFileIds(files: unknown): string[] | null { diff --git a/packages/api/src/stream/__tests__/generationSettled.spec.ts b/packages/api/src/stream/__tests__/generationSettled.spec.ts new file mode 100644 index 00000000000..473b548dced --- /dev/null +++ b/packages/api/src/stream/__tests__/generationSettled.spec.ts @@ -0,0 +1,145 @@ +/** + * Waiting completion deliveries are expedited when a generation settles, so the + * settled notification must fire exactly once for every terminal path and name + * the generation's owner — and a failing listener must never disturb cleanup. + */ +import type { GenerationSettledEvent } from '../GenerationJobManager'; +import { buildPendingAction, buildToolApprovalPayload } from '~/agents/hitl/policy'; + +/** Suppress winston Console transport output (survives jest.resetModules) */ +jest.spyOn(console, 'log').mockImplementation(); + +async function configureManager() { + const { GenerationJobManager } = await import('../GenerationJobManager'); + const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore'); + const { InMemoryEventTransport } = await import('../implementations/InMemoryEventTransport'); + + GenerationJobManager.configure({ + jobStore: new InMemoryJobStore(), + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + GenerationJobManager.initialize(); + return GenerationJobManager; +} + +describe('generation settled notifications', () => { + beforeEach(() => { + jest.resetModules(); + }); + + it('announces a completed generation once, with its owner', async () => { + const manager = await configureManager(); + const events: GenerationSettledEvent[] = []; + manager.onGenerationSettled((event) => events.push(event)); + const job = await manager.createJob('settled-complete', 'user-1', 'conversation-1'); + + await manager.completeJob('settled-complete', undefined, job.createdAt); + await manager.completeJob('settled-complete', undefined, job.createdAt); + + expect(events).toEqual([ + { + streamId: 'settled-complete', + conversationId: 'conversation-1', + userId: 'user-1', + status: 'complete', + }, + ]); + await manager.destroy(); + }); + + it('announces a failed generation', async () => { + const manager = await configureManager(); + const events: GenerationSettledEvent[] = []; + manager.onGenerationSettled((event) => events.push(event)); + const job = await manager.createJob('settled-error', 'user-2'); + + await manager.completeJob('settled-error', 'provider failed', job.createdAt); + + expect(events).toEqual([ + expect.objectContaining({ streamId: 'settled-error', userId: 'user-2', status: 'error' }), + ]); + await manager.destroy(); + }); + + it('announces an aborted generation', async () => { + const manager = await configureManager(); + const events: GenerationSettledEvent[] = []; + manager.onGenerationSettled((event) => events.push(event)); + await manager.createJob('settled-abort', 'user-3'); + + const result = await manager.abortJob('settled-abort'); + + expect(result.success).toBe(true); + expect(events).toEqual([ + expect.objectContaining({ streamId: 'settled-abort', userId: 'user-3', status: 'aborted' }), + ]); + await manager.destroy(); + }); + + it('announces a generation whose approval expired', async () => { + const manager = await configureManager(); + const events: GenerationSettledEvent[] = []; + manager.onGenerationSettled((event) => events.push(event)); + await manager.createJob('settled-approval', 'user-6', 'conversation-6'); + const action = buildPendingAction( + buildToolApprovalPayload([ + { name: 'shell', arguments: { command: 'ls' }, tool_call_id: 'call_expire' }, + ]), + { + streamId: 'settled-approval', + conversationId: 'conversation-6', + runId: 'run-1', + responseMessageId: 'msg-1', + }, + ); + await manager.approvals.pause('settled-approval', action); + expect(events).toEqual([]); + + await expect(manager.expireApproval('settled-approval')).resolves.toBe(true); + await expect(manager.expireApproval('settled-approval')).resolves.toBe(false); + /** The next expiry sweep relays the terminal job without announcing it again. */ + await (manager as unknown as { cleanup(): Promise }).cleanup(); + + expect(events).toEqual([ + { + streamId: 'settled-approval', + conversationId: 'conversation-6', + userId: 'user-6', + status: 'aborted', + }, + ]); + await manager.destroy(); + }); + + it('keeps cleanup and later listeners running when one listener throws', async () => { + const manager = await configureManager(); + const events: GenerationSettledEvent[] = []; + manager.onGenerationSettled(() => { + throw new Error('listener failed'); + }); + manager.onGenerationSettled((event) => events.push(event)); + const job = await manager.createJob('settled-throwing-listener', 'user-4'); + + await expect( + manager.completeJob('settled-throwing-listener', undefined, job.createdAt), + ).resolves.toBe(true); + + expect(events).toHaveLength(1); + await manager.destroy(); + }); + + it('stops notifying after unsubscribe', async () => { + const manager = await configureManager(); + const listener = jest.fn(); + const unsubscribe = manager.onGenerationSettled(listener); + unsubscribe(); + const job = await manager.createJob('settled-unsubscribed', 'user-5'); + + await manager.completeJob('settled-unsubscribed', undefined, job.createdAt); + + expect(listener).not.toHaveBeenCalled(); + await manager.destroy(); + }); +}); diff --git a/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts b/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts index 3b5d964d9df..1c8578d86d2 100644 --- a/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts +++ b/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts @@ -1070,48 +1070,86 @@ describe('GenerationJobManager start-generation claim', () => { expect(Object.getOwnPropertyDescriptor(durable!, 'replacedJobs')).toBeUndefined(); }); - it('terminalizes the exact committed epoch when legacy verification fails during recovery', async () => { - const streamId = 'stream-lost-create-legacy-failure'; - const clientRequestId = 'req-lost-create-legacy-failure'; - const claim = await manager.claimGeneration('user-1', clientRequestId, streamId, streamId, 2); - const actualCreate = store.createJob.bind(store); - jest.spyOn(store, 'createJob').mockImplementationOnce(async (...args) => { - await actualCreate(...args); - throw new Error('simulated lost atomic create reply'); - }); - const actualClaim = store.claimIdempotencyKey.bind(store); - let failLegacyProbe = true; - jest.spyOn(store, 'claimIdempotencyKey').mockImplementation((key, value, ttlSeconds) => { - if (failLegacyProbe && key === `{user-1:${clientRequestId}}`) { - failLegacyProbe = false; - return Promise.reject(new Error('simulated legacy probe outage')); - } - return actualClaim(key, value, ttlSeconds); - }); + it.each([ + { loseTerminalReply: false, replaced: false }, + { loseTerminalReply: true, replaced: false }, + { loseTerminalReply: true, replaced: true }, + ])( + 'announces only the exact recovered epoch (lost reply: $loseTerminalReply, replaced: $replaced)', + async ({ loseTerminalReply, replaced }) => { + const settled = jest.fn(); + manager.onGenerationSettled(settled); + const transition = store.transitionStatus.bind(store); + jest.spyOn(store, 'transitionStatus').mockImplementation(async (id, input) => { + const committed = await transition(id, input); + if (loseTerminalReply && input.to === 'error' && committed) { + if (replaced) { + const previous = await store.getJob(id); + const replacement = await actualCreate(id, 'user-1', id); + await transition(id, { + from: 'running', + to: 'error', + expectCreatedAt: replacement.createdAt, + patch: { error: previous?.error, finalEvent: previous?.finalEvent }, + }); + } + throw new Error('lost terminal CAS reply'); + } + return committed; + }); + const streamId = 'stream-lost-create-legacy-failure'; + const clientRequestId = 'req-lost-create-legacy-failure'; + const claim = await manager.claimGeneration('user-1', clientRequestId, streamId, streamId, 2); + const actualCreate = store.createJob.bind(store); + jest.spyOn(store, 'createJob').mockImplementationOnce(async (...args) => { + await actualCreate(...args); + throw new Error('simulated lost atomic create reply'); + }); + const actualClaim = store.claimIdempotencyKey.bind(store); + let failLegacyProbe = true; + jest.spyOn(store, 'claimIdempotencyKey').mockImplementation((key, value, ttlSeconds) => { + if (failLegacyProbe && key === `{user-1:${clientRequestId}}`) { + failLegacyProbe = false; + return Promise.reject(new Error('simulated legacy probe outage')); + } + return actualClaim(key, value, ttlSeconds); + }); - await expect( - manager.createJob(streamId, 'user-1', streamId, { - idempotencyClientRequestId: clientRequestId, - idempotencyClaimToken: claim.existing!.claimToken, - initialMetadata: { generationProtocolVersion: 2 }, - }), - ).rejects.toThrow('simulated lost atomic create reply'); + await expect( + manager.createJob(streamId, 'user-1', streamId, { + idempotencyClientRequestId: clientRequestId, + idempotencyClaimToken: claim.existing!.claimToken, + initialMetadata: { generationProtocolVersion: 2 }, + }), + ).rejects.toThrow('simulated lost atomic create reply'); - const durable = await store.getJob(streamId); - expect(durable).toMatchObject({ - status: 'error', - error: 'Generation idempotency rollout fence could not be recovered', - finalEvent: expect.stringContaining('terminal_payload_missing'), - }); - const retry = await manager.claimGeneration('user-1', clientRequestId, streamId, streamId, 2); - expect(retry).toMatchObject({ - claimed: false, - existing: { startedAt: durable!.createdAt }, - }); - await expect( - manager.takeoverGeneration('user-1', clientRequestId, streamId, retry.existing!), - ).resolves.toMatchObject({ claimed: false }); - }); + const durable = await store.getJob(streamId); + expect(durable).toMatchObject({ + status: 'error', + error: 'Generation idempotency rollout fence could not be recovered', + finalEvent: expect.stringContaining('terminal_payload_missing'), + }); + if (replaced) { + expect(settled).not.toHaveBeenCalled(); + return; + } + expect(settled).toHaveBeenCalledTimes(1); + expect(settled).toHaveBeenCalledWith({ + streamId, + conversationId: streamId, + userId: 'user-1', + status: 'error', + }); + const retry = await manager.claimGeneration('user-1', clientRequestId, streamId, streamId, 2); + expect(retry).toMatchObject({ + claimed: false, + existing: { startedAt: durable!.createdAt }, + }); + await expect( + manager.takeoverGeneration('user-1', clientRequestId, streamId, retry.existing!), + ).resolves.toMatchObject({ claimed: false }); + }, + ); it('terminalizes and preserves the primary fence when the legacy started mark fails', async () => { const claim = await manager.claimGeneration( diff --git a/packages/api/src/stream/__tests__/recovery.spec.ts b/packages/api/src/stream/__tests__/recovery.spec.ts new file mode 100644 index 00000000000..287909a9a61 --- /dev/null +++ b/packages/api/src/stream/__tests__/recovery.spec.ts @@ -0,0 +1,152 @@ +import { logger } from '@librechat/data-schemas'; +import type { RecoveredSteerPayload } from '../SteerRecovery'; +import { getSteerRecoveryFailure, RecoveredSteerPayloadMismatchError } from '../SteerRecovery'; +import { InMemoryJobStore, PARKED_STEERS_TTL_MS } from '../implementations/InMemoryJobStore'; + +const proof: RecoveredSteerPayload = { text: 'original words', fileIds: [], quotes: [] }; +const source = { steerId: 'source', text: proof.text, createdAt: 1 }; + +describe('recovery admission failures', () => { + let store: InMemoryJobStore; + beforeEach(() => { + store = new InMemoryJobStore(); + }); + afterEach(async () => { + jest.restoreAllMocks(); + await store.destroy(); + }); + const recover = (subject: InMemoryJobStore, payload = proof) => + subject.createJob( + 'conversation', + 'owner', + 'conversation', + undefined, + { generationProtocolVersion: 2 }, + 'source', + undefined, + undefined, + undefined, + payload, + ); + const park = async (subject: InMemoryJobStore, generationProtocolVersion: 1 | 2 = 2) => { + const job = await subject.createJob('conversation', 'owner', 'conversation', undefined, { + generationProtocolVersion, + }); + await subject.parkSteers( + 'conversation', + JSON.stringify({ userId: 'owner', steers: [source] }), + job.createdAt, + ); + return job; + }; + + test('missing source is not diagnosed as changed text', async () => { + await expect(recover(store)).rejects.toMatchObject({ + code: 'RECOVERY_PAYLOAD_MISMATCH', + reason: 'source_missing', + }); + expect(await store.getJob('conversation')).toBeNull(); + }); + + test('an expired source reports source_missing', async () => { + const now = Date.now(); + jest.spyOn(Date, 'now').mockReturnValue(now); + await park(store); + jest.spyOn(Date, 'now').mockReturnValue(now + PARKED_STEERS_TTL_MS + 1); + await expect(recover(store)).rejects.toMatchObject({ reason: 'source_missing' }); + }); + + test('already consumed source reports source_missing and preserves the winning job', async () => { + await park(store); + const job = await recover(store); + expect( + await store.consumeParkedSteer('conversation', 'source', 'owner', undefined, job.createdAt), + ).toBe(true); + await expect(recover(store)).rejects.toMatchObject({ reason: 'source_missing' }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + }); + + test('a legacy source reports protocol_mismatch without upgrading or deleting it', async () => { + const job = await park(store, 1); + await expect(recover(store)).rejects.toMatchObject({ reason: 'protocol_mismatch' }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + const parked = await store.claimParkedSteersDetailed('conversation', 'owner', undefined, 2); + expect(parked?.generationProtocolVersion).toBe(1); + }); + + test.each([ + { ...proof, text: 'different words' }, + { ...proof, quotes: ['local quote never accepted by the source'] }, + { ...proof, fileIds: ['extra-file'] }, + ])('a genuine mismatch preserves the source and predecessor', async (payload) => { + const job = await park(store); + await expect(recover(store, payload)).rejects.toMatchObject({ reason: 'payload_mismatch' }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + const parked = await store.claimParkedSteersDetailed('conversation', 'owner', undefined, 2); + expect(JSON.parse(parked!.payload).steers).toEqual([source]); + }); + + test('a foreign recovery cannot replace the owner or consume its source', async () => { + const job = await park(store); + await expect( + store.createJob( + 'conversation', + 'another-owner', + 'conversation', + undefined, + { generationProtocolVersion: 2 }, + 'source', + undefined, + undefined, + undefined, + proof, + ), + ).rejects.toMatchObject({ reason: 'owner_mismatch' }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + expect( + JSON.parse( + (await store.claimParkedSteersDetailed('conversation', 'owner', undefined, 2))!.payload, + ).steers, + ).toEqual([source]); + }); + + test('missing proof reports invalid_payload before a mutation', async () => { + await expect( + store.createJob( + 'conversation', + 'owner', + 'conversation', + undefined, + { generationProtocolVersion: 2 }, + 'source', + ), + ).rejects.toMatchObject({ reason: 'invalid_payload' }); + expect(await store.getJob('conversation')).toBeNull(); + }); + + test.each([ + 'source_missing', + 'protocol_mismatch', + 'owner_mismatch', + 'invalid_payload', + 'payload_mismatch', + ] as const)( + 'preserves the legacy wire code and records a payload-free diagnostic (%s)', + (reason) => { + const warning = jest.spyOn(logger, 'warn').mockImplementation(); + expect( + getSteerRecoveryFailure(new RecoveredSteerPayloadMismatchError(reason), { + conversationId: 'conversation', + streamId: 'stream', + recoveredSteerId: 'source', + }), + ).toMatchObject({ code: 'RECOVERY_PAYLOAD_MISMATCH', reason }); + expect(warning).toHaveBeenCalledWith('[SteerRecovery] Recovery rejected', { + conversationId: 'conversation', + streamId: 'stream', + recoveredSteerId: 'source', + reason, + }); + }, + ); +}); diff --git a/packages/api/src/stream/__tests__/recovery.stream_integration.spec.ts b/packages/api/src/stream/__tests__/recovery.stream_integration.spec.ts new file mode 100644 index 00000000000..47d46d30fc3 --- /dev/null +++ b/packages/api/src/stream/__tests__/recovery.stream_integration.spec.ts @@ -0,0 +1,133 @@ +import { randomUUID } from 'crypto'; +import type { RecoveredSteerPayload } from '../SteerRecovery'; +import type { RedisTestClient } from './helpers/redis'; +import { clearRedisTestPrefix, createRedisTestClient } from './helpers/redis'; +import { RedisJobStore } from '../implementations/RedisJobStore'; + +const proof: RecoveredSteerPayload = { text: 'original words', fileIds: [], quotes: [] }; +const source = { steerId: 'source', text: proof.text, createdAt: 1 }; + +describe('Redis recovery rejection reasons', () => { + const prefix = `recovery-test:${randomUUID()}:`; + let redis: RedisTestClient; + let store: RedisJobStore; + + beforeAll(async () => { + redis = createRedisTestClient(prefix); + await redis.ping(); + }); + beforeEach(() => { + store = new RedisJobStore(redis); + }); + afterEach(async () => { + await store.destroy(); + await clearRedisTestPrefix(redis, prefix); + }); + afterAll(async () => { + await redis.quit(); + }); + + const recover = (payload = proof, generationProtocolVersion: 1 | 2 = 2, userId = 'owner') => + store.createJob( + 'conversation', + userId, + 'conversation', + undefined, + { generationProtocolVersion }, + 'source', + undefined, + undefined, + undefined, + payload, + ); + + const park = async (generationProtocolVersion: 1 | 2 = 2) => { + const job = await store.createJob('conversation', 'owner', 'conversation', undefined, { + generationProtocolVersion, + }); + await store.parkSteers( + 'conversation', + JSON.stringify({ userId: 'owner', steers: [source] }), + job.createdAt, + ); + return job; + }; + + test('missing source is a specific failure with no job creation', async () => { + await expect(recover()).rejects.toMatchObject({ + code: 'RECOVERY_PAYLOAD_MISMATCH', + reason: 'source_missing', + }); + expect(await store.getJob('conversation')).toBeNull(); + }); + + test('expired source does not become a genuine payload mismatch', async () => { + await park(); + await redis.pexpire('stream:{conversation}:parked', 1); + await new Promise((resolve) => setTimeout(resolve, 20)); + await expect(recover()).rejects.toMatchObject({ reason: 'source_missing' }); + }); + + test('consumed source cannot start a second generation', async () => { + await park(); + const job = await recover(); + expect( + await store.consumeParkedSteer('conversation', 'source', 'owner', undefined, job.createdAt), + ).toBe(true); + await expect(recover()).rejects.toMatchObject({ reason: 'source_missing' }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + }); + + test.each([ + [1, 2], + [2, 1], + ] as const)( + 'source v%s with requester v%s reports protocol_mismatch', + async (sourceVersion, requestVersion) => { + const job = await park(sourceVersion); + await expect(recover(proof, requestVersion)).rejects.toMatchObject({ + reason: 'protocol_mismatch', + }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + expect(await redis.exists('stream:{conversation}:parked')).toBe(1); + }, + ); + + test('foreign recovery owner cannot mutate the source', async () => { + const job = await park(); + await expect(recover(proof, 2, 'another-owner')).rejects.toMatchObject({ + reason: 'owner_mismatch', + }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + expect(await redis.exists('stream:{conversation}:parked')).toBe(1); + }); + + test.each([ + { ...proof, text: 'changed words' }, + { ...proof, quotes: ['local excerpt never accepted'] }, + { ...proof, fileIds: ['another-file'] }, + ])('actual payload mismatch leaves the source recoverable', async (payload) => { + const job = await park(); + await expect(recover(payload)).rejects.toMatchObject({ reason: 'payload_mismatch' }); + expect((await store.getJob('conversation'))?.createdAt).toBe(job.createdAt); + expect( + JSON.parse( + (await store.claimParkedSteersDetailed('conversation', 'owner', undefined, 2))!.payload, + ).steers, + ).toEqual([source]); + }); + + test('invalid proof reports invalid_payload before changing state', async () => { + await expect( + store.createJob( + 'conversation', + 'owner', + 'conversation', + undefined, + { generationProtocolVersion: 2 }, + 'source', + ), + ).rejects.toMatchObject({ reason: 'invalid_payload' }); + expect(await store.getJob('conversation')).toBeNull(); + }); +}); diff --git a/packages/api/src/stream/__tests__/settled.spec.ts b/packages/api/src/stream/__tests__/settled.spec.ts new file mode 100644 index 00000000000..ff760d45ddd --- /dev/null +++ b/packages/api/src/stream/__tests__/settled.spec.ts @@ -0,0 +1,276 @@ +import type { GenerationSettledEvent, GenerationSettledListener } from '../GenerationJobManager'; +import type { GenerationJobStatus } from '../../types/stream'; +import { waitForGenerationSettled } from '../settled'; + +function createSource(initial: Record) { + const statuses = new Map(Object.entries(initial)); + const listeners = new Set(); + const epochs = new Map(); + const pending = new Set(); + return { + statuses, + epochs, + pending, + listeners, + getGenerationSettlementState: jest.fn(async (streamId: string) => { + const status = statuses.get(streamId); + return status == null + ? undefined + : { + createdAt: epochs.get(streamId) ?? 1, + status, + terminalPersistencePending: pending.has(streamId), + }; + }), + onGenerationSettled: (listener: GenerationSettledListener) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + settle(event: Omit) { + statuses.set(event.conversationId, event.status as GenerationJobStatus); + for (const listener of listeners) { + listener({ ...event, streamId: event.conversationId, userId: 'user-1' }); + } + }, + }; +} + +describe('waitForGenerationSettled', () => { + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it.each([ + ['no generation', {}], + ['a completed generation', { 'conversation-1': 'complete' as const }], + ])('resolves false at once for %s', async (_label, initial) => { + const source = createSource(initial); + + await expect(waitForGenerationSettled(source, 'conversation-1')).resolves.toBe(false); + expect(source.listeners.size).toBe(0); + }); + + it('resolves when this conversation settles, ignoring others', async () => { + const source = createSource({ 'conversation-1': 'running', 'conversation-2': 'running' }); + let resolved: boolean | undefined; + void waitForGenerationSettled(source, 'conversation-1').then((settled) => { + resolved = settled; + }); + await jest.advanceTimersByTimeAsync(0); + + source.settle({ conversationId: 'conversation-2', status: 'complete' }); + await jest.advanceTimersByTimeAsync(0); + expect(resolved).toBeUndefined(); + + source.settle({ conversationId: 'conversation-1', status: 'complete' }); + await jest.advanceTimersByTimeAsync(0); + expect(resolved).toBe(true); + expect(source.listeners.size).toBe(0); + }); + + it('waits through an approval pause', async () => { + const source = createSource({ 'conversation-1': 'requires_action' }); + let resolved: boolean | undefined; + void waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }).then( + (settled) => { + resolved = settled; + }, + ); + + await jest.advanceTimersByTimeAsync(5_000); + expect(resolved).toBeUndefined(); + }); + + it('observes a generation that settles on another replica', async () => { + const source = createSource({ 'conversation-1': 'running' }); + const waiting = waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }); + await jest.advanceTimersByTimeAsync(0); + + source.statuses.set('conversation-1', 'complete'); + await jest.advanceTimersByTimeAsync(1_000); + + await expect(waiting).resolves.toBe(true); + expect(source.listeners.size).toBe(0); + }); + + it('keeps waiting for as long as the generation stays active by default', async () => { + const source = createSource({ 'conversation-1': 'running' }); + let resolved: boolean | undefined; + void waitForGenerationSettled(source, 'conversation-1', { recheckMs: 60 * 60 * 1_000 }).then( + (settled) => { + resolved = settled; + }, + ); + + await jest.advanceTimersByTimeAsync(48 * 60 * 60 * 1_000); + expect(resolved).toBeUndefined(); + + source.settle({ conversationId: 'conversation-1', status: 'complete' }); + await jest.advanceTimersByTimeAsync(0); + expect(resolved).toBe(true); + }); + + it('gives up after the maximum wait', async () => { + const source = createSource({ 'conversation-1': 'running' }); + const waiting = waitForGenerationSettled(source, 'conversation-1', { + recheckMs: 1_000, + maxWaitMs: 10_000, + }); + + await jest.advanceTimersByTimeAsync(10_000); + + await expect(waiting).resolves.toBe(false); + expect(source.listeners.size).toBe(0); + }); + + it('stops waiting when the caller aborts', async () => { + const source = createSource({ 'conversation-1': 'running' }); + const controller = new AbortController(); + const waiting = waitForGenerationSettled(source, 'conversation-1', { + signal: controller.signal, + }); + await jest.advanceTimersByTimeAsync(0); + + controller.abort(); + + await expect(waiting).resolves.toBe(false); + expect(source.listeners.size).toBe(0); + }); + + it('recovers from initial and later read failures without unsubscribing', async () => { + const source = createSource({ 'conversation-1': 'running' }); + source.getGenerationSettlementState.mockRejectedValueOnce(new Error('job store unavailable')); + const waiting = waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }); + await jest.advanceTimersByTimeAsync(1_000); + expect(source.getGenerationSettlementState).toHaveBeenCalledTimes(2); + expect(source.listeners.size).toBe(1); + source.getGenerationSettlementState.mockRejectedValueOnce(new Error('later outage')); + await jest.advanceTimersByTimeAsync(1_000); + expect(source.listeners.size).toBe(1); + source.statuses.set('conversation-1', 'complete'); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toBe(true); + expect(source.listeners.size).toBe(0); + }); + + it.each(['running', 'requires_action'] as const)( + 'does not follow a %s successor on another replica', + async (status) => { + const source = createSource({ 'conversation-1': 'running' }); + const waiting = waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }); + await jest.advanceTimersByTimeAsync(0); + source.statuses.set('conversation-1', status); + source.epochs.set('conversation-1', 2); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toBe(true); + expect(source.listeners.size).toBe(0); + }, + ); + + it('uses the dispatch epoch even when the first successful read sees its successor', async () => { + const source = createSource({ 'conversation-1': 'requires_action' }); + source.epochs.set('conversation-1', 2); + source.getGenerationSettlementState.mockRejectedValueOnce(new Error('outage')); + const waiting = waitForGenerationSettled(source, 'conversation-1', { + generationCreatedAt: 1, + recheckMs: 1_000, + }); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toBe(true); + }); + + it.each(['complete', 'error', 'aborted'] as const)( + 'waits for %s terminal persistence even after a stale settlement event', + async (status) => { + const source = createSource({ 'conversation-1': status }); + source.pending.add('conversation-1'); + let resolved = false; + const waiting = waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }).then( + (value) => { + resolved = true; + return value; + }, + ); + await jest.advanceTimersByTimeAsync(120_000); + source.settle({ conversationId: 'conversation-1', status }); + await jest.advanceTimersByTimeAsync(0); + expect(resolved).toBe(false); + source.pending.delete('conversation-1'); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toBe(true); + }, + ); + + it('finishes when the observed job disappears', async () => { + const source = createSource({ 'conversation-1': 'running' }); + const waiting = waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }); + await jest.advanceTimersByTimeAsync(0); + source.statuses.delete('conversation-1'); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toBe(true); + }); + + it('resolves a pre-aborted wait without subscribing or creating timers', async () => { + const source = createSource({ 'conversation-1': 'running' }); + const controller = new AbortController(); + controller.abort(); + await expect( + waitForGenerationSettled(source, 'conversation-1', { signal: controller.signal }), + ).resolves.toBe(false); + expect(source.getGenerationSettlementState).not.toHaveBeenCalled(); + expect(source.listeners.size).toBe(0); + expect(jest.getTimerCount()).toBe(0); + }); + + it('serializes reads and remembers a settlement event during an in-flight read', async () => { + const source = createSource({ 'conversation-1': 'running' }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + source.getGenerationSettlementState.mockImplementationOnce(async () => { + await gate; + return { createdAt: 1, status: 'running', terminalPersistencePending: false }; + }); + const waiting = waitForGenerationSettled(source, 'conversation-1', { recheckMs: 1_000 }); + await jest.advanceTimersByTimeAsync(10_000); + source.settle({ conversationId: 'conversation-1', status: 'complete' }); + source.settle({ conversationId: 'conversation-1', status: 'complete' }); + expect(source.getGenerationSettlementState).toHaveBeenCalledTimes(1); + release(); + await expect(waiting).resolves.toBe(true); + expect(source.getGenerationSettlementState).toHaveBeenCalledTimes(2); + expect(jest.getTimerCount()).toBe(0); + }); + + it('removes abort listeners and ignores reads completing after cancellation', async () => { + const source = createSource({ 'conversation-1': 'running' }); + const controller = new AbortController(); + const remove = jest.spyOn(controller.signal, 'removeEventListener'); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + source.getGenerationSettlementState.mockImplementationOnce(async () => { + await gate; + return undefined; + }); + const waiting = waitForGenerationSettled(source, 'conversation-1', { + signal: controller.signal, + maxWaitMs: 500, + }); + controller.abort(); + await expect(waiting).resolves.toBe(false); + release(); + await jest.advanceTimersByTimeAsync(1_000); + expect(remove).toHaveBeenCalledWith('abort', expect.any(Function)); + expect(source.listeners.size).toBe(0); + expect(jest.getTimerCount()).toBe(0); + }); +}); diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index f711884b09b..14c290c5979 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -320,7 +320,9 @@ export class InMemoryJobStore implements IJobStoreV2 { existingJob.userId !== userId || (existingJob.tenantId != null && existingJob.tenantId !== tenantId)) ) { - throw new Error('Generation job owner mismatch'); + throw recoveredSteerId != null + ? new RecoveredSteerPayloadMismatchError('owner_mismatch') + : new Error('Generation job owner mismatch'); } const parked = this.parkedSteers.get(streamId); @@ -351,7 +353,9 @@ export class InMemoryJobStore implements IJobStoreV2 { throw new Error('Generation recovery state is corrupt'); } if (parsed.userId !== userId || (parsed.tenantId != null && parsed.tenantId !== tenantId)) { - throw new Error('Generation job owner mismatch'); + throw recoveredSteerId != null + ? new RecoveredSteerPayloadMismatchError('owner_mismatch') + : new Error('Generation job owner mismatch'); } }; @@ -398,15 +402,15 @@ export class InMemoryJobStore implements IJobStoreV2 { const assertRecoveryCompatible = (): void => { if (recoveredSteerId == null) { if (recoveredSteerPayload != null) { - throw new RecoveredSteerPayloadMismatchError(); + throw new RecoveredSteerPayloadMismatchError('invalid_payload'); } return; } if (!isRecoveredSteerPayload(recoveredSteerPayload)) { - throw new RecoveredSteerPayloadMismatchError(); + throw new RecoveredSteerPayloadMismatchError('invalid_payload'); } if (initialMetadata.generationProtocolVersion === 1) { - throw new RecoveredSteerPayloadMismatchError(); + throw new RecoveredSteerPayloadMismatchError('protocol_mismatch'); } let candidate: ReturnType | undefined; @@ -437,11 +441,13 @@ export class InMemoryJobStore implements IJobStoreV2 { candidate = item == null ? undefined : toPendingSteer(item); } - if ( - sourceProtocol !== 2 || - candidate == null || - !recoveredSteerPayloadMatches(candidate, recoveredSteerPayload) - ) { + if (candidate == null) { + throw new RecoveredSteerPayloadMismatchError('source_missing'); + } + if (sourceProtocol !== 2) { + throw new RecoveredSteerPayloadMismatchError('protocol_mismatch'); + } + if (!recoveredSteerPayloadMatches(candidate, recoveredSteerPayload)) { throw new RecoveredSteerPayloadMismatchError(); } }; diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index 20675b3ec18..36848f35225 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -473,19 +473,19 @@ const JOB_CREATE_LUA = 'if replacedProviderDrained then replaced.providerDrained = replacedProviderDrained == "1" end ' + 'replacementChain[#replacementChain + 1] = replaced replacementSeen[tostring(replacedEpoch)] = true end ' + 'local recoveredSteerId = ARGV[5] local expectedRecovery = nil ' + - 'if recoveredSteerId ~= "" and ARGV[10] ~= "2" then return { "", "", "0", "recovery_payload_mismatch" } end ' + + 'if recoveredSteerId ~= "" and ARGV[10] ~= "2" then return { "", "", "0", "recovery_protocol_mismatch" } end ' + 'if recoveredSteerId ~= "" then local ok, decoded = pcall(cjson.decode, ARGV[9]) ' + 'if not ok or type(decoded) ~= "table" or type(decoded.text) ~= "string" ' + - 'or not isDenseArray(decoded.fileIds) then return { "", "", "0", "recovery_payload_mismatch" } end ' + + 'or not isDenseArray(decoded.fileIds) then return { "", "", "0", "recovery_invalid_payload" } end ' + 'local expectedSeen = {} for i = 1, #decoded.fileIds do local fileId = decoded.fileIds[i] ' + 'if type(fileId) ~= "string" or fileId == "" or expectedSeen[fileId] then ' + - 'return { "", "", "0", "recovery_payload_mismatch" } end expectedSeen[fileId] = true end ' + + 'return { "", "", "0", "recovery_invalid_payload" } end expectedSeen[fileId] = true end ' + 'if decoded.quotes ~= nil then if not isDenseArray(decoded.quotes) then ' + - 'return { "", "", "0", "recovery_payload_mismatch" } end ' + + 'return { "", "", "0", "recovery_invalid_payload" } end ' + 'for i = 1, #decoded.quotes do if type(decoded.quotes[i]) ~= "string" or decoded.quotes[i] == "" then ' + - 'return { "", "", "0", "recovery_payload_mismatch" } end end end ' + + 'return { "", "", "0", "recovery_invalid_payload" } end end end ' + 'expectedRecovery = decoded elseif ARGV[9] ~= "" then ' + - 'return { "", "", "0", "recovery_payload_mismatch" } end ' + + 'return { "", "", "0", "recovery_invalid_payload" } end ' + 'local function recoveryMatches(item, expected) ' + 'if not expected or type(item.text) ~= "string" or item.text ~= expected.text then return false end ' + // Quotes are model-bound like the text: order-significant identity, with a @@ -518,7 +518,7 @@ const JOB_CREATE_LUA = 'return { "", "", "0", "recovery_corrupt" } end end ' + 'if recoveredSteerId ~= "" and parked.generationProtocolVersion ~= 2 then ' + 'for i = 1, #parked.steers do if parked.steers[i].steerId == recoveredSteerId then ' + - 'return { "", "", "0", "recovery_payload_mismatch" } end end end ' + + 'return { "", "", "0", "recovery_protocol_mismatch" } end end end ' + 'if parked.userId ~= ARGV[6] or (parked.tenantId and parked.tenantId ~= ARGV[7]) then ' + 'return { "", "", "0", "owner_mismatch" } end ' + 'parkedUserId = parked.userId parkedTenantId = parked.tenantId ' + @@ -531,7 +531,7 @@ const JOB_CREATE_LUA = 'local sources = { claimedRows, redis.call("LRANGE", KEYS[4], 0, -1) } ' + 'for s = 1, #sources do for i = 1, #sources[s] do local ok, item = pcall(cjson.decode, sources[s][i]) ' + 'if ok and recoveredSteerId ~= "" and item.steerId == recoveredSteerId and replacedProtocol ~= "2" then ' + - 'return { "", "", "0", "recovery_payload_mismatch" } end ' + + 'return { "", "", "0", "recovery_protocol_mismatch" } end ' + 'if ok and item.steerId and not seen[item.steerId] then seen[item.steerId] = true ' + 'local projected = { steerId = item.steerId, text = item.text, createdAt = item.createdAt } ' + 'if item.clientSteerId then projected.clientSteerId = item.clientSteerId end ' + @@ -545,10 +545,11 @@ const JOB_CREATE_LUA = 'for i = 1, #merged do local item = merged[i] ' + 'item.recoveringCreatedAt = nil ' + 'if recoveredSteerId ~= "" and item.steerId == recoveredSteerId then ' + - 'if not recoveryOwnerMatches or not recoveryMatches(item, expectedRecovery) then ' + + 'if not recoveryOwnerMatches then return { "", "", "0", "recovery_owner_mismatch" } end ' + + 'if not recoveryMatches(item, expectedRecovery) then ' + 'return { "", "", "0", "recovery_payload_mismatch" } end ' + 'item.recoveringCreatedAt = createdAt recoveryFound = true end end ' + - 'if not recoveryFound then return { "", "", "0", "recovery_payload_mismatch" } end ' + + 'if not recoveryFound then return { "", "", "0", "recovery_source_missing" } end ' + 'for i = 1, #receiptUpdates do local item = receiptUpdates[i] ' + 'if replacedProtocol == "2" and item.clientSteerId then local raw = redis.call("HGET", KEYS[8], item.clientSteerId) ' + 'if raw then local receiptOk, receipt = pcall(cjson.decode, raw) ' + @@ -2101,13 +2102,26 @@ export class RedisJobStore implements IJobStoreV2 { throw new Error('Generation idempotency claim was taken over before job creation'); } if (Array.isArray(previousOwner) && previousOwner[3] === 'owner_mismatch') { - throw new Error('Generation job owner mismatch'); + throw recoveredSteerId != null + ? new RecoveredSteerPayloadMismatchError('owner_mismatch') + : new Error('Generation job owner mismatch'); } if (Array.isArray(previousOwner) && previousOwner[3] === 'recovery_corrupt') { throw new Error('Generation recovery state is corrupt'); } - if (Array.isArray(previousOwner) && previousOwner[3] === 'recovery_payload_mismatch') { - throw new RecoveredSteerPayloadMismatchError(); + if (Array.isArray(previousOwner)) { + switch (previousOwner[3]) { + case 'recovery_source_missing': + throw new RecoveredSteerPayloadMismatchError('source_missing'); + case 'recovery_protocol_mismatch': + throw new RecoveredSteerPayloadMismatchError('protocol_mismatch'); + case 'recovery_owner_mismatch': + throw new RecoveredSteerPayloadMismatchError('owner_mismatch'); + case 'recovery_invalid_payload': + throw new RecoveredSteerPayloadMismatchError('invalid_payload'); + case 'recovery_payload_mismatch': + throw new RecoveredSteerPayloadMismatchError(); + } } if (Array.isArray(previousOwner) && previousOwner[3] === 'replacement_receipt_corrupt') { throw new Error('Generation replacement receipt is corrupt'); diff --git a/packages/api/src/stream/index.ts b/packages/api/src/stream/index.ts index b6c97e60988..504fafa586a 100644 --- a/packages/api/src/stream/index.ts +++ b/packages/api/src/stream/index.ts @@ -4,6 +4,8 @@ export { type CreateGenerationJobOptions, type GenerationJobManagerOptions, type TerminalJobClaim, + type GenerationSettledEvent, + type GenerationSettledListener, GENERATION_RECOVERY_FAILED_ERROR, TERMINAL_PUBLICATION_RECONNECT_ERROR, } from './GenerationJobManager'; @@ -56,6 +58,7 @@ export type { JobStoreV2RequiredMethod } from './jobStoreCapabilities'; export { buildRecoveredSteerPayload, canonicalRecoveryFileIds, + getSteerRecoveryFailure, RecoveredSteerPayloadMismatchError, } from './SteerRecovery'; export type { RecoveredSteerPayload } from './SteerRecovery'; @@ -83,3 +86,7 @@ export { RedisJobStore } from './implementations/RedisJobStore'; export { RedisEventTransport } from './implementations/RedisEventTransport'; export { waitForGenerationPersistence } from './persistence'; +export { waitForGenerationSettled } from './settled'; +export type { GenerationSettledSource, GenerationSettledWaitOptions } from './settled'; + +export type { GenerationSettlementState } from './interfaces/IJobStore'; diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 4507661aeb9..2e7d471b360 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -103,9 +103,13 @@ export type JobStatus = 'running' | 'complete' | 'error' | 'aborted' | 'requires * Missing markers on pre-rollout records are interpreted as protocol v1. */ export type GenerationProtocolVersion = 1 | 2; -/** - * Serializable job data - no object references, suitable for Redis/external storage - */ +/** Read-only durable evidence, without attaching a runtime or recovering a slow save. */ +export type GenerationSettlementState = Pick< + SerializableJobData, + 'createdAt' | 'status' | 'terminalPersistencePending' +>; + +/** Serializable job data without object references, suitable for external storage. */ export interface SerializableJobData { streamId: string; userId: string; diff --git a/packages/api/src/stream/settled.ts b/packages/api/src/stream/settled.ts new file mode 100644 index 00000000000..38fc7a092cd --- /dev/null +++ b/packages/api/src/stream/settled.ts @@ -0,0 +1,114 @@ +import { logger } from '@librechat/data-schemas'; +import type { GenerationSettledListener } from './GenerationJobManager'; +import type { GenerationSettlementState } from './interfaces/IJobStore'; + +const DEFAULT_RECHECK_MS = 60_000; + +export interface GenerationSettledSource { + getGenerationSettlementState: ( + streamId: string, + ) => Promise; + onGenerationSettled: (listener: GenerationSettledListener) => () => void; +} + +export interface GenerationSettledWaitOptions { + recheckMs?: number; + /** Pin the dispatch generation even if it was replaced before the first read. */ + generationCreatedAt?: number; + /** Optional caller deadline. Store errors are not evidence of settlement. */ + maxWaitMs?: number; + /** Stops waiting (resolving `false`) once the caller no longer needs the answer. */ + signal?: AbortSignal; +} + +/** Waits for the dispatch generation, including terminal persistence. Events request + * a fresh read; they cannot settle an unrelated epoch or bypass an unfinished save. + * The first successful read pins the epoch if the caller did not supply it. Reads + * are single-flight, with a bounded periodic fallback for remote events and outages. */ +export function waitForGenerationSettled( + source: GenerationSettledSource, + conversationId: string, + { + recheckMs = DEFAULT_RECHECK_MS, + generationCreatedAt, + maxWaitMs, + signal, + }: GenerationSettledWaitOptions = {}, +): Promise { + if (signal?.aborted) { + return Promise.resolve(false); + } + return new Promise((resolve) => { + let done = false; + let reading = false; + let readRequested = false; + let epoch = generationCreatedAt; + let observedActive = epoch != null; + let unsubscribe: (() => void) | undefined = undefined; + let recheck: NodeJS.Timeout | undefined; + let deadline: NodeJS.Timeout | undefined; + const abort = (): void => finish(false); + const finish = (settled: boolean): void => { + if (done) return; + done = true; + unsubscribe?.(); + signal?.removeEventListener('abort', abort); + clearTimeout(recheck); + clearTimeout(deadline); + resolve(settled); + }; + const readState = async (): Promise => { + if (done) return; + if (reading) { + readRequested = true; + return; + } + reading = true; + clearTimeout(recheck); + try { + const job = await source.getGenerationSettlementState(conversationId); + if (done) return; + if (job == null || (epoch != null && epoch !== job.createdAt)) { + finish(observedActive); + return; + } + epoch = job.createdAt; + const active = + job.status === 'running' || + job.status === 'requires_action' || + job.terminalPersistencePending === true; + if (!active) { + finish(observedActive); + return; + } + observedActive = true; + } catch (error) { + logger.warn( + `[GenerationSettled] Failed to read generation state for ${conversationId}:`, + error, + ); + } finally { + reading = false; + if (!done) { + if (readRequested) { + readRequested = false; + void readState(); + } else { + recheck = setTimeout(() => void readState(), recheckMs); + recheck.unref?.(); + } + } + } + }; + if (maxWaitMs != null) { + deadline = setTimeout(() => finish(false), maxWaitMs); + deadline.unref?.(); + } + signal?.addEventListener('abort', abort, { once: true }); + unsubscribe = source.onGenerationSettled((event) => { + if (event.conversationId === conversationId) void readState(); + }); + if (done) unsubscribe(); + else void readState(); + }); +} diff --git a/packages/api/src/utils/tokenizer.spec.ts b/packages/api/src/utils/tokenizer.spec.ts index bda83078ad2..feaddeb64a5 100644 --- a/packages/api/src/utils/tokenizer.spec.ts +++ b/packages/api/src/utils/tokenizer.spec.ts @@ -1,5 +1,6 @@ import { Tokenizer as AiTokenizer } from 'ai-tokenizer'; -import Tokenizer from './tokenizer'; +import type { EncodingName } from './tokenizer'; +import Tokenizer, { countTokens } from './tokenizer'; jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -32,6 +33,16 @@ describe('Tokenizer', () => { expect(count).toBeGreaterThan(0); }); + it('keeps the oversized cold-start fallback and loads the encoding for the next count', async () => { + await jest.isolateModulesAsync(async () => { + const cold = (await import('./tokenizer')).default; + const text = 'word '.repeat(4096); + expect(cold.getTokenCount(text, 'o200k_base')).toBe(Buffer.byteLength(text, 'utf8')); + await cold.initEncoding('o200k_base'); + expect(cold.getTokenCount(text, 'o200k_base')).toBeLessThan(text.length / 4); + }); + }); + it('should deduplicate concurrent init calls', async () => { const [, , count] = await Promise.all([ Tokenizer.initEncoding('o200k_base'), @@ -60,18 +71,29 @@ describe('Tokenizer', () => { expect(count).toBeGreaterThan(0); }); - it.each([ - { label: 'an uninterrupted non-whitespace run', text: '_'.repeat(4 * 1024 + 1) }, - { label: 'an oversized input', text: 'word '.repeat(1024) }, - { label: 'multibyte input', text: '界'.repeat(4 * 1024 + 1) }, - ])('uses a conservative estimate without tokenizing $label', ({ text }) => { - const count = jest.spyOn(AiTokenizer.prototype, 'count'); + it('uses the same bounded count through the asynchronous public helper', async () => { + const text = 'word '.repeat(4096); + expect(await countTokens(text)).toBe(Tokenizer.getTokenCount(text, 'o200k_base')); + expect(await countTokens(text)).toBeLessThan(text.length / 4); + }); + + it('uses the established fallback for the whole input after a later chunk fails', async () => { + const text = 'word '.repeat(4096); + const original = AiTokenizer.prototype.count; + const count = jest + .spyOn(AiTokenizer.prototype, 'count') + .mockImplementationOnce(original) + .mockImplementationOnce(() => { + throw new Error('second chunk failed'); + }); try { expect(Tokenizer.getTokenCount(text, 'o200k_base')).toBe(Buffer.byteLength(text, 'utf8')); - expect(count).not.toHaveBeenCalled(); + expect(count).toHaveBeenCalledTimes(2); } finally { count.mockRestore(); + await Tokenizer.initEncoding('o200k_base'); } + expect(Tokenizer.getTokenCount(text, 'o200k_base')).toBeLessThan(text.length / 4); }); }); @@ -89,17 +111,103 @@ describe('Tokenizer', () => { } }); - it('uses a conservative estimate for unsafe input without invoking the tokenizer', async () => { + it('never caches a partial count after a later chunk fails and recovers on retry', async () => { const counter = await Tokenizer.createExactTokenCounter('o200k_base'); + const text = 'word '.repeat(4096); + const original = AiTokenizer.prototype.count; + const count = jest + .spyOn(AiTokenizer.prototype, 'count') + .mockImplementationOnce(original) + .mockImplementationOnce(() => { + throw new Error('second chunk failed'); + }); + try { + expect(() => counter(text)).toThrow('second chunk failed'); + expect(count).toHaveBeenCalledTimes(2); + } finally { + count.mockRestore(); + await Tokenizer.initEncoding('o200k_base'); + } + const reloaded = await Tokenizer.createExactTokenCounter('o200k_base'); + expect(counter(text)).toBe(reloaded(text)); + }); + }); + + const encodings: EncodingName[] = ['o200k_base', 'claude']; + describe.each(encodings)('bounded counting with %s', (encoding) => { + let counter: (text: string) => number; + beforeAll(async () => { + counter = await Tokenizer.createExactTokenCounter(encoding); + }); + + it.each([4095, 4096, 4097, 8192, 8193, 65536])( + 'counts %i characters without a byte-count cliff', + (length) => { + const text = 'word '.repeat(Math.ceil(length / 5)).slice(0, length); + const exact = Tokenizer.countExactTokens(text, encoding)!; + expect(counter(text)).toBeGreaterThanOrEqual(exact * 0.99); + expect(counter(text)).toBeLessThanOrEqual(exact * 1.01); + expect(Tokenizer.getTokenCount(text, encoding)).toBe(counter(text)); + }, + ); + + it.each([ + { + label: 'minified JSON', + text: JSON.stringify( + Array.from({ length: 200 }, (_, id) => ({ + id, + description: 'Search project documents', + parameters: { type: 'object', required: ['query'] }, + })), + ), + }, + { label: 'CJK', text: '界'.repeat(4097) }, + { label: 'emoji', text: '😀'.repeat(4097) }, + { label: 'mixed density', text: 'word '.repeat(1024) + '界 '.repeat(2048) }, + ])('measures all of $label rather than extrapolating a prefix', ({ text }) => { + const exact = Tokenizer.countExactTokens(text, encoding)!; + expect(counter(text)).toBeGreaterThanOrEqual(exact * 0.98); + expect(counter(text)).toBeLessThanOrEqual(exact * 1.02); + expect(Tokenizer.getTokenCount(text, encoding)).toBe(counter(text)); + }); + + it.each([ + { label: 'English', text: 'word '.repeat(4096) }, + { label: 'uninterrupted punctuation', text: '_'.repeat(8193) }, + { label: 'uninterrupted letters', text: 'a'.repeat(8193) }, + { label: 'whitespace', text: ' '.repeat(8193) }, + { label: 'Unicode whitespace', text: '\u2003'.repeat(8193) }, + { label: 'CJK', text: '界'.repeat(8193) }, + { + label: 'surrogate at input boundary', + text: 'word '.repeat(819) + '😀' + 'word '.repeat(819), + }, + { label: 'surrogate at run boundary', text: 'a'.repeat(255) + '😀' + 'b'.repeat(500) }, + ])('bounds work without dropping or corrupting $label', ({ text }) => { const count = jest.spyOn(AiTokenizer.prototype, 'count'); try { - const text = '_'.repeat(4 * 1024 + 1); - expect(counter(text)).toBe(Buffer.byteLength(text, 'utf8')); - expect(count).not.toHaveBeenCalled(); + counter(text); + const chunks = count.mock.calls.map(([chunk]) => chunk); + expect(chunks.join('')).toBe(text); + expect(chunks.length).toBeGreaterThan(1); + for (const chunk of chunks) { + expect(chunk.length).toBeLessThanOrEqual(4096); + expect(chunk.length).toBeGreaterThan(0); + expect(/^[\uDC00-\uDFFF]|[\uD800-\uDBFF]$/u.test(chunk)).toBe(false); + for (const run of chunk.match(/\s+|\S+/gu) ?? []) { + expect(run.length).toBeLessThanOrEqual(256); + } + } } finally { count.mockRestore(); } }); + + it('keeps empty input empty', () => { + expect(counter('')).toBe(0); + expect(Tokenizer.getTokenCount('', encoding)).toBe(0); + }); }); describe('countExactTokens', () => { @@ -107,11 +215,8 @@ describe('Tokenizer', () => { await Tokenizer.initEncoding('o200k_base'); }); - it('tokenizes oversized input whole instead of falling back to byte length', () => { - /** The fast-path estimate is byte length, several times the real count on - * ordinary text; a figure ADDED to provider accounting cannot carry that. - * Nor can it carry a sum of slices: a BPE merge spanning a seam is charged - * twice, so the whole input is tokenized in one pass. */ + it("tokenizes oversized input whole without the budgeting counter's seam approximation", () => { + /** Provider accounting must not carry a sum of slices: BPE merges can change at seams. */ const text = 'word '.repeat(4096); const exact = Tokenizer.countExactTokens(text, 'o200k_base'); expect(exact).toBeGreaterThan(0); diff --git a/packages/api/src/utils/tokenizer.ts b/packages/api/src/utils/tokenizer.ts index 06de186c6d8..83148ba86d3 100644 --- a/packages/api/src/utils/tokenizer.ts +++ b/packages/api/src/utils/tokenizer.ts @@ -6,17 +6,41 @@ export type EncodingName = 'o200k_base' | 'claude'; type EncodingData = ConstructorParameters[0]; const MAX_TOKENIZER_INPUT_LENGTH = 4 * 1024; +/** Bound BPE's quadratic merge work on uninterrupted text, including whitespace. */ +const MAX_TOKENIZER_RUN_LENGTH = 256; +const WHITESPACE = /\s/u; -function estimateBoundedTokenCount(text: string): number { - return Buffer.byteLength(text, 'utf8'); +/** Near-exact budgeting count; seams can change BPE merges, but no content is omitted. */ +function countBoundedTokens(text: string, tokenizer: AiTokenizer): number { + let tokens = 0; + for (let offset = 0; offset < text.length; ) { + let end = Math.min(offset + MAX_TOKENIZER_INPUT_LENGTH, text.length); + let runLength = 0; + let previousWhitespace = false; + for (let index = offset; index < end; index++) { + const whitespace = WHITESPACE.test(text[index]); + runLength = whitespace === previousWhitespace ? runLength + 1 : 1; + previousWhitespace = whitespace; + if (runLength >= MAX_TOKENIZER_RUN_LENGTH) { + end = index + 1; + break; + } + } + const last = text.charCodeAt(end - 1); + const next = text.charCodeAt(end); + if (last >= 0xd800 && last <= 0xdbff && next >= 0xdc00 && next <= 0xdfff) { + end--; + } + tokens += tokenizer.count(text.slice(offset, end)); + offset = end; + } + return tokens; } function estimateUnavailableTokenCount(text: string): number { - return Math.ceil(text.length / 4); -} - -function requiresTokenEstimate(text: string): boolean { - return text.length > MAX_TOKENIZER_INPUT_LENGTH; + return text.length > MAX_TOKENIZER_INPUT_LENGTH + ? Buffer.byteLength(text, 'utf8') + : Math.ceil(text.length / 4); } class Tokenizer { @@ -41,7 +65,7 @@ class Tokenizer { return this.loadingPromises[encoding]; } - /** Returns a counter that avoids expensive tokenization for oversized content. */ + /** Returns an initialized, deterministic budgeting counter that throws on tokenizer failure. */ async createExactTokenCounter(encoding: EncodingName): Promise<(text: string) => number> { await this.initEncoding(encoding); const tokenizer = this.tokenizersCache[encoding]; @@ -49,11 +73,8 @@ class Tokenizer { throw new Error(`Tokenizer encoding failed to initialize: ${encoding}`); } return (text: string): number => { - if (requiresTokenEstimate(text)) { - return estimateBoundedTokenCount(text); - } try { - return tokenizer.count(text); + return countBoundedTokens(text, tokenizer); } catch (error) { this.handleCountError(encoding, tokenizer, error); throw error; @@ -62,16 +83,13 @@ class Tokenizer { } getTokenCount(text: string, encoding: EncodingName = 'o200k_base'): number { - if (requiresTokenEstimate(text)) { - return estimateBoundedTokenCount(text); - } const tokenizer = this.tokenizersCache[encoding]; if (!tokenizer) { this.initEncoding(encoding); return estimateUnavailableTokenCount(text); } try { - return tokenizer.count(text); + return countBoundedTokens(text, tokenizer); } catch (error) { this.handleCountError(encoding, tokenizer, error); return estimateUnavailableTokenCount(text); @@ -83,13 +101,11 @@ class Tokenizer { * cannot be produced, so the caller can omit a figure instead of publishing a * guess. Used for a value that is ADDED to exact provider accounting — see the * retained tool results of a turn stopped at the tool-call limit — where - * {@link getTokenCount}'s two fallbacks would corrupt the sum: the cold-start - * estimate is a character ratio, and the oversized-input shortcut returns byte - * length, several times the real count on ordinary text. + * {@link getTokenCount}'s cold-start/error estimate and chunk-boundary + * approximation must not be added to exact provider accounting. * * The input is tokenized whole. Counting it in slices would be cheaper but not - * exact — a BPE merge spanning a seam is charged twice, ~1 token per 4 KiB - * measured — and this is the one caller that needs the real number. It costs + * exact — seams can change BPE merges — and this caller needs the real number. It costs * ~60 ms/MB, so the caller decides how much content is worth counting * (`endpoints.agents.maxRetainedToolCountChars`) rather than a bound here. */ diff --git a/packages/api/src/utils/tokens.spec.ts b/packages/api/src/utils/tokens.spec.ts index 62801ece029..d7b2ca717ba 100644 --- a/packages/api/src/utils/tokens.spec.ts +++ b/packages/api/src/utils/tokens.spec.ts @@ -1,6 +1,56 @@ import { EModelEndpoint } from 'librechat-data-provider'; import type { EndpointTokenConfig } from '~/types'; -import { getModelMaxTokens, getModelMaxOutputTokens } from './tokens'; +import { getModelMaxTokens, getModelMaxOutputTokens, findMatchingPattern } from './tokens'; + +describe('Bedrock OpenAI context windows', () => { + const models = [ + 'gpt-6-sol', + 'gpt-6-luna', + 'gpt-6-astra', + 'gpt-5.6-sol', + 'gpt-5.6-terra', + 'gpt-5.6-luna', + ]; + + it.each(models)('resolves %s across Bedrock inference profiles', (model) => { + for (const prefix of ['', 'us.', 'global.']) { + expect(getModelMaxTokens(`${prefix}openai.${model}`, EModelEndpoint.bedrock)).toBe(950000); + } + }); + + it('matches versioned profiles without changing the native OpenAI context window', () => { + expect(getModelMaxTokens('us.openai.gpt-6-sol-2026-09-22', EModelEndpoint.bedrock)).toBe( + 950000, + ); + expect(getModelMaxTokens('gpt-6-sol', EModelEndpoint.openAI)).toBe(1050000); + }); + + it('keeps GPT-OSS limits and unknown-model fallback unchanged', () => { + for (const model of ['openai.gpt-oss-20b-1:0', 'us.openai.gpt-oss-120b-1:0']) { + expect(getModelMaxTokens(model, EModelEndpoint.bedrock)).toBe(128000); + } + expect(getModelMaxTokens('us.openai.unknown-model', EModelEndpoint.bedrock)).toBeUndefined(); + }); + + it('preserves explicit endpoint overrides while falling back for unlisted profiles', () => { + const override: EndpointTokenConfig = { + 'openai.gpt-6-sol': { context: 64000, prompt: 1, completion: 1 }, + }; + expect(getModelMaxTokens('us.openai.gpt-6-sol', EModelEndpoint.bedrock, override)).toBe(64000); + expect(getModelMaxTokens('global.openai.gpt-6-luna', EModelEndpoint.bedrock, override)).toBe( + 950000, + ); + }); + + it('prefers the GPT-5.6 key over shorter keys proposed for Bedrock Mantle', () => { + const keys = { + 'openai.gpt-5': 272000, + 'openai.gpt-5.5': 272000, + 'openai.gpt-5.6': 950000, + }; + expect(findMatchingPattern('us.openai.gpt-5.6-sol', keys)).toBe('openai.gpt-5.6'); + }); +}); describe('getModelMaxTokens partial-override fallback', () => { const partialOverride: EndpointTokenConfig = { diff --git a/packages/api/src/utils/tokens.ts b/packages/api/src/utils/tokens.ts index f9ad7c0792a..07ce2637d94 100644 --- a/packages/api/src/utils/tokens.ts +++ b/packages/api/src/utils/tokens.ts @@ -419,9 +419,16 @@ const amazonModels = { 'nova-2-pro': 995000, // -5000 from max }; +/** Bedrock can reject near-limit GPT prompts before its published window is reached. */ +const bedrockOpenAIContext = 950000; + const openAIBedrockModels = { 'openai.gpt-oss-20b': 128000, 'openai.gpt-oss-120b': 128000, + 'openai.gpt-5.6': bedrockOpenAIContext, + 'openai.gpt-6-astra': bedrockOpenAIContext, + 'openai.gpt-6-sol': bedrockOpenAIContext, + 'openai.gpt-6-luna': bedrockOpenAIContext, }; const bedrockModels = { diff --git a/packages/data-provider/src/code/workspace.ts b/packages/data-provider/src/code/workspace.ts index b431232d96d..52beb3c6f65 100644 --- a/packages/data-provider/src/code/workspace.ts +++ b/packages/data-provider/src/code/workspace.ts @@ -5,6 +5,8 @@ export const CODE_WORKSPACE_MAX_COUNT = 32; export const CODE_ENVIRONMENT_DECISION_VERSION = 1 as const; /** API/client protocol for an owner's explicit move of a sealed environment decision. */ export const CODE_ENVIRONMENT_MOVE_VERSION = 1 as const; +/** Additive capability for replacing a missing workspace without disabling moves in V1 clients. */ +export const CODE_WORKSPACE_RECOVERY_VERSION = 1 as const; export const CODE_WORKSPACE_OPERATIONS = [ 'read_file', 'search_text', diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index d41e03d260b..7ff73ec7a80 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -673,6 +673,34 @@ describe('agent background completion batch config', () => { }); describe('agent event runtime config', () => { + it('defaults and bounds durable-worker idle polling intervals', () => { + const parsed = configSchema.parse({ + version: '1.0', + endpoints: { agents: { eventDriven: { idlePolling: {} } } }, + }); + expect(parsed.endpoints?.agents?.eventDriven?.idlePolling).toEqual({ + deliveryMaxIntervalMs: 15_000, + queuedTurnMaxIntervalMs: 120_000, + maintenanceMaxIntervalMs: 120_000, + completionWaitMaxIntervalMs: 60_000, + }); + for (const [key, value] of [ + ['deliveryMaxIntervalMs', 0], + ['queuedTurnMaxIntervalMs', 29_999], + ['maintenanceMaxIntervalMs', 300_001], + ['maintenanceMaxIntervalMs', 30_000.5], + ['completionWaitMaxIntervalMs', 4_999], + ['completionWaitMaxIntervalMs', 300_001], + ] as const) { + expect( + configSchema.safeParse({ + version: '1.0', + endpoints: { agents: { eventDriven: { idlePolling: { [key]: value } } } }, + }).success, + ).toBe(false); + } + }); + it('accepts the routing choice and ignores removed rollout fields', () => { const result = configSchema.safeParse({ version: '1.0', diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 521863b2d46..3acf59cd6b6 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -30,7 +30,11 @@ import { MAX_SUBAGENTS_CEILING, DEFAULT_MAX_RETAINED_TOOL_COUNT_CHARS, } from './limits'; -import { CODE_ENVIRONMENT_DECISION_VERSION, CODE_ENVIRONMENT_MOVE_VERSION } from './code/workspace'; +import { + CODE_ENVIRONMENT_DECISION_VERSION, + CODE_ENVIRONMENT_MOVE_VERSION, + CODE_WORKSPACE_RECOVERY_VERSION, +} from './code/workspace'; import { ComponentTypes, SettingTypes, OptionTypes } from './generate'; import { STATEFUL_CODE_ENVIRONMENTS } from './stateful-code'; import { specsConfigSchema, TSpecsConfig } from './models'; @@ -1374,7 +1378,8 @@ export const agentsEndpointSchema = baseEndpointSchema }) .optional(), /** Server-only policy letting a conversation's owner move its sealed attached decision - * onto the environments its agents now use. Omit to keep sealed decisions immovable. */ + * onto the environments its agents now use, or recover a missing workspace. Omit to keep + * sealed decisions immovable. */ conversationMoves: z .object({ enabled: z.boolean().optional(), @@ -1507,6 +1512,43 @@ export const agentsEndpointSchema = baseEndpointSchema eventDriven: z .object({ selfUrl: z.string().url().optional(), + /** Every replica keeps polling Mongo as a crash-recovery fallback. Wakes + * keep local delivery prompt; these caps bound work missed across replicas. */ + idlePolling: z + .object({ + deliveryMaxIntervalMs: z + .number() + .int() + .min(1_000) + .max(300_000) + .optional() + .default(15_000), + queuedTurnMaxIntervalMs: z + .number() + .int() + .min(30_000) + .max(300_000) + .optional() + .default(120_000), + maintenanceMaxIntervalMs: z + .number() + .int() + .min(30_000) + .max(300_000) + .optional() + .default(120_000), + /** Longest a background or subagent completion re-checks whether its + * result and parent turn are ready. The events it waits on expedite it, + * so this bounds missed signals rather than normal delivery latency. */ + completionWaitMaxIntervalMs: z + .number() + .int() + .min(5_000) + .max(300_000) + .optional() + .default(60_000), + }) + .optional(), }) .optional(), /** Conversational background-task delivery policy. Automatic completion wakeups are @@ -2390,6 +2432,9 @@ export type TStartupConfig = { /** Owner moves of a sealed code-environment decision supported by the API. Clients must not * offer to move a conversation unless this is advertised. */ codeEnvironmentMoveVersion?: typeof CODE_ENVIRONMENT_MOVE_VERSION; + /** Additive recovery support. Clients require this and the move capability before replacing + * a missing workspace. Keeping it separate preserves exact-version checks in older clients. */ + codeWorkspaceRecoveryVersion?: typeof CODE_WORKSPACE_RECOVERY_VERSION; interface?: TInterfaceConfig; turnstile?: TTurnstileConfig; balance?: TBalanceConfig; diff --git a/packages/data-provider/src/filters.spec.ts b/packages/data-provider/src/filters.spec.ts index c92df1ea8a1..eaae6af7841 100644 --- a/packages/data-provider/src/filters.spec.ts +++ b/packages/data-provider/src/filters.spec.ts @@ -57,7 +57,7 @@ describe('filtersConfigSchema', () => { ).toBe(false); }); - it('accepts omitted, audit, and block actions while rejecting unsupported actions', () => { + it('accepts opt-in redact, audit, and block actions while rejecting unsupported actions', () => { expect(filtersConfigSchema.parse({ messages: { pii: {} } })).toEqual({ messages: { pii: {} }, }); @@ -71,8 +71,44 @@ describe('filtersConfigSchema', () => { prompts: { pii: { action: 'block' } }, }); expect(filtersConfigSchema.safeParse({ messages: { pii: { action: 'redact' } } }).success).toBe( + true, + ); + expect( + filtersConfigSchema.safeParse({ messages: { pii: { action: 'forward' } } }).success, + ).toBe(false); + }); + + it('bounds opt-in transformation controls and validates placeholder categories', () => { + expect(filtersConfigSchema.parse({ messages: { pii: { action: 'redact' } } })).toEqual({ + messages: { pii: { action: 'redact' } }, + }); + expect(filtersConfigSchema.safeParse({ messages: { pii: { maxCharacters: 0 } } }).success).toBe( false, ); + expect( + filtersConfigSchema.safeParse({ messages: { pii: { maxMatches: 4_097 } } }).success, + ).toBe(false); + expect( + filtersConfigSchema.safeParse({ + messages: { + pii: { + action: 'redact', + customPatterns: [{ id: 'email', label: 'Email', regex: 'a@b', category: 'email' }], + }, + }, + }).success, + ).toBe(true); + expect( + filtersConfigSchema.safeParse({ + messages: { + pii: { + customPatterns: [ + { id: 'email', label: 'Email', regex: 'a@b', category: 'secret-value' }, + ], + }, + }, + }).success, + ).toBe(false); }); it('combines active-pattern and selected-field checks', () => { diff --git a/packages/data-provider/src/filters.ts b/packages/data-provider/src/filters.ts index 9d010c727e0..4368777d199 100644 --- a/packages/data-provider/src/filters.ts +++ b/packages/data-provider/src/filters.ts @@ -162,7 +162,10 @@ export const fileFilterFieldSchema = z.enum(FILE_FILTER_FIELDS); export const toolArgumentFilterFieldSchema = z.enum(TOOL_ARGUMENT_FILTER_FIELDS); export const modelParameterFilterFieldSchema = z.enum(MODEL_PARAMETER_FILTER_FIELDS); export const filterPiiStarterPatternSchema = z.enum(FILTER_PII_STARTER_PATTERNS); -export const filterPiiActionSchema = z.enum(['block', 'audit']); +export const filterPiiActionSchema = z.enum(['block', 'audit', 'redact']); +export const filterPiiCategorySchema = z.enum(['email', 'phone', 'name', 'credential', 'custom']); +export const DEFAULT_PII_REDACTION_MAX_CHARACTERS = 65_536; +export const DEFAULT_PII_REDACTION_MAX_MATCHES = 256; export const actionMetadataFilterFieldSchema = z.enum(ACTION_METADATA_FILTER_FIELDS); export const unattributedAssistantContentSchema = z.enum(['model_output', 'inspect']); export type UnattributedAssistantContent = z.infer; @@ -180,6 +183,7 @@ export type ToolArgumentFilterField = z.infer; export type ActionMetadataFilterField = z.infer; export type FilterPiiAction = z.infer; +export type FilterPiiCategory = z.infer; export const userSubmittedMessageFieldPathSchema = z .object({ @@ -279,6 +283,7 @@ export const filterPiiCustomPatternSchema = z id: z.string().min(1).max(MAX_PII_PATTERN_ID_LENGTH), label: z.string().min(1).max(MAX_PII_PATTERN_LABEL_LENGTH), regex: filterPiiRegexSchema, + category: filterPiiCategorySchema.optional(), }) .strict(); @@ -288,6 +293,8 @@ function createPiiFilterSchema(fieldSchema: Field) { return z .object({ action: filterPiiActionSchema.optional(), + maxCharacters: z.number().int().positive().max(262_144).optional(), + maxMatches: z.number().int().positive().max(4_096).optional(), fields: z.array(fieldSchema).min(1).max(MAX_PII_PATTERNS_PER_SOURCE).optional(), starterPatterns: z .array(filterPiiStarterPatternSchema) diff --git a/packages/data-schemas/src/methods/agent.spec.ts b/packages/data-schemas/src/methods/agent.spec.ts index 0fbe6b0e491..1eecbcdeb59 100644 --- a/packages/data-schemas/src/methods/agent.spec.ts +++ b/packages/data-schemas/src/methods/agent.spec.ts @@ -5029,6 +5029,37 @@ describe('Support Contact Field', () => { }); }); + test('unrestricted listing still applies filters and paginates without overlap', async () => { + const first = await getListAgentsByAccess({ + accessibleIds: null, + otherParams: { name: /^Agent A/ }, + limit: 1, + }); + expect(first.data).toHaveLength(1); + expect(first.has_more).toBe(true); + expect(first.after).toBeTruthy(); + + const second = await getListAgentsByAccess({ + accessibleIds: null, + otherParams: { name: /^Agent A/ }, + limit: 1, + after: first.after, + }); + expect(second.data).toHaveLength(1); + expect(second.has_more).toBe(true); + + const third = await getListAgentsByAccess({ + accessibleIds: null, + otherParams: { name: /^Agent A/ }, + limit: 1, + after: second.after, + }); + expect(third.has_more).toBe(false); + expect( + [first, second, third].flatMap((page) => page.data.map((agent) => agent.id)).sort(), + ).toEqual([agentA1.id, agentA2.id, agentA3.id].sort()); + }); + test('should return empty list when user has no accessible agents (empty accessibleIds)', async () => { // User B has no agents and no shared agents const result = await getListAgentsByAccess({ @@ -5042,6 +5073,90 @@ describe('Support Contact Field', () => { expect(result.last_id).toBeNull(); }); + test('unrestricted listing cannot cross tenant boundaries', async () => { + const tenantA = `tenant-a-${uuidv4()}`; + const tenantB = `tenant-b-${uuidv4()}`; + const name = 'Shared Name'; + const agentInA = await tenantStorage.run({ tenantId: tenantA }, () => + createAgent({ + id: `agent_${uuidv4()}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userA, + }), + ); + await tenantStorage.run({ tenantId: tenantB }, () => + createAgent({ + id: `agent_${uuidv4()}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userB, + }), + ); + + const result = await tenantStorage.run({ tenantId: tenantA }, () => + getListAgentsByAccess({ accessibleIds: null, tenantId: tenantA, otherParams: { name } }), + ); + expect(result.data.map((agent) => agent.id)).toEqual([agentInA.id]); + const denied = await tenantStorage.run({ tenantId: tenantA }, () => + getListAgentsByAccess({ accessibleIds: [], otherParams: { name } }), + ); + expect(denied.data).toHaveLength(0); + }); + + test('unrestricted list uses the supplied tenant even without request context', async () => { + const tenantA = `tenant-a-${uuidv4()}`; + const tenantB = `tenant-b-${uuidv4()}`; + const name = 'Same Agent Name'; + const legacy = await createAgent({ + id: `agent_${uuidv4()}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userA, + }); + const agentInA = await tenantStorage.run({ tenantId: tenantA }, () => + createAgent({ + id: `agent_${uuidv4()}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userA, + }), + ); + await tenantStorage.run({ tenantId: tenantB }, () => + createAgent({ + id: `agent_${uuidv4()}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userB, + }), + ); + + const scoped = await getListAgentsByAccess({ + accessibleIds: null, + tenantId: tenantA, + otherParams: { name }, + }); + expect(scoped.data.map((agent) => agent.id)).toEqual([agentInA.id]); + const legacyOnly = await getListAgentsByAccess({ + accessibleIds: null, + tenantId: null, + otherParams: { name }, + }); + expect(legacyOnly.data.map((agent) => agent.id)).toEqual([legacy.id]); + const omitted = await getListAgentsByAccess({ accessibleIds: null, otherParams: { name } }); + expect(omitted.data.map((agent) => agent.id)).toEqual([legacy.id]); + }); + + test('accepts plain string IDs from the API resolver without weakening ACL filtering', async () => { + const result = await getListAgentsByAccess({ accessibleIds: [agentA2._id.toString()] }); + expect(result.data.map((agent) => agent.id)).toEqual([agentA2.id]); + }); + test('should not return other users agents when accessibleIds is empty', async () => { // User B trying to list agents with empty accessibleIds should not see User A's agents const result = await getListAgentsByAccess({ diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index 701a5056580..4295248995b 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -606,13 +606,16 @@ export function createAgentMethods( getListAgentsByAccess: ({ accessibleIds, otherParams, + tenantId, limit, after, includeSkillConfig, includeExecutionConfig, }: { - accessibleIds?: Types.ObjectId[]; + accessibleIds?: Array | null; otherParams?: Record; + /** Authenticated tenant for unrestricted listings; null/omitted restricts to legacy agents. */ + tenantId?: string | null; limit?: number | null; after?: string | null; includeSkillConfig?: boolean; @@ -1423,19 +1426,23 @@ export function createAgentMethods( } /** - * Get agents by accessible IDs with cursor pagination. Defaults to a 100-page - * limit (max 1000); pass `limit: null` to opt out entirely. + * Get agents by accessible IDs with cursor pagination. Pass `accessibleIds: null` + * only after a management-capability check, with the authenticated tenantId + * (or null for legacy agents); `[]` and omitted IDs match nothing. + * Defaults to a 100-page limit (max 1000); pass `limit: null` to opt out entirely. */ async function getListAgentsByAccess({ accessibleIds = [], otherParams = {}, + tenantId, limit = 100, after = null, includeSkillConfig = false, includeExecutionConfig = false, }: { - accessibleIds?: Types.ObjectId[]; + accessibleIds?: Array | null; otherParams?: Record; + tenantId?: string | null; limit?: number | null; after?: string | null; includeSkillConfig?: boolean; @@ -1456,7 +1463,9 @@ export function createAgentMethods( const baseQuery: Record = { ...otherParams, - _id: { $in: accessibleIds }, + ...(accessibleIds === null + ? { tenantId: tenantId ?? null } + : { _id: { $in: accessibleIds } }), }; if (after) { diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index 6c1b2eea3d1..704327518d3 100644 --- a/packages/data-schemas/src/methods/conversation.spec.ts +++ b/packages/data-schemas/src/methods/conversation.spec.ts @@ -6449,7 +6449,9 @@ describe('Conversation Operations', () => { }, ]); - await expect(methods.expireLegacyAgentEventActorReceipts(now, 1)).resolves.toBe(1); + const activity = { found: false }; + await expect(methods.expireLegacyAgentEventActorReceipts(now, 1, activity)).resolves.toBe(1); + expect(activity.found).toBe(true); await expect( Conversation.findOne({ conversationId: oldConversationId }) .select('+agentEventActorReconciliations') @@ -6625,9 +6627,20 @@ describe('Conversation Operations', () => { awaitTerminalHandling: true, }); - /** The first raw page contains no expired receipt. The second contains - * a protected receipt, and only the third reaches removable work. */ - await expect(methods.expireLegacyAgentEventActorReceipts(now, 1)).resolves.toBe(0); + /** The first raw page contains no expired receipt. It must still keep + * discovery active, without querying delivery protection for an empty ID set. + * The second is protected, and only the third reaches removable work. */ + const activity = { found: false }; + const protection = jest.spyOn(Delivery, 'find'); + try { + await expect(methods.expireLegacyAgentEventActorReceipts(now, 1, activity)).resolves.toBe( + 0, + ); + expect(activity.found).toBe(true); + expect(protection).not.toHaveBeenCalled(); + } finally { + protection.mockRestore(); + } await expect(methods.expireLegacyAgentEventActorReceipts(now, 1)).resolves.toBe(0); await expect(methods.expireLegacyAgentEventActorReceipts(now, 1)).resolves.toBe(1); await expect( diff --git a/packages/data-schemas/src/methods/conversation.ts b/packages/data-schemas/src/methods/conversation.ts index f81e52fb117..71811f2f440 100644 --- a/packages/data-schemas/src/methods/conversation.ts +++ b/packages/data-schemas/src/methods/conversation.ts @@ -449,7 +449,11 @@ export interface ConversationMethods { getAgentEventActorReconciliationStorageMetrics( now: Date, ): Promise; - expireLegacyAgentEventActorReceipts(now: Date, limit?: number): Promise; + expireLegacyAgentEventActorReceipts( + now: Date, + limit?: number, + activity?: { found: boolean }, + ): Promise; reserveSubagentThread(input: { user: string; conversationId: string; @@ -1782,7 +1786,11 @@ export function createConversationMethods( /** Bounded mixed-version cleanup for terminal receipts embedded by older * builds. New receipts expire through the delivery collection TTL index, * but dormant legacy conversations need an independent retirement path. */ - async function expireLegacyAgentEventActorReceipts(now: Date, limit = 100): Promise { + async function expireLegacyAgentEventActorReceipts( + now: Date, + limit = 100, + activity?: { found: boolean }, + ): Promise { if (Number.isNaN(now.getTime())) { throw new TypeError('now must be a valid date'); } @@ -1821,6 +1829,9 @@ export function createConversationMethods( return 0; } legacyReceiptExpiryCursor = candidates[candidates.length - 1]._id; + // A full page is not an empty sweep: later conversations may contain + // expired receipts. Keep scanning at base cadence until the page walk ends. + if (activity != null && candidates.length === boundedLimit) activity.found = true; const expiredInvocationIds = [ ...new Set( candidates.flatMap((candidate) => @@ -1830,6 +1841,8 @@ export function createConversationMethods( ), ), ]; + if (expiredInvocationIds.length === 0) return 0; + if (activity != null) activity.found = true; const protectedInvocationIds = new Set( await Delivery.find({ deliveryKey: { $in: expiredInvocationIds }, diff --git a/packages/data-schemas/src/methods/memory.spec.ts b/packages/data-schemas/src/methods/memory.spec.ts index ec1f3d4362e..09d9d7dec20 100644 --- a/packages/data-schemas/src/methods/memory.spec.ts +++ b/packages/data-schemas/src/methods/memory.spec.ts @@ -256,6 +256,31 @@ describe('memory partitions', () => { expect(partitionA.withKeys).not.toContain('personal one'); }); + it('signals a read failure rather than reporting an empty memory partition', async () => { + const find = jest.spyOn(MemoryEntry, 'find').mockImplementationOnce(() => { + throw new Error('simulated memory read failure'); + }); + try { + const failed = await methods.getFormattedMemories({ userId }); + expect(failed).toEqual({ + withKeys: undefined, + withoutKeys: undefined, + totalTokens: 0, + tokenCountsByKey: new Map(), + readFailed: true, + }); + } finally { + find.mockRestore(); + } + const empty = await methods.getFormattedMemories({ userId }); + expect(empty).toEqual({ + withKeys: '', + withoutKeys: '', + totalTokens: 0, + tokenCountsByKey: new Map(), + }); + }); + it('returns every partition from getAllUserMemories and wipes them all on deleteAllUserMemories', async () => { await methods.setMemory({ userId, key: 'one', value: 'personal', tokenCount: 1 }); await methods.setMemory({ userId, key: 'two', value: 'agent a', tokenCount: 1, agentId }); diff --git a/packages/data-schemas/src/methods/memory.ts b/packages/data-schemas/src/methods/memory.ts index ebff3a21601..c41f3763dad 100644 --- a/packages/data-schemas/src/methods/memory.ts +++ b/packages/data-schemas/src/methods/memory.ts @@ -311,11 +311,14 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { return { withKeys, withoutKeys, totalTokens, tokenCountsByKey }; } catch (error) { logger.error('Failed to get formatted memories:', error); + /** Keep the current error handling without presenting a read failure + * as an eligible, empty memory partition to prompt consumers. */ return { - withKeys: '', - withoutKeys: '', + withKeys: undefined, + withoutKeys: undefined, totalTokens: 0, tokenCountsByKey: new Map(), + readFailed: true, }; } } diff --git a/packages/data-schemas/src/methods/queuedTurn.spec.ts b/packages/data-schemas/src/methods/queuedTurn.spec.ts index f4718dca8e6..74c2f95f54e 100644 --- a/packages/data-schemas/src/methods/queuedTurn.spec.ts +++ b/packages/data-schemas/src/methods/queuedTurn.spec.ts @@ -220,6 +220,31 @@ describe('agent queued turn methods', () => { ).resolves.toMatchObject({ status: 'queued', sequence: expect.any(Number) }); }); + it('reports reservation-only recovery activity when deleted lanes yield no deliveries', async () => { + const first = await methods.enqueueAgentQueuedTurn( + enqueueInput({ clientRequestId: 'lost-lane-first' }), + ); + const second = await methods.enqueueAgentQueuedTurn( + enqueueInput({ clientRequestId: 'lost-lane-second' }), + ); + await Turn.updateMany( + { _id: { $in: [first.turn.queuedTurnId, second.turn.queuedTurnId] } }, + { + $set: { status: 'reserving' }, + $unset: { sequence: 1 }, + }, + ); + await Sequence.deleteMany({}); + for (let i = 0; i < 2; i++) { + const activity = { found: false }; + expect(await methods.findQueuedTurnsNeedingDelivery(1, activity)).toEqual([]); + expect(activity.found).toBe(true); + } + const empty = { found: false }; + expect(await methods.findQueuedTurnsNeedingDelivery(1, empty)).toEqual([]); + expect(empty.found).toBe(false); + }); + it('caps each conversation at 100 active turns while preserving exact replay', async () => { const inputs = Array.from({ length: 101 }, (_, index) => enqueueInput({ clientRequestId: `capacity-${index}`, text: `turn ${index}` }), diff --git a/packages/data-schemas/src/methods/queuedTurn.ts b/packages/data-schemas/src/methods/queuedTurn.ts index 22085ba0d99..f10bcb17ae4 100644 --- a/packages/data-schemas/src/methods/queuedTurn.ts +++ b/packages/data-schemas/src/methods/queuedTurn.ts @@ -203,9 +203,13 @@ export interface AgentQueuedTurnMethods { listAgentQueuedTurnReceipts: ( input: AgentQueuedTurnConversationScope & { clientRequestIds?: readonly string[] }, ) => Promise; - findQueuedTurnsNeedingDelivery: (limit?: number) => Promise; + findQueuedTurnsNeedingDelivery: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; claimQueuedTurnsForAdmissionReconciliation: ( input: ClaimAgentQueuedTurnReconciliationInput, + activity?: { found: boolean }, ) => Promise; deferAgentQueuedTurnAdmissionReconciliation: ( input: AgentQueuedTurnConversationScope & { @@ -1381,7 +1385,10 @@ export function createAgentQueuedTurnMethods( return { outcome: 'not_cancellable', turn: toRecord(current) }; } - async function findQueuedTurnsNeedingDelivery(limit = 100): Promise { + async function findQueuedTurnsNeedingDelivery( + limit = 100, + activity?: { found: boolean }, + ): Promise { if (!Number.isSafeInteger(limit) || limit <= 0 || limit > 1000) { throw new TypeError('Agent queued turn recovery limit must be between 1 and 1000'); } @@ -1390,6 +1397,7 @@ export function createAgentQueuedTurnMethods( .sort({ createdAt: 1, _id: 1 }) .limit(limit) .lean(); + if (activity != null && reservations.length > 0) activity.found = true; for (const reservation of reservations) { const scopeInput: AgentQueuedTurnConversationScope = { user: reservation.user, @@ -1460,6 +1468,7 @@ export function createAgentQueuedTurnMethods( async function claimQueuedTurnsForAdmissionReconciliation( input: ClaimAgentQueuedTurnReconciliationInput, + activity?: { found: boolean }, ): Promise { const limit = input.limit ?? 100; if (!Number.isSafeInteger(limit) || limit <= 0 || limit > 1000) { @@ -1519,6 +1528,7 @@ export function createAgentQueuedTurnMethods( if (candidates.length === 0) { return []; } + if (activity != null) activity.found = true; await Turn().updateMany( { ...eligible, diff --git a/packages/data-schemas/src/methods/triggerDelivery.spec.ts b/packages/data-schemas/src/methods/triggerDelivery.spec.ts index 1f0fcc74f6a..a6ad23cb447 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.spec.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.spec.ts @@ -453,6 +453,594 @@ describe('agent trigger delivery methods', () => { ).resolves.toEqual({ status: 'live', leaseUntil: renewedUntil }); }); + describe('expediteAgentTriggerDeliveries', () => { + const background = { id: 'background-tool-completion', type: 'internal' }; + const later = new Date(START.getTime() + 30_000); + const capable = { + workerId: 'background-capable-worker', + claimToken: 'background-capable-claim', + now: START, + leaseUntil: new Date(START.getTime() + 60_000), + workerCapabilities: [AGENT_TRIGGER_WORKER_CAPABILITY_BACKGROUND_COMPLETION_RECEIPT_V2], + }; + const waiting = (overrides: Partial[0]> = {}) => + methods.enqueueAgentTriggerDelivery( + enqueueInput({ + orderingKey: `background-lane-${counter + 1}`, + envelope: { event: { source: background } }, + requiredWorkerCapability: + AGENT_TRIGGER_WORKER_CAPABILITY_BACKGROUND_COMPLETION_RECEIPT_V2, + availableAt: later, + ...overrides, + }), + ); + + it('makes a deferred delivery claimable now once its result is durable', async () => { + const target = await waiting(); + const sibling = await waiting(); + await expect(methods.claimNextAgentTriggerDelivery(capable)).resolves.toBeNull(); + + await expect( + methods.expediteAgentTriggerDeliveries({ + deliveryKeys: [target.delivery.deliveryKey], + sourceIds: [background.id], + now: START, + }), + ).resolves.toEqual({ expedited: 1, held: 0 }); + + await expect(methods.claimNextAgentTriggerDelivery(capable)).resolves.toMatchObject({ + id: target.delivery.id, + }); + await expect( + methods.claimNextAgentTriggerDelivery({ ...capable, claimToken: 'second-claim' }), + ).resolves.toBeNull(); + const untouched = await Delivery.findById(sibling.delivery.id).lean(); + expect(untouched?.availableAt).toEqual(later); + }); + + it("moves only the principal's waiting rows from the named sources", async () => { + const user = new mongoose.Types.ObjectId(); + const mine = await waiting({ user }); + const otherUser = await waiting(); + const otherSource = await waiting({ + user, + envelope: { event: { source: { id: 'agent-queued-turn', type: 'internal' } } }, + }); + const external = await waiting({ + user, + envelope: { event: { source: { id: background.id, type: 'webhook' } } }, + }); + const due = await waiting({ user, availableAt: START }); + + await expect( + methods.expediteAgentTriggerDeliveries({ user, sourceIds: [background.id], now: START }), + ).resolves.toEqual({ expedited: 1, held: 0 }); + + const rows = await Delivery.find({ + _id: { + $in: [mine, otherUser, otherSource, external, due].map((row) => row.delivery.id), + }, + }).lean(); + const availableAt = new Map(rows.map((row) => [String(row._id), row.availableAt])); + expect(availableAt.get(mine.delivery.id)).toEqual(START); + expect(availableAt.get(otherUser.delivery.id)).toEqual(later); + expect(availableAt.get(otherSource.delivery.id)).toEqual(later); + expect(availableAt.get(external.delivery.id)).toEqual(later); + expect(availableAt.get(due.delivery.id)).toEqual(START); + }); + + it('marks a delivery a worker currently holds instead of moving it', async () => { + const user = new mongoose.Types.ObjectId(); + const held = await waiting({ user, availableAt: START }); + const claim = await methods.claimNextAgentTriggerDelivery(capable); + expect(claim).toMatchObject({ id: held.delivery.id }); + const before = await Delivery.findById(held.delivery.id).lean(); + + await expect( + methods.expediteAgentTriggerDeliveries({ + user, + sourceIds: [background.id], + now: new Date(START.getTime() - 60_000), + }), + ).resolves.toEqual({ expedited: 0, held: 1 }); + + const after = await Delivery.findById(held.delivery.id).lean(); + expect(after?.status).toBe(before?.status); + expect(after?.availableAt).toEqual(before?.availableAt); + expect(after?.wakeRequestedAt).toEqual(new Date(START.getTime() - 60_000)); + }); + + it('re-checks at once when readiness changed while the delivery was held', async () => { + const user = new mongoose.Types.ObjectId(); + const held = await waiting({ user, availableAt: START }); + const claim = await methods.claimNextAgentTriggerDelivery(capable); + expect(claim).toMatchObject({ id: held.delivery.id }); + const fence = { + id: held.delivery.id, + workerId: capable.workerId, + claimToken: capable.claimToken, + }; + const attempt = await methods.beginAgentTriggerDeliveryAttempt({ ...fence, now: START }); + await expect( + methods.expediteAgentTriggerDeliveries({ user, sourceIds: [background.id], now: START }), + ).resolves.toEqual({ expedited: 0, held: 1 }); + + const beforeDefer = Date.now(); + await expect( + methods.deferAgentTriggerDeliveryAttempt({ + ...fence, + attempt: attempt!, + availableAt: later, + }), + ).resolves.toBe('expedited'); + + const deferred = await Delivery.findById(held.delivery.id).lean(); + expect(deferred?.wakeRequestedAt).toBeUndefined(); + expect(deferred?.availableAt.getTime()).toBeGreaterThanOrEqual(beforeDefer); + expect(deferred?.availableAt).not.toEqual(later); + + const reclaimed = await methods.claimNextAgentTriggerDelivery({ + ...capable, + claimToken: 'second-claim', + now: new Date(), + leaseUntil: new Date(Date.now() + 60_000), + }); + expect(reclaimed).toMatchObject({ id: held.delivery.id }); + const secondFence = { ...fence, claimToken: 'second-claim' }; + const secondAttempt = await methods.beginAgentTriggerDeliveryAttempt({ + ...secondFence, + now: new Date(), + }); + await expect( + methods.deferAgentTriggerDeliveryAttempt({ + ...secondFence, + attempt: secondAttempt!, + availableAt: later, + }), + ).resolves.toBe(true); + expect((await Delivery.findById(held.delivery.id).lean())?.availableAt).toEqual(later); + }); + + it.each(['ordinary', 'legacy', 'shielded'] as const)( + 'honors a held wake marker on ordering release for %s leases', + async (profile) => { + const row = await waiting({ + availableAt: START, + ...(profile === 'ordinary' && { requiredWorkerCapability: undefined }), + }); + if (profile === 'legacy') { + await Delivery.updateOne( + { _id: row.delivery.id }, + { + $set: { status: 'capability_pending', availableAt: START }, + $unset: { capabilityStatus: 1, leaseUntil: 1 }, + }, + ); + } + const claimed = await methods.claimNextAgentTriggerDelivery(capable); + expect(claimed?.id).toBe(row.delivery.id); + const fence = { + id: row.delivery.id, + workerId: capable.workerId, + claimToken: capable.claimToken, + }; + await methods.expediteAgentTriggerDeliveries({ + deliveryKeys: [row.delivery.deliveryKey], + sourceIds: [background.id], + now: START, + }); + await expect( + methods.releaseAgentTriggerDelivery({ + ...fence, + claimToken: 'stale', + availableAt: later, + }), + ).resolves.toBe(false); + expect((await Delivery.findById(row.delivery.id).lean())?.wakeRequestedAt).toEqual(START); + await expect( + methods.releaseAgentTriggerDelivery({ ...fence, availableAt: later }), + ).resolves.toBe(true); + const released = await Delivery.findById(row.delivery.id).lean(); + expect(released?.wakeRequestedAt).toBeUndefined(); + expect(released?.availableAt.getTime()).toBeLessThanOrEqual(Date.now()); + expect(released?.claimAvailableAt).toEqual(released?.availableAt); + expect(released?.attempts).toBe(0); + const reclaimed = await methods.claimNextAgentTriggerDelivery({ + ...capable, + now: new Date(), + claimToken: 'next', + }); + expect(reclaimed?.id).toBe(row.delivery.id); + await methods.releaseAgentTriggerDelivery({ + ...fence, + claimToken: 'next', + availableAt: later, + }); + expect((await Delivery.findById(row.delivery.id).lean())?.availableAt).toEqual(later); + }, + ); + + it.each(['release', 'defer'] as const)( + 'does not lose an expedite between the unmarked and marked %s writes', + async (operation) => { + const row = await waiting({ availableAt: START }); + await methods.claimNextAgentTriggerDelivery(capable); + const fence = { + id: row.delivery.id, + workerId: capable.workerId, + claimToken: capable.claimToken, + }; + const attempt = + operation === 'defer' + ? await methods.beginAgentTriggerDeliveryAttempt({ ...fence, now: START }) + : undefined; + const updateOne = Delivery.collection.updateOne.bind(Delivery.collection); + let injected = false; + const spy = jest + .spyOn(Delivery.collection, 'updateOne') + .mockImplementation(async (filter, update, options) => { + if ( + !injected && + filter.capabilityClaimToken === capable.claimToken && + (filter.wakeRequestedAt as { $exists?: boolean } | undefined)?.$exists === false + ) { + injected = true; + await methods.expediteAgentTriggerDeliveries({ + deliveryKeys: [row.delivery.deliveryKey], + sourceIds: [background.id], + now: START, + }); + } + return updateOne(filter, update, options); + }); + try { + const input = { ...fence, availableAt: later }; + if (operation === 'defer') { + await expect( + methods.deferAgentTriggerDeliveryAttempt({ ...input, attempt: attempt! }), + ).resolves.toBe('expedited'); + } else { + await expect(methods.releaseAgentTriggerDelivery(input)).resolves.toBe(true); + } + } finally { + spy.mockRestore(); + } + expect(injected).toBe(true); + const released = await Delivery.findById(row.delivery.id).lean(); + expect(released?.wakeRequestedAt).toBeUndefined(); + expect(released?.availableAt.getTime()).toBeLessThanOrEqual(Date.now()); + expect(released?.attempts).toBe(0); + }, + ); + + it('leaves unfinished siblings backed off when only named child tasks settle', async () => { + const user = new mongoose.Types.ObjectId(); + const source = { id: 'subagent-completion', type: 'internal' }; + const task = (taskId: string) => + waiting({ + user, + envelope: { + event: { source, payload: { taskId } }, + target: { conversationId: 'parent' }, + }, + }); + const original = await task('original'); + const recovered = await task('recovered'); + const sibling = await task('sibling'); + await expect( + methods.expediteAgentTriggerDeliveries({ + user, + conversationId: 'parent', + taskIds: ['original', 'recovered'], + sourceIds: [source.id], + now: START, + }), + ).resolves.toEqual({ expedited: 2, held: 0 }); + expect((await Delivery.findById(original.delivery.id).lean())?.availableAt).toEqual(START); + expect((await Delivery.findById(recovered.delivery.id).lean())?.availableAt).toEqual(START); + expect((await Delivery.findById(sibling.delivery.id).lean())?.availableAt).toEqual(later); + await expect( + methods.expediteAgentTriggerDeliveries({ + user, + conversationId: 'parent', + taskIds: [], + sourceIds: [source.id], + now: START, + }), + ).rejects.toThrow(TypeError); + }); + + it("narrows a principal's selection to the conversation being resumed", async () => { + const user = new mongoose.Types.ObjectId(); + const target = await waiting({ + user, + envelope: { event: { source: background }, target: { conversationId: 'settled-convo' } }, + }); + const other = await waiting({ + user, + envelope: { event: { source: background }, target: { conversationId: 'other-convo' } }, + }); + + await expect( + methods.expediteAgentTriggerDeliveries({ + user, + conversationId: 'settled-convo', + sourceIds: [background.id], + now: START, + }), + ).resolves.toEqual({ expedited: 1, held: 0 }); + + expect((await Delivery.findById(target.delivery.id).lean())?.availableAt).toEqual(START); + expect((await Delivery.findById(other.delivery.id).lean())?.availableAt).toEqual(later); + }); + + it('refuses an unbounded or malformed selection', async () => { + await expect( + methods.expediteAgentTriggerDeliveries({ sourceIds: [background.id], now: START }), + ).rejects.toThrow(TypeError); + await expect( + methods.expediteAgentTriggerDeliveries({ + user: new mongoose.Types.ObjectId(), + sourceIds: [], + now: START, + }), + ).rejects.toThrow(TypeError); + await expect( + methods.expediteAgentTriggerDeliveries({ + deliveryKeys: ['trigger_x'], + sourceIds: [background.id], + now: new Date(Number.NaN), + }), + ).rejects.toThrow(TypeError); + }); + }); + + describe('listPendingAgentBackgroundToolCompletions', () => { + const background = { id: 'background-tool-completion', type: 'internal' }; + const completion = ( + user: mongoose.Types.ObjectId, + taskId: string, + overrides: Partial[0]> = {}, + ) => + methods.enqueueAgentTriggerDelivery( + enqueueInput({ + user, + orderingKey: `background-lane-${taskId}`, + envelope: { + event: { + source: background, + payload: { taskId, toolCallId: `call-${taskId}`, toolName: 'slow_task' }, + }, + target: { conversationId: 'conversation-1' }, + }, + requiredWorkerCapability: + AGENT_TRIGGER_WORKER_CAPABILITY_BACKGROUND_COMPLETION_RECEIPT_V2, + ...overrides, + }), + ); + + it('lists what is still going to arrive, running or settled, without result content', async () => { + const user = new mongoose.Types.ObjectId(); + const running = await completion(user, 'task-running'); + const settled = await completion(user, 'task-settled'); + await methods.persistAgentBackgroundToolResult({ + deliveryKey: settled.delivery.deliveryKey, + sourceId: background.id, + result: { status: 'completed', output: 'secret output', settledAt: START }, + }); + + const pending = await methods.listPendingAgentBackgroundToolCompletions({ + user, + conversationId: 'conversation-1', + sourceId: background.id, + }); + + expect(pending.truncated).toBe(false); + expect(pending.completions).toEqual([ + { + deliveryKey: running.delivery.deliveryKey, + taskId: 'task-running', + toolCallId: 'call-task-running', + toolName: 'slow_task', + dispatchedAt: expect.any(Date), + claimedByWakeup: false, + }, + { + deliveryKey: settled.delivery.deliveryKey, + taskId: 'task-settled', + toolCallId: 'call-task-settled', + toolName: 'slow_task', + dispatchedAt: expect.any(Date), + result: { status: 'completed', settledAt: START }, + claimedByWakeup: false, + }, + ]); + expect(JSON.stringify(pending)).not.toContain('secret output'); + }); + + it('excludes delivered rows and everything outside the conversation, user, and source', async () => { + const user = new mongoose.Types.ObjectId(); + const delivered = await completion(user, 'task-delivered'); + await Delivery.updateOne({ _id: delivered.delivery.id }, { $set: { status: 'succeeded' } }); + await completion(new mongoose.Types.ObjectId(), 'task-other-user'); + await completion(user, 'task-other-conversation', { + envelope: { + event: { + source: background, + payload: { + taskId: 'task-other-conversation', + toolCallId: 'call', + toolName: 'slow_task', + }, + }, + target: { conversationId: 'conversation-2' }, + }, + }); + await completion(user, 'task-other-source', { + envelope: { + event: { + source: { id: 'agent-queued-turn', type: 'internal' }, + payload: { taskId: 'task-other-source', toolCallId: 'call', toolName: 'slow_task' }, + }, + target: { conversationId: 'conversation-1' }, + }, + }); + const waiting = await completion(user, 'task-waiting'); + + const pending = await methods.listPendingAgentBackgroundToolCompletions({ + user, + conversationId: 'conversation-1', + sourceId: background.id, + }); + + expect(pending.completions.map((entry) => entry.deliveryKey)).toEqual([ + waiting.delivery.deliveryKey, + ]); + }); + + it('looks up one task and reports a truncated listing', async () => { + const user = new mongoose.Types.ObjectId(); + await completion(user, 'task-first'); + const second = await completion(user, 'task-second'); + + const one = await methods.listPendingAgentBackgroundToolCompletions({ + user, + conversationId: 'conversation-1', + sourceId: background.id, + taskId: 'task-second', + }); + expect(one).toEqual({ + completions: [expect.objectContaining({ deliveryKey: second.delivery.deliveryKey })], + dead: [], + truncated: false, + }); + + const page = await methods.listPendingAgentBackgroundToolCompletions({ + user, + conversationId: 'conversation-1', + sourceId: background.id, + limit: 1, + }); + expect(page.completions.map((entry) => entry.taskId)).toEqual(['task-first']); + expect(page.truncated).toBe(true); + }); + + it('omits legacy rows whose results live only on the parent message', async () => { + const user = new mongoose.Types.ObjectId(); + await completion(user, 'task-legacy', { + requiredWorkerCapability: AGENT_TRIGGER_WORKER_CAPABILITY_BACKGROUND_COMPLETION_V1, + }); + + const pending = await methods.listPendingAgentBackgroundToolCompletions({ + user, + conversationId: 'conversation-1', + sourceId: background.id, + }); + + expect(pending.completions).toEqual([]); + }); + + it('reports dead-lettered tasks apart from pending ones', async () => { + const user = new mongoose.Types.ObjectId(); + const dead = await completion(user, 'task-dead'); + await Delivery.updateOne( + { _id: dead.delivery.id }, + { $set: { status: 'leased', capabilityStatus: 'dead' } }, + ); + const deadLetter = await completion(user, 'task-dead-letter'); + await Delivery.updateOne({ _id: deadLetter.delivery.id }, { $set: { status: 'dead' } }); + + const pending = await methods.listPendingAgentBackgroundToolCompletions({ + user, + conversationId: 'conversation-1', + sourceId: background.id, + }); + + expect(pending.completions).toEqual([]); + expect(pending.dead.map(({ taskId }) => taskId).sort()).toEqual([ + 'task-dead', + 'task-dead-letter', + ]); + expect(pending.dead[0]).toEqual( + expect.objectContaining({ toolName: 'slow_task', dispatchedAt: expect.any(Date) }), + ); + }); + + it("lists a conversation's undelivered task ids for one source", async () => { + const user = new mongoose.Types.ObjectId(); + const subagent = { id: 'subagent-completion', type: 'internal' }; + await completion(user, 'child-waiting', { + envelope: { + event: { source: subagent, payload: { taskId: 'child-waiting' } }, + target: { conversationId: 'conversation-1' }, + }, + }); + const delivered = await completion(user, 'child-delivered', { + envelope: { + event: { source: subagent, payload: { taskId: 'child-delivered' } }, + target: { conversationId: 'conversation-1' }, + }, + }); + await Delivery.updateOne({ _id: delivered.delivery.id }, { $set: { status: 'succeeded' } }); + + await expect( + methods.listUndeliveredAgentTriggerTaskIds({ + user, + conversationId: 'conversation-1', + sourceId: subagent.id, + }), + ).resolves.toEqual({ taskIds: ['child-waiting'], truncated: false }); + }); + + it('distinguishes retiring a completion from finding it already delivered', async () => { + const user = new mongoose.Types.ObjectId(); + const delivered = await completion(user, 'task-already-delivered'); + await Delivery.updateOne({ _id: delivered.delivery.id }, { $set: { status: 'succeeded' } }); + const retire = (requireTransition?: true) => + methods.retireAgentTriggerDelivery({ + deliveryKey: delivered.delivery.deliveryKey, + sourceId: background.id, + settledAt: START, + reason: 'background result discarded by its owner', + onlyIfUnclaimed: true, + ...(requireTransition != null && { requireTransition }), + }); + + await expect(retire()).resolves.toBe(true); + await expect(retire(true)).resolves.toBe(false); + + const waiting = await completion(user, 'task-still-waiting'); + await expect( + methods.retireAgentTriggerDelivery({ + deliveryKey: waiting.delivery.deliveryKey, + sourceId: background.id, + settledAt: START, + reason: 'background result discarded by its owner', + onlyIfUnclaimed: true, + requireTransition: true, + }), + ).resolves.toBe(true); + }); + + it('refuses a malformed lookup', async () => { + await expect( + methods.listPendingAgentBackgroundToolCompletions({ + user: new mongoose.Types.ObjectId(), + conversationId: '', + sourceId: background.id, + }), + ).rejects.toThrow(TypeError); + await expect( + methods.listPendingAgentBackgroundToolCompletions({ + user: new mongoose.Types.ObjectId(), + conversationId: 'conversation-1', + sourceId: background.id, + limit: 0, + }), + ).rejects.toThrow(TypeError); + }); + }); + it('persists one private background result receipt independently of message rows', async () => { const source = { id: 'background-tool-completion', type: 'internal' }; const queued = await methods.enqueueAgentTriggerDelivery( @@ -1744,6 +2332,74 @@ describe('agent trigger delivery methods', () => { }); }); + it.each(['complete', 'dead', 'cleanup'] as const)( + 'signals maintenance after %s finalization fails without undoing root settlement', + async (mode) => { + const user = new mongoose.Types.ObjectId(); + const coalesceUntil = new Date(Date.now() + 60_000); + const shared = { + user, + orderingKey: 'inline-recovery', + coalesceKey: 'inline-batch', + coalesceFrom: new Date(coalesceUntil.getTime() - 750), + coalesceUntil, + availableAt: coalesceUntil, + envelopeBytes: 128, + }; + const root = await methods.enqueueAgentTriggerDelivery(enqueueInput(shared)); + if (mode !== 'cleanup') await methods.enqueueAgentTriggerDelivery(enqueueInput(shared)); + const claimed = await methods.claimNextAgentTriggerDelivery({ + workerId: 'worker-1', + claimToken: 'inline-finalization-claim', + now: coalesceUntil, + leaseUntil: new Date(coalesceUntil.getTime() + 60_000), + }); + expect(claimed?.id).toBe(root.delivery.id); + const input = { + id: claimed!.id, + workerId: 'worker-1', + claimToken: claimed!.claimToken!, + attempt: 1, + settledAt: coalesceUntil, + }; + const failure = + mode === 'cleanup' + ? jest + .spyOn(LaneSequence, 'updateOne') + .mockRejectedValueOnce(new Error('cleanup interrupted')) + : jest + .spyOn(Delivery, 'updateMany') + .mockRejectedValueOnce(new Error('batch interrupted')); + const recovery = { required: false }; + try { + const settled = + mode === 'dead' + ? await methods.deadLetterAgentTriggerDelivery( + { ...input, error: transientFailure({ retryable: false }) }, + recovery, + ) + : await methods.completeAgentTriggerDelivery( + { ...input, result: { accepted: true } }, + recovery, + ); + expect(settled).toBe(true); + expect(recovery.required).toBe(true); + } finally { + failure.mockRestore(); + } + expect((await Delivery.findById(root.delivery.id).lean())?.status).toBe( + mode === 'dead' ? 'dead' : 'succeeded', + ); + await methods.recoverAgentTriggerBatchReceipts(); + await methods.reclaimInactiveAgentTriggerLanes(); + const rows = await Delivery.find({ orderingKey: shared.orderingKey }).lean(); + expect(rows.every((row) => row.status === (mode === 'dead' ? 'dead' : 'succeeded'))).toBe( + true, + ); + if (mode !== 'dead') expect(await LaneSequence.findById(shared.orderingKey)).toBeNull(); + }, + ); + it('recovers batch receipts and lane cleanup after root settlement was interrupted', async () => { const user = new mongoose.Types.ObjectId(); const coalesceUntil = new Date(Date.now() + 60_000); @@ -2244,6 +2900,9 @@ describe('agent trigger delivery methods', () => { }); it('leaves staging unpublished while its durable user purge marker exists', async () => { + const emptyActivity = { found: false }; + await expect(methods.recoverAgentTriggerLanePublications(1, emptyActivity)).resolves.toBe(0); + expect(emptyActivity.found).toBe(false); const user = new mongoose.Types.ObjectId(); const orderingKey = 'purge-fenced-staging'; await UserPurge.create({ _id: user, fenceStartedAt: START, tenantId: 'tenant-1' }); @@ -2256,7 +2915,9 @@ describe('agent trigger delivery methods', () => { stagingRecoveryAt: START, }); - await expect(methods.recoverAgentTriggerLanePublications(1)).resolves.toBe(0); + const activity = { found: false }; + await expect(methods.recoverAgentTriggerLanePublications(1, activity)).resolves.toBe(0); + expect(activity.found).toBe(true); await expect(Delivery.findById(staged._id).lean()).resolves.toMatchObject({ status: 'staging', laneSequence: 0, @@ -2395,6 +3056,18 @@ describe('agent trigger delivery methods', () => { ); }); + it("indexes a user's waiting deliveries for settle-time expediting", async () => { + const deliveryIndexes = await Delivery.collection.indexes(); + const userIndex = deliveryIndexes.find( + (index) => + JSON.stringify(index.key) === JSON.stringify({ user: 1, status: 1, availableAt: 1 }), + ); + expect(userIndex).toBeDefined(); + /** A sparse index would skip every row the expedite and listing reads need. */ + expect(userIndex?.sparse).toBeUndefined(); + expect(userIndex?.partialFilterExpression).toBeUndefined(); + }); + it('publishes an idempotent replay on its persisted ordering lane', async () => { const user = new mongoose.Types.ObjectId(); const input = enqueueInput({ user, orderingKey: 'original-lane' }); diff --git a/packages/data-schemas/src/methods/triggerDelivery.ts b/packages/data-schemas/src/methods/triggerDelivery.ts index 0137a6dd43d..441977f8186 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.ts @@ -44,6 +44,24 @@ export const CLAIM_CAS_MAX_ATTEMPTS = 16; /** Candidates fetched per claim read; losing claimers advance through the * batch instead of re-reading the same head-of-queue row. */ const CLAIM_CANDIDATE_BATCH = 8; +/** Matches the per-conversation background task limit, so a listing is complete + * unless durable rows outlived that limit across restarts; it then says so. */ +const MAX_PENDING_BACKGROUND_COMPLETIONS = 200; +/** Every status before a delivery settles, i.e. whose result has not reached its conversation. */ +const DEAD_STATUSES: IAgentTriggerDelivery['status'][] = ['dead', 'capability_dead']; +/** Dead to every worker version, including a capability row a legacy worker still sees leased. */ +function isDeadDelivery(row: Pick): boolean { + return DEAD_STATUSES.includes(row.status) || row.capabilityStatus === 'dead'; +} +const UNDELIVERED_STATUSES: IAgentTriggerDelivery['status'][] = [ + 'staging', + 'capability_staging', + 'batched', + 'pending', + 'capability_pending', + 'leased', + 'capability_leased', +]; /** Capability work is inert to legacy claimers while preserving their lane * behavior: publishing is `staging`; queued work is `leased` without a lease * owner/deadline; execution adds a private lease; dead work is terminal. */ @@ -246,6 +264,53 @@ export interface AgentEventActorReceiptStorageMetrics { deadDeliveries: number; } +/** A background tool completion that has not reached its conversation yet. */ +export interface PendingAgentBackgroundToolCompletion { + deliveryKey: string; + taskId: string; + toolCallId: string; + toolName: string; + dispatchedAt: Date; + /** The tool's terminal outcome once it settled; absent while it still runs. */ + result?: { status: AgentBackgroundToolResultReceipt['status']; settledAt: Date }; + /** An automatic delivery holds the result and is starting its turn. */ + claimedByWakeup: boolean; +} + +export interface PendingAgentBackgroundToolCompletions { + completions: PendingAgentBackgroundToolCompletion[]; + /** Completions whose delivery dead-lettered: never delivered, recoverable only by a poll. */ + dead: PendingAgentBackgroundToolCompletion[]; + /** More undelivered completions exist than were returned. */ + truncated: boolean; +} + +export interface UndeliveredAgentTriggerTaskIds { + taskIds: string[]; + truncated: boolean; +} + +/** Selects deferred internal deliveries whose readiness condition just changed. */ +export interface ExpediteAgentTriggerDeliveriesInput { + sourceIds: readonly string[]; + /** Exact deliveries, e.g. the one whose result just became durable. */ + deliveryKeys?: readonly string[]; + /** Every matching delivery of one principal, e.g. after one of its generations settled. */ + user?: string | Types.ObjectId; + /** Narrows a principal's selection to deliveries that resume this conversation. */ + conversationId?: string; + /** Exact tasks within a principal's conversation, including a repaired attempt's predecessor. */ + taskIds?: readonly string[]; + now: Date; +} + +export interface ExpediteAgentTriggerDeliveriesResult { + /** Deferred, unheld deliveries moved to `now`. */ + expedited: number; + /** Deliveries a worker held, marked so their next deferral re-checks at once. */ + held: number; +} + export interface AgentTriggerDeliveryMethods { ensureAgentTriggerDeliveryIndexes: () => Promise; enqueueAgentTriggerDelivery: ( @@ -270,9 +335,10 @@ export interface AgentTriggerDeliveryMethods { beginAgentTriggerDeliveryAttempt: ( input: AgentTriggerDeliveryFence & { now: Date }, ) => Promise; + /** `expedited` when readiness changed while the delivery was held: it is due now instead. */ deferAgentTriggerDeliveryAttempt: ( input: AgentTriggerDeliveryFence & { attempt: number; availableAt: Date }, - ) => Promise; + ) => Promise; completeAgentTriggerDelivery: ( input: AgentTriggerDeliveryFence & { attempt: number; @@ -281,18 +347,25 @@ export interface AgentTriggerDeliveryMethods { handling?: AgentTriggerHandlingState; awaitTerminalHandling?: true; }, + recovery?: { required: boolean }, + ) => Promise; + retireAgentTriggerDelivery: ( + input: { + deliveryKey: string; + sourceId: string; + settledAt: Date; + reason: string; + onlyIfUnclaimed?: boolean; + onlyIfDead?: boolean; + /** Accept transport success without a terminal handling receipt, unless the + * delivery explicitly keeps its lane open for terminal handling. */ + allowSucceeded?: boolean; + /** True only when this call retired the delivery, not when it had already + * succeeded, e.g. delivered by a resolver that won the race. */ + requireTransition?: boolean; + }, + recovery?: { required: boolean }, ) => Promise; - retireAgentTriggerDelivery: (input: { - deliveryKey: string; - sourceId: string; - settledAt: Date; - reason: string; - onlyIfUnclaimed?: boolean; - onlyIfDead?: boolean; - /** Accept transport success without a terminal handling receipt, unless the - * delivery explicitly keeps its lane open for terminal handling. */ - allowSucceeded?: boolean; - }) => Promise; renewAgentTriggerDeliveryProducerLease: (input: { deliveryKey: string; sourceId: string; @@ -303,6 +376,23 @@ export interface AgentTriggerDeliveryMethods { sourceId: string; now: Date; }) => Promise; + listPendingAgentBackgroundToolCompletions: (input: { + user: string | Types.ObjectId; + conversationId: string; + sourceId: string; + /** One task's completion, e.g. to discard it. */ + taskId?: string; + limit?: number; + }) => Promise; + /** Task ids of one conversation's undelivered internal deliveries from one source. */ + listUndeliveredAgentTriggerTaskIds: (input: { + user: string | Types.ObjectId; + conversationId: string; + sourceId: string; + }) => Promise; + expediteAgentTriggerDeliveries: ( + input: ExpediteAgentTriggerDeliveriesInput, + ) => Promise; persistAgentBackgroundToolResult: ( input: PersistAgentBackgroundToolResultInput, ) => Promise; @@ -385,6 +475,7 @@ export interface AgentTriggerDeliveryMethods { settledAt: Date; receiptRetryAt?: Date; }, + recovery?: { required: boolean }, ) => Promise; getAgentTriggerDelivery: (deliveryKey: string) => Promise; getAgentTriggerDeliveryStatus: ( @@ -402,9 +493,18 @@ export interface AgentTriggerDeliveryMethods { user: string | Types.ObjectId, now: Date, ) => Promise; - recoverAgentTriggerLanePublications: (limit?: number) => Promise; - recoverAgentTriggerBatchReceipts: (limit?: number) => Promise; - reclaimInactiveAgentTriggerLanes: (limit?: number) => Promise; + recoverAgentTriggerLanePublications: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; + recoverAgentTriggerBatchReceipts: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; + reclaimInactiveAgentTriggerLanes: ( + limit?: number, + activity?: { found: boolean }, + ) => Promise; prepareAgentTriggerUserPurge: ( user: string | Types.ObjectId, fenceStartedAt: Date, @@ -414,7 +514,7 @@ export interface AgentTriggerDeliveryMethods { user: string | Types.ObjectId, fenceStartedAt: Date, ) => Promise; - recoverAgentTriggerUserPurges: (limit?: number) => Promise; + recoverAgentTriggerUserPurges: (limit?: number, activity?: { found: boolean }) => Promise; deleteAgentTriggerDeliveriesByUser: (user: string | Types.ObjectId) => Promise; eraseAgentTriggerDeliveryConversationResults: ( user: string | Types.ObjectId, @@ -1092,6 +1192,7 @@ export function createAgentTriggerDeliveryMethods( /** Repairs abandoned reservations and staging rows left by crashed writers. */ async function recoverAgentTriggerLanePublications( limit = DEFAULT_PURGE_RECOVERY_LIMIT, + activity?: { found: boolean }, ): Promise { if (!Number.isSafeInteger(limit) || limit <= 0) { throw new TypeError('Agent trigger lane recovery limit must be a positive integer'); @@ -1105,6 +1206,7 @@ export function createAgentTriggerDeliveryMethods( .sort({ publisherStartedAt: 1, _id: 1 }) .limit(boundedLimit) .lean(); + if (activity != null && lanes.length > 0) activity.found = true; let recovered = 0; const recoveryCursor = new Date(); for (const lane of lanes) { @@ -1175,6 +1277,9 @@ export function createAgentTriggerDeliveryMethods( .limit(remaining) .lean() : []; + if (activity != null && (legacyStaged.length > 0 || indexedStaged.length > 0)) { + activity.found = true; + } const staged = [ ...legacyStaged.map((delivery) => ({ ...delivery, @@ -1293,6 +1398,7 @@ export function createAgentTriggerDeliveryMethods( /** Bounds high-cardinality ordering metadata after the final retained job settles. */ async function reclaimInactiveAgentTriggerLanes( limit = DEFAULT_PURGE_RECOVERY_LIMIT, + activity?: { found: boolean }, ): Promise { if (!Number.isSafeInteger(limit) || limit <= 0) { throw new TypeError('Agent trigger lane reclamation limit must be a positive integer'); @@ -1304,6 +1410,7 @@ export function createAgentTriggerDeliveryMethods( .limit(boundedLimit) .select('_id orderingKey laneCleanupPendingAt') .lean>>(); + if (activity != null && pendingCleanup.length > 0) activity.found = true; let reclaimed = 0; for (const delivery of pendingCleanup) { if (await fulfillLaneCleanupRequest(delivery)) { @@ -1321,6 +1428,7 @@ export function createAgentTriggerDeliveryMethods( .limit(remaining) .select('_id') .lean>>(); + if (activity != null && lanes.length > 0) activity.found = true; for (const lane of lanes) { if (await reclaimLaneIfInactive(lane._id)) { reclaimed += 1; @@ -1760,6 +1868,7 @@ export function createAgentTriggerDeliveryMethods( async function recoverAgentTriggerBatchReceipts( limit = DEFAULT_PURGE_RECOVERY_LIMIT, + activity?: { found: boolean }, ): Promise { if (!Number.isSafeInteger(limit) || limit <= 0) { throw new TypeError('Agent trigger batch recovery limit must be a positive integer'); @@ -1773,6 +1882,7 @@ export function createAgentTriggerDeliveryMethods( .sort({ settledAt: 1, _id: 1 }) .limit(Math.min(limit, MAX_PURGE_RECOVERY_LIMIT)) .lean(); + if (activity != null && roots.length > 0) activity.found = true; let recovered = 0; for (const root of roots) { if (root.settledAt == null || (root.status !== 'succeeded' && root.status !== 'dead')) { @@ -1822,57 +1932,61 @@ export function createAgentTriggerDeliveryMethods( $or: [ordinaryFence(input), legacyCapabilityFence(input), shieldCapabilityFence(input)], }); + /** Readiness signals and the lease release must meet in one fenced write. Try + * the unmarked state first, then the marked state: a concurrent expedite can + * only add a marker while this lease is held. After release, expedite's second + * update sees an unheld row. No pipeline updates or unfenced follow-up needed. */ + async function releaseWaitingDelivery( + input: AgentTriggerDeliveryFence & { availableAt: Date; attempt?: number }, + ): Promise { + const attemptFence = input.attempt == null ? {} : { attempts: input.attempt }; + const attemptChange = input.attempt == null ? {} : { $inc: { attempts: -1 } }; + const profiles = [ + { filter: shieldCapabilityFence(input), status: 'leased', capabilityStatus: 'pending' }, + { filter: legacyCapabilityFence(input), status: 'capability_pending' }, + { filter: ordinaryFence(input), status: 'pending' }, + ] as const; + for (const marked of [false, true]) { + for (const profile of profiles) { + const availableAt = marked ? new Date() : input.availableAt; + const result = await Delivery().updateOne( + { + _id: input.id, + ...profile.filter, + ...attemptFence, + wakeRequestedAt: { $exists: marked }, + }, + { + ...attemptChange, + $set: { + status: profile.status, + ...('capabilityStatus' in profile && { capabilityStatus: profile.capabilityStatus }), + availableAt, + claimAvailableAt: availableAt, + }, + $unset: { + leaseBy: 1, + leaseUntil: 1, + claimToken: 1, + capabilityLeaseBy: 1, + capabilityLeaseUntil: 1, + capabilityClaimToken: 1, + wakeRequestedAt: 1, + }, + }, + ); + if (result.modifiedCount === 1) { + return marked ? 'expedited' : true; + } + } + } + return false; + } + async function releaseAgentTriggerDelivery( input: AgentTriggerDeliveryFence & { availableAt: Date }, ): Promise { - const shieldResult = await Delivery().updateOne( - { _id: input.id, ...shieldCapabilityFence(input) }, - { - $set: { - status: 'leased', - availableAt: input.availableAt, - capabilityStatus: 'pending', - claimAvailableAt: input.availableAt, - }, - $unset: { - leaseBy: 1, - leaseUntil: 1, - claimToken: 1, - capabilityLeaseBy: 1, - capabilityLeaseUntil: 1, - capabilityClaimToken: 1, - }, - }, - ); - if (shieldResult.modifiedCount === 1) { - return true; - } - const capabilityResult = await Delivery().updateOne( - { _id: input.id, ...legacyCapabilityFence(input) }, - { - $set: { - status: 'capability_pending', - availableAt: input.availableAt, - claimAvailableAt: input.availableAt, - }, - $unset: { leaseBy: 1, leaseUntil: 1, claimToken: 1 }, - }, - ); - if (capabilityResult.modifiedCount === 1) { - return true; - } - const result = await Delivery().updateOne( - { _id: input.id, ...ordinaryFence(input) }, - { - $set: { - status: 'pending', - availableAt: input.availableAt, - claimAvailableAt: input.availableAt, - }, - $unset: { leaseBy: 1, leaseUntil: 1, claimToken: 1 }, - }, - ); - return result.modifiedCount === 1; + return (await releaseWaitingDelivery(input)) !== false; } async function beginAgentTriggerDeliveryAttempt( @@ -1902,64 +2016,11 @@ export function createAgentTriggerDeliveryMethods( /** Releases a pre-dispatch deferral and restores the attempt consumed by beginAttempt. */ async function deferAgentTriggerDeliveryAttempt( input: AgentTriggerDeliveryFence & { attempt: number; availableAt: Date }, - ): Promise { + ): Promise { if (!Number.isSafeInteger(input.attempt) || input.attempt <= 0) { throw new TypeError('attempt must be a positive integer'); } - const update = { - $inc: { attempts: -1 }, - $set: { - availableAt: input.availableAt, - claimAvailableAt: input.availableAt, - }, - $unset: { leaseBy: 1, leaseUntil: 1, claimToken: 1 }, - }; - const shieldResult = await Delivery().updateOne( - { - _id: input.id, - ...shieldCapabilityFence(input), - attempts: input.attempt, - }, - { - $inc: update.$inc, - $set: { - status: 'leased', - availableAt: input.availableAt, - capabilityStatus: 'pending', - claimAvailableAt: input.availableAt, - }, - $unset: { - leaseBy: 1, - leaseUntil: 1, - claimToken: 1, - capabilityLeaseBy: 1, - capabilityLeaseUntil: 1, - capabilityClaimToken: 1, - }, - }, - ); - if (shieldResult.modifiedCount === 1) { - return true; - } - const capabilityResult = await Delivery().updateOne( - { - _id: input.id, - ...legacyCapabilityFence(input), - attempts: input.attempt, - }, - { - ...update, - $set: { ...update.$set, status: 'capability_pending' }, - }, - ); - if (capabilityResult.modifiedCount === 1) { - return true; - } - const result = await Delivery().updateOne( - { _id: input.id, ...ordinaryFence(input), attempts: input.attempt }, - { ...update, $set: { ...update.$set, status: 'pending' } }, - ); - return result.modifiedCount === 1; + return releaseWaitingDelivery(input); } async function completeAgentTriggerDelivery( @@ -1970,6 +2031,7 @@ export function createAgentTriggerDeliveryMethods( handling?: AgentTriggerHandlingState; awaitTerminalHandling?: true; }, + recovery?: { required: boolean }, ): Promise { const awaitsTerminalHandling = input.awaitTerminalHandling === true && input.handling?.status === 'started'; @@ -2050,6 +2112,7 @@ export function createAgentTriggerDeliveryMethods( }); await fulfillLaneCleanupRequest(completed); } catch (error) { + if (recovery != null) recovery.required = true; // Root success is authoritative. Maintenance retries both constituent // receipt settlement and the existing durable lane-cleanup marker. logger.warn('[agent-triggers] failed to finalize a completed trigger batch', { @@ -2064,17 +2127,23 @@ export function createAgentTriggerDeliveryMethods( * that the result will never become dispatchable. This transition is keyed * by the immutable delivery identity rather than a worker lease so the * producer can unblock the lane even while a resolver is deferring it. */ - async function retireAgentTriggerDelivery(input: { - deliveryKey: string; - sourceId: string; - settledAt: Date; - reason: string; - onlyIfUnclaimed?: boolean; - onlyIfDead?: boolean; - /** Accept transport success without a terminal handling receipt, unless the - * delivery explicitly keeps its lane open for terminal handling. */ - allowSucceeded?: boolean; - }): Promise { + async function retireAgentTriggerDelivery( + input: { + deliveryKey: string; + sourceId: string; + settledAt: Date; + reason: string; + onlyIfUnclaimed?: boolean; + onlyIfDead?: boolean; + /** Accept transport success without a terminal handling receipt, unless the + * delivery explicitly keeps its lane open for terminal handling. */ + allowSucceeded?: boolean; + /** True only when this call retired the delivery, not when it had already + * succeeded, e.g. delivered by a resolver that won the race. */ + requireTransition?: boolean; + }, + recovery?: { required: boolean }, + ): Promise { if ( input.deliveryKey.length === 0 || input.deliveryKey.length > 256 || @@ -2151,6 +2220,7 @@ export function createAgentTriggerDeliveryMethods( try { await fulfillLaneCleanupRequest(retired); } catch (error) { + if (recovery != null) recovery.required = true; logger.warn('[agent-triggers] failed to finalize a retired internal delivery', { deliveryKey: input.deliveryKey, error: error instanceof Error ? error.message : String(error), @@ -2158,6 +2228,9 @@ export function createAgentTriggerDeliveryMethods( } return true; } + if (input.requireTransition === true) { + return false; + } return ( (await Delivery().exists({ deliveryKey: input.deliveryKey, @@ -2257,6 +2330,202 @@ export function createAgentTriggerDeliveryMethods( : { status: 'expired', leaseUntil: delivery.producerLeaseUntil }; } + /** Every background completion of one conversation that has not been + * delivered yet: still running, or settled and waiting for a wake-up. The + * durable delivery row outlives the process-local task registry (another + * replica, a restart, or the registry's retention), so it is the record of + * what is still going to arrive. Result content is never returned here. */ + async function listPendingAgentBackgroundToolCompletions(input: { + user: string | Types.ObjectId; + conversationId: string; + sourceId: string; + taskId?: string; + limit?: number; + }): Promise { + const limit = Math.min( + input.limit ?? MAX_PENDING_BACKGROUND_COMPLETIONS, + MAX_PENDING_BACKGROUND_COMPLETIONS, + ); + if ( + input.conversationId.length === 0 || + input.conversationId.length > 256 || + input.sourceId.length === 0 || + input.sourceId.length > 256 || + (input.taskId != null && (input.taskId.length === 0 || input.taskId.length > 256)) || + !Number.isSafeInteger(limit) || + limit <= 0 + ) { + throw new TypeError('Invalid pending background completion lookup'); + } + const rows = await Delivery() + .find({ + user: input.user, + 'envelope.event.source.type': 'internal', + 'envelope.event.source.id': input.sourceId, + 'envelope.target.conversationId': input.conversationId, + ...(input.taskId != null && { 'envelope.event.payload.taskId': input.taskId }), + /** Legacy rows keep results only on the parent message, so a missing + * receipt cannot tell running from finished; they drain on their own path. */ + requiredWorkerCapability: AGENT_TRIGGER_WORKER_CAPABILITY_BACKGROUND_COMPLETION_RECEIPT_V2, + /** Dead letters are read too, so a caller can tell "delivered" from "failed". */ + status: { $in: [...UNDELIVERED_STATUSES, ...DEAD_STATUSES] }, + }) + .select( + 'deliveryKey createdAt status capabilityStatus envelope.event.payload ' + + 'backgroundToolResult.status backgroundToolResult.settledAt backgroundToolResult.resultClaim', + ) + .sort({ createdAt: 1, _id: 1 }) + .limit(limit + 1) + .lean< + Array< + Pick< + IAgentTriggerDelivery, + 'deliveryKey' | 'createdAt' | 'backgroundToolResult' | 'status' | 'capabilityStatus' + > & { + envelope?: { event?: { payload?: Record } }; + } + > + >(); + const dead: PendingAgentBackgroundToolCompletion[] = []; + const completions = rows.slice(0, limit).flatMap((row) => { + const payload = row.envelope?.event?.payload; + const taskId = payload?.taskId; + const toolCallId = payload?.toolCallId; + const toolName = payload?.toolName; + if ( + row.createdAt == null || + typeof taskId !== 'string' || + typeof toolCallId !== 'string' || + typeof toolName !== 'string' + ) { + return []; + } + const receipt = row.backgroundToolResult; + return [ + { + deliveryKey: row.deliveryKey, + taskId, + toolCallId, + toolName, + dispatchedAt: row.createdAt, + ...(receipt != null && { + result: { status: receipt.status, settledAt: receipt.settledAt }, + }), + claimedByWakeup: receipt?.resultClaim != null, + }, + ].filter((completion) => { + if (!isDeadDelivery(row)) { + return true; + } + dead.push(completion); + return false; + }); + }); + return { completions, dead, truncated: rows.length > limit }; + } + + async function listUndeliveredAgentTriggerTaskIds(input: { + user: string | Types.ObjectId; + conversationId: string; + sourceId: string; + }): Promise { + if ( + input.conversationId.length === 0 || + input.conversationId.length > 256 || + input.sourceId.length === 0 || + input.sourceId.length > 256 + ) { + throw new TypeError('Invalid undelivered task lookup'); + } + const rows = await Delivery() + .find({ + user: input.user, + 'envelope.event.source.type': 'internal', + 'envelope.event.source.id': input.sourceId, + 'envelope.target.conversationId': input.conversationId, + status: { $in: UNDELIVERED_STATUSES }, + capabilityStatus: { $ne: 'dead' }, + }) + .select('envelope.event.payload.taskId') + .sort({ createdAt: 1, _id: 1 }) + .limit(MAX_PENDING_BACKGROUND_COMPLETIONS + 1) + .lean>(); + const taskIds = rows + .slice(0, MAX_PENDING_BACKGROUND_COMPLETIONS) + .map((row) => row.envelope?.event?.payload?.taskId) + .filter((taskId): taskId is string => typeof taskId === 'string'); + return { taskIds, truncated: rows.length > MAX_PENDING_BACKGROUND_COMPLETIONS }; + } + + /** Pulls deferred deliveries back to `now` when the condition they were + * waiting on has changed, so a waiting delivery can back off without delaying + * the moment it becomes deliverable. Unclaimed rows move now; held rows retain + * a signal consumed atomically by readiness deferral or ordering release. */ + async function expediteAgentTriggerDeliveries( + input: ExpediteAgentTriggerDeliveriesInput, + ): Promise { + const deliveryKeys = input.deliveryKeys ?? []; + if ( + (input.taskIds != null && + (input.user == null || + input.conversationId == null || + input.taskIds.length === 0 || + input.taskIds.some((id) => id.length === 0 || id.length > 256))) || + input.sourceIds.length === 0 || + input.sourceIds.some((id) => id.length === 0 || id.length > 256) || + deliveryKeys.some((key) => key.length === 0 || key.length > 256) || + (deliveryKeys.length === 0 && input.user == null) || + (input.conversationId != null && + (input.conversationId.length === 0 || input.conversationId.length > 256)) || + !(input.now instanceof Date) || + !Number.isFinite(input.now.getTime()) + ) { + throw new TypeError('Invalid agent trigger delivery expedite'); + } + const selection = { + 'envelope.event.source.type': 'internal', + 'envelope.event.source.id': { $in: [...input.sourceIds] }, + ...(deliveryKeys.length > 0 && { deliveryKey: { $in: [...deliveryKeys] } }), + ...(input.taskIds != null && { + 'envelope.event.payload.taskId': { $in: [...input.taskIds] }, + }), + ...(input.user != null && { user: input.user }), + ...(input.conversationId != null && { + 'envelope.target.conversationId': input.conversationId, + }), + }; + /** Classic operators only: aggregation-pipeline updates are not portable. A + * held row may be deferred on readiness its worker read before the change, + * so it keeps a marker that its deferral honors instead of moving now. The + * marker is written first: a row released after it was read as held is then + * seen unheld by the move, so no release between the two escapes both. */ + const held = await Delivery().updateMany( + { + ...selection, + status: { $in: ['leased', 'capability_leased'] }, + $or: [{ leaseBy: { $exists: true } }, { capabilityLeaseBy: { $exists: true } }], + }, + { $set: { wakeRequestedAt: input.now } }, + ); + const moved = await Delivery().updateMany( + { + ...selection, + availableAt: { $gt: input.now }, + leaseBy: { $exists: false }, + $or: [ + { status: { $in: ['pending', 'capability_pending'] } }, + { + status: 'leased', + capabilityStatus: 'pending', + capabilityLeaseBy: { $exists: false }, + }, + ], + }, + { $set: { availableAt: input.now, claimAvailableAt: input.now } }, + ); + return { expedited: moved.modifiedCount, held: held.matchedCount }; + } + /** Stores terminal output on the pre-admitted delivery before attempting the * parent-message projection. The first terminal receipt wins; exact retries * are idempotent and conflicting rewrites fail closed. */ @@ -3509,6 +3778,7 @@ export function createAgentTriggerDeliveryMethods( settledAt: Date; receiptRetryAt?: Date; }, + recovery?: { required: boolean }, ): Promise { const error = normalizeFailure(input.error); if ( @@ -3603,6 +3873,7 @@ export function createAgentTriggerDeliveryMethods( error, }); } catch (settlementError) { + if (recovery != null) recovery.required = true; logger.warn('[agent-triggers] failed to settle batch dead-letter receipts', { deliveryId: String(dead._id), error: settlementError instanceof Error ? settlementError.message : String(settlementError), @@ -3881,6 +4152,7 @@ export function createAgentTriggerDeliveryMethods( /** Recovers cleanup markers whose users are gone; active-user markers are never destructive. */ async function recoverAgentTriggerUserPurges( limit = DEFAULT_PURGE_RECOVERY_LIMIT, + activity?: { found: boolean }, ): Promise { if (!Number.isSafeInteger(limit) || limit <= 0) { throw new TypeError('Agent trigger purge recovery limit must be a positive integer'); @@ -3893,6 +4165,7 @@ export function createAgentTriggerDeliveryMethods( .sort({ backgroundToolResultDeletionPendingAt: 1, _id: 1 }) .limit(Math.min(limit, MAX_PURGE_RECOVERY_LIMIT)) .lean(); + if (activity != null && pending.length > 0) activity.found = true; for (const row of pending) { const conversationId = (row.envelope as { target?: { conversationId?: string } }).target ?.conversationId; @@ -3915,6 +4188,7 @@ export function createAgentTriggerDeliveryMethods( .sort({ updatedAt: 1, _id: 1 }) .limit(Math.min(limit, MAX_PURGE_RECOVERY_LIMIT)) .lean(); + if (activity != null && markers.length > 0) activity.found = true; let recovered = 0; for (const marker of markers) { const user = await mongoose.models.User.findById(marker._id) @@ -4014,6 +4288,9 @@ export function createAgentTriggerDeliveryMethods( retireAgentTriggerDelivery, renewAgentTriggerDeliveryProducerLease, getAgentTriggerDeliveryProducerLease, + listPendingAgentBackgroundToolCompletions, + listUndeliveredAgentTriggerTaskIds, + expediteAgentTriggerDeliveries, persistAgentBackgroundToolResult, getAgentBackgroundToolResult, getAgentBackgroundToolResultClaim, diff --git a/packages/data-schemas/src/schema/triggerDelivery.ts b/packages/data-schemas/src/schema/triggerDelivery.ts index 1de097809d1..d9b503e6c59 100644 --- a/packages/data-schemas/src/schema/triggerDelivery.ts +++ b/packages/data-schemas/src/schema/triggerDelivery.ts @@ -218,6 +218,7 @@ const triggerDeliverySchema: Schema = new Schema( requeueCount: { type: Number, default: 0, min: 0 }, stagingRecoveryAt: { type: Date }, laneCleanupPendingAt: { type: Date }, + wakeRequestedAt: { type: Date }, }, { timestamps: true }, ); @@ -253,6 +254,8 @@ triggerDeliverySchema.index( { sparse: true }, ); triggerDeliverySchema.index({ status: 1, updatedAt: -1 }); +/** One user's waiting deliveries, read when that user's generation settles. */ +triggerDeliverySchema.index({ user: 1, status: 1, availableAt: 1 }); triggerDeliverySchema.index({ 'actorReceipt.resolution': 1 }, { sparse: true }); triggerDeliverySchema.index({ user: 1, actorActionAdmittedAt: 1 }, { sparse: true }); triggerDeliverySchema.index({ stagingRecoveryAt: 1 }, { sparse: true }); diff --git a/packages/data-schemas/src/types/memory.ts b/packages/data-schemas/src/types/memory.ts index a8ccce8fe47..89a3123cb28 100644 --- a/packages/data-schemas/src/types/memory.ts +++ b/packages/data-schemas/src/types/memory.ts @@ -72,9 +72,20 @@ export interface SetMemoryByIdResult extends MemoryResult { memory?: IMemoryEntryLean; } -export interface FormattedMemoriesResult { - withKeys: string; - withoutKeys: string; - totalTokens?: number; - tokenCountsByKey?: Map; -} +/** A failed read is not an empty memory partition. Consumers can suppress + * memory guidance and writes without changing the shape of successful reads. */ +export type FormattedMemoriesResult = + | { + withKeys: string; + withoutKeys: string; + totalTokens?: number; + tokenCountsByKey?: Map; + readFailed?: false; + } + | { + withKeys: undefined; + withoutKeys: undefined; + totalTokens?: number; + tokenCountsByKey?: Map; + readFailed: true; + }; diff --git a/packages/data-schemas/src/types/triggerDelivery.ts b/packages/data-schemas/src/types/triggerDelivery.ts index be019d65a22..5865d882679 100644 --- a/packages/data-schemas/src/types/triggerDelivery.ts +++ b/packages/data-schemas/src/types/triggerDelivery.ts @@ -171,6 +171,8 @@ export interface IAgentTriggerDelivery { stagingRecoveryAt?: Date; /** Durable proof that successful settlement still owes lane cleanup publication. */ laneCleanupPendingAt?: Date; + /** Readiness changed while a worker held this delivery; its next deferral re-checks at once. */ + wakeRequestedAt?: Date; createdAt?: Date; updatedAt?: Date; }