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/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/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/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, From 56301c1b5a32f37e5c59e09f0c8567f559ee0b76 Mon Sep 17 00:00:00 2001 From: "lia-by-librechat[bot]" <328778573+lia-by-librechat[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 10:10:42 -0400 Subject: [PATCH 22/73] =?UTF-8?q?=F0=9F=AA=9A=20fix:=20Count=20Large=20Pro?= =?UTF-8?q?mpts=20in=20Bounded=20Tokenizer=20Chunks=20(#16305)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tokenize budgeting inputs in bounded chunks instead of charging a token per UTF-8 byte above 4 KiB. Subdivide long whitespace/non-whitespace runs to bound BPE merge work and preserve surrogate pairs. Keep strict provider-accounting counts and unavailable-tokenizer fallbacks separate. Cover both encodings, SDK slice composition, programmatic tool prompts, summarization, real overflow, cache reuse, and failure recovery. Related to #16278. Co-authored-by: Lia --- packages/api/src/agents/client.spec.ts | 128 ++++++++++++++++++++- packages/api/src/utils/tokenizer.spec.ts | 139 ++++++++++++++++++++--- packages/api/src/utils/tokenizer.ts | 58 ++++++---- 3 files changed, 284 insertions(+), 41 deletions(-) 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/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. */ From edddf92eadf157aa275ba736fd5170839f763053 Mon Sep 17 00:00:00 2001 From: "lia-by-librechat[bot]" <328778573+lia-by-librechat[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:10:44 -0400 Subject: [PATCH 23/73] =?UTF-8?q?=F0=9F=AA=80=20fix:=20Preserve=20Coding?= =?UTF-8?q?=20Approval=20Mode=20on=20Background=20Returns=20(#16322)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Lia --- .../agents/backgroundCompletionWakeup.spec.ts | 74 ++++++++++++++++++- .../src/agents/backgroundCompletionWakeup.ts | 3 + .../agents/subagentCompletionWakeup.spec.ts | 30 +++++++- .../src/agents/subagentCompletionWakeup.ts | 1 + packages/api/src/agents/triggers/host.spec.ts | 31 ++++++++ packages/api/src/agents/triggers/host.ts | 8 +- 6 files changed, 142 insertions(+), 5 deletions(-) diff --git a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts index eef75fc6777..1e87d6a2f55 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts @@ -1,5 +1,6 @@ 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, @@ -45,12 +46,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', @@ -440,6 +444,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({ diff --git a/packages/api/src/agents/backgroundCompletionWakeup.ts b/packages/api/src/agents/backgroundCompletionWakeup.ts index e107118bffb..bfdb9f21d3b 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.ts @@ -386,6 +386,7 @@ export function createBackgroundToolCompletionWakeupResolver({ return { status: 'ready', parentMessageId, + ...(parent.codeApprovalMode != null && { codeApprovalMode: parent.codeApprovalMode }), input, releaseOnDefiniteFailure: async () => { const released = await methods.releaseBackgroundToolResultClaims({ @@ -455,6 +456,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 +492,7 @@ export function createBackgroundToolCompletionWakeupResolver({ return { status: 'ready', parentMessageId, + ...(parent.codeApprovalMode != null && { codeApprovalMode: parent.codeApprovalMode }), input: buildWakeupInput([ { ...registration, status: receipt.status, output: receipt.output }, ]), diff --git a/packages/api/src/agents/subagentCompletionWakeup.spec.ts b/packages/api/src/agents/subagentCompletionWakeup.spec.ts index d2fd1fc95eb..c8317a6906a 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'; @@ -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({ diff --git a/packages/api/src/agents/subagentCompletionWakeup.ts b/packages/api/src/agents/subagentCompletionWakeup.ts index 612e22c0c0a..02240dcfb9d 100644 --- a/packages/api/src/agents/subagentCompletionWakeup.ts +++ b/packages/api/src/agents/subagentCompletionWakeup.ts @@ -778,6 +778,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({ 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, }), From 3788fe05d92efe8b06315991a65cb9464d67cc8f Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 24 Sep 2026 21:36:10 -0400 Subject: [PATCH 24/73] =?UTF-8?q?=F0=9F=90=A2=20fix:=20Back=20Off=20Waitin?= =?UTF-8?q?g=20Completion=20Wake-ups=20and=20Keep=20Their=20Trace=20Out=20?= =?UTF-8?q?of=20Requests=20(#16349)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/services/Agents/triggers.js | 2 + .../agents/backgroundCompletionWakeup.spec.ts | 82 +++++++++++++++++++ .../src/agents/backgroundCompletionWakeup.ts | 16 +++- .../agents/subagentCompletionWakeup.spec.ts | 53 ++++++++++++ .../src/agents/subagentCompletionWakeup.ts | 16 +++- packages/api/src/agents/triggers/README.md | 5 ++ .../api/src/agents/triggers/backoff.spec.ts | 28 +++++++ packages/api/src/agents/triggers/backoff.ts | 26 ++++++ .../api/src/agents/triggers/engine.spec.ts | 75 +++++++++++++++++ packages/api/src/agents/triggers/engine.ts | 81 ++++++++++++------ .../agents/triggers/service.delivery.spec.ts | 12 +++ packages/api/src/agents/triggers/service.ts | 8 ++ packages/data-provider/src/config.spec.ts | 3 + packages/data-provider/src/config.ts | 10 +++ .../src/methods/triggerDelivery.spec.ts | 12 +++ .../src/schema/triggerDelivery.ts | 2 + 16 files changed, 402 insertions(+), 29 deletions(-) create mode 100644 packages/api/src/agents/triggers/backoff.spec.ts create mode 100644 packages/api/src/agents/triggers/backoff.ts diff --git a/api/server/services/Agents/triggers.js b/api/server/services/Agents/triggers.js index 7b58b54a9e8..329bfd46356 100644 --- a/api/server/services/Agents/triggers.js +++ b/api/server/services/Agents/triggers.js @@ -24,11 +24,13 @@ const getGenerationAdmissionEvidence = (userId, clientRequestId, streamId, conve const subagentCompletionAdapter = createSubagentCompletionWakeupResolver({ methods, getGenerationJob: (conversationId) => GenerationJobManager.getJob(conversationId), + getWaitMaxIntervalMs: () => service.getCompletionWaitMaxIntervalMs(), }); const backgroundToolCompletionAdapter = createBackgroundToolCompletionWakeupResolver({ methods, getGenerationJob: (conversationId) => GenerationJobManager.getJob(conversationId), getResultBatchSize: () => service.getBackgroundCompletionResultBatchSize(), + getWaitMaxIntervalMs: () => service.getCompletionWaitMaxIntervalMs(), }); const eventActorAdapter = createAgentEventContinueResolver({ methods, diff --git a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts index 1e87d6a2f55..ceb95fe5409 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts @@ -756,6 +756,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: [] }); diff --git a/packages/api/src/agents/backgroundCompletionWakeup.ts b/packages/api/src/agents/backgroundCompletionWakeup.ts index bfdb9f21d3b..1ca6fca26c9 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.ts @@ -21,6 +21,7 @@ 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 { createAgentTriggerEnvelope } from './triggers/envelope'; import { AgentTriggerExecutionError } from './triggers/host'; @@ -94,6 +95,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 +155,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 +261,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 +300,7 @@ export function createBackgroundToolCompletionWakeupResolver({ code: 'PARENT_NOT_READY', retryable: true, status: 409, - retryAfter: '1', + retryAfter: isParentWorking(parentJob) ? waitingRetry(envelope.receivedAt) : '1', deferWithoutAttempt: true, }); } @@ -526,7 +538,7 @@ export function createBackgroundToolCompletionWakeupResolver({ code: 'BACKGROUND_TOOL_RESULT_NOT_READY', retryable: true, status: 409, - retryAfter: '1', + retryAfter: waitingRetry(envelope.receivedAt), deferWithoutAttempt: true, }); }; diff --git a/packages/api/src/agents/subagentCompletionWakeup.spec.ts b/packages/api/src/agents/subagentCompletionWakeup.spec.ts index c8317a6906a..16c54710c9b 100644 --- a/packages/api/src/agents/subagentCompletionWakeup.spec.ts +++ b/packages/api/src/agents/subagentCompletionWakeup.spec.ts @@ -1287,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 02240dcfb9d..93285187bd1 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, }); } diff --git a/packages/api/src/agents/triggers/README.md b/packages/api/src/agents/triggers/README.md index faef699ef5d..875d797218a 100644 --- a/packages/api/src/agents/triggers/README.md +++ b/packages/api/src/agents/triggers/README.md @@ -99,6 +99,11 @@ experimental clustered startup: | `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, so a long wait costs about one re-check per cap instead of one every five seconds. 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. 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/engine.spec.ts b/packages/api/src/agents/triggers/engine.spec.ts index e00a600654b..c5532179a60 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,80 @@ 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('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..ac5bf350e2e 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'; @@ -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/service.delivery.spec.ts b/packages/api/src/agents/triggers/service.delivery.spec.ts index 8d571b0f92f..1130b7f6050 100644 --- a/packages/api/src/agents/triggers/service.delivery.spec.ts +++ b/packages/api/src/agents/triggers/service.delivery.spec.ts @@ -209,6 +209,18 @@ describe('durable agent trigger service', () => { 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({ diff --git a/packages/api/src/agents/triggers/service.ts b/packages/api/src/agents/triggers/service.ts index a37e1821f35..f208336e576 100644 --- a/packages/api/src/agents/triggers/service.ts +++ b/packages/api/src/agents/triggers/service.ts @@ -33,6 +33,7 @@ import { selfOriginFromAddress } from '../../app/origin'; import { createAgentTriggerExecutionHost } from './host'; import { parseAgentTriggerEnvelope } from './envelope'; import { createIdleRecoveryLoop } from '../recovery'; +import { WAITING_RETRY_CAP_MS } from './backoff'; export const AGENT_TRIGGER_TOKEN_TTL = '60s'; const DEFAULT_USER_DRAIN_TIMEOUT_MS = 35_000; @@ -48,6 +49,7 @@ export interface AgentTriggerServiceOptions { queuedTurnMaxIntervalMs?: number; maintenanceMaxIntervalMs?: number; deliveryMaxIntervalMs?: number; + completionWaitMaxIntervalMs?: number; }; } @@ -206,6 +208,8 @@ export interface AgentTriggerService { input: Parameters[0], ) => ReturnType; getBackgroundCompletionResultBatchSize: () => number; + /** Longest a waiting completion delivery re-checks readiness. */ + getCompletionWaitMaxIntervalMs: () => number; releaseBackgroundToolResultClaims: AgentTriggerDeliveryMethods['releaseAgentBackgroundToolResultClaims']; drainUser: (userId: string) => Promise; prepareUserPurge: (userId: string, fenceStartedAt: Date, tenantId?: string) => Promise; @@ -323,6 +327,7 @@ 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; @@ -553,6 +558,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(); @@ -728,6 +735,7 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A return getClaim == null ? null : getClaim(input); }), getBackgroundCompletionResultBatchSize: () => backgroundCompletionResultBatchSize, + getCompletionWaitMaxIntervalMs: () => completionWaitMaxIntervalMs, releaseBackgroundToolResultClaims: (input) => runAsSystem(async () => { const release = requireMethods().releaseAgentBackgroundToolResultClaims; diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index 02676ab8ad4..7ff73ec7a80 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -682,12 +682,15 @@ describe('agent event runtime config', () => { 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({ diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 698a5a6406c..3acf59cd6b6 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1537,6 +1537,16 @@ export const agentsEndpointSchema = baseEndpointSchema .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(), }) diff --git a/packages/data-schemas/src/methods/triggerDelivery.spec.ts b/packages/data-schemas/src/methods/triggerDelivery.spec.ts index c2b61109308..3c6850d7b38 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.spec.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.spec.ts @@ -2468,6 +2468,18 @@ describe('agent trigger delivery methods', () => { ); }); + it("indexes a user's waiting deliveries for user-scoped readiness reads", 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/schema/triggerDelivery.ts b/packages/data-schemas/src/schema/triggerDelivery.ts index 1de097809d1..2d08f6d14c3 100644 --- a/packages/data-schemas/src/schema/triggerDelivery.ts +++ b/packages/data-schemas/src/schema/triggerDelivery.ts @@ -253,6 +253,8 @@ triggerDeliverySchema.index( { sparse: true }, ); triggerDeliverySchema.index({ status: 1, updatedAt: -1 }); +/** One user's waiting deliveries, read when their readiness changes. */ +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 }); From 402c19ce32884f78be9c528fb732060f7a6efe8c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 24 Sep 2026 21:50:50 -0400 Subject: [PATCH 25/73] =?UTF-8?q?=F0=9F=A7=AA=20ci:=20Count=20System=20Ins?= =?UTF-8?q?tructions=20in=20Mock=20Model=20Usage=20(#16350)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The e2e fake model reported input tokens for the chat messages alone, because the SDK hands a test override model the pruned messages without the systemRunnable pipe. Once #16301 gave every memory-enabled chat a system prompt, the calibrated context snapshot put used tokens below the instructions and the gauge dropped its Messages row. Count input over the complete prompt, as a real provider bills it. --- e2e/setup/fake-model.js | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) 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, From 7238be3e5ea1eceaa9cc59496625c8ea55eb60a4 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 24 Sep 2026 22:28:28 -0400 Subject: [PATCH 26/73] =?UTF-8?q?=E2=9A=A1=20fix:=20Deliver=20Waiting=20Co?= =?UTF-8?q?mpletion=20Wake-ups=20the=20Moment=20They=20Are=20Ready=20(#163?= =?UTF-8?q?39)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐢 fix: Back Off Waiting Completion Wake-ups and Deliver Them When Ready * fix: Signal every readiness path and make the completion wait cap configurable * fix: Mark held completion deliveries instead of pipeline expedite, and announce store-won approval expiry * fix: Scope settle expedites to the resumed conversation and close the remaining signal gaps * fix: Announce a won approval expiry once and release subagent wake-up registrations * 🧪 ci: Count System Instructions in Mock Model Usage (#16350) The e2e fake model reported input tokens for the chat messages alone, because the SDK hands a test override model the pruned messages without the systemRunnable pipe. Once #16301 gave every memory-enabled chat a system prompt, the calibrated context snapshot put used tokens below the instructions and the gauge dropped its Messages row. Count input over the complete prompt, as a real provider bills it. --------- Co-authored-by: Lia --- api/server/services/Agents/triggers.js | 2 + .../Endpoints/agents/backgroundCompletion.js | 2 + .../Endpoints/agents/subagentThreadStore.js | 4 +- .../agents/subagentThreadStore.spec.js | 11 + .../api/src/agents/backgroundCompletion.ts | 3 + .../agents/backgroundCompletionWakeup.spec.ts | 19 + .../src/agents/backgroundCompletionWakeup.ts | 2 + .../src/agents/handlers.background.spec.ts | 38 ++ packages/api/src/agents/handlers.ts | 2 + .../agents/subagentCompletionWakeup.spec.ts | 4 +- .../src/agents/subagentCompletionWakeup.ts | 5 +- .../api/src/agents/subagentThreads.spec.ts | 122 ++++++- packages/api/src/agents/subagentThreads.ts | 57 ++- packages/api/src/agents/triggers/README.md | 8 +- .../api/src/agents/triggers/engine.spec.ts | 86 +++++ packages/api/src/agents/triggers/engine.ts | 2 +- .../agents/triggers/service.delivery.spec.ts | 137 +++++++ packages/api/src/agents/triggers/service.ts | 71 +++- .../api/src/stream/GenerationJobManager.ts | 92 ++++- .../__tests__/generationSettled.spec.ts | 145 ++++++++ .../stream/__tests__/idempotencyClaim.spec.ts | 118 ++++-- packages/api/src/stream/index.ts | 2 + .../src/methods/triggerDelivery.spec.ts | 344 +++++++++++++++++- .../src/methods/triggerDelivery.ts | 254 +++++++------ .../src/schema/triggerDelivery.ts | 3 +- .../data-schemas/src/types/triggerDelivery.ts | 2 + 26 files changed, 1373 insertions(+), 162 deletions(-) create mode 100644 packages/api/src/stream/__tests__/generationSettled.spec.ts diff --git a/api/server/services/Agents/triggers.js b/api/server/services/Agents/triggers.js index 329bfd46356..f8944c6ea8d 100644 --- a/api/server/services/Agents/triggers.js +++ b/api/server/services/Agents/triggers.js @@ -53,6 +53,7 @@ service = createAgentTriggerService({ ), isPrincipalActive: methods.isAgentTriggerPrincipalActive, supportsDetachedActionCompletion: () => GenerationJobManager.supportsDetachedAgentEventActions, + subscribeGenerationSettled: (listener) => GenerationJobManager.onGenerationSettled(listener), settleSourceBeforeDeadLetter: queuedTurnLifecycle.settleBeforeDeadLetter, prepareContinue: createAgentContinuationResolver({ eventActor: eventActorAdapter, @@ -92,6 +93,7 @@ module.exports = { retireAgentTrigger: service.retire, renewAgentTriggerProducerLease: service.renewProducerLease, persistAgentBackgroundToolResult: service.persistBackgroundToolResult, + expediteCompletionWakeups: service.expediteCompletionWakeups, getAgentBackgroundToolResultClaim: service.getBackgroundToolResultClaim, releaseAgentBackgroundToolResultClaims: service.releaseBackgroundToolResultClaims, drainAgentTriggerDeliveriesForUser: service.drainUser, diff --git a/api/server/services/Endpoints/agents/backgroundCompletion.js b/api/server/services/Endpoints/agents/backgroundCompletion.js index 8ca1cae6efc..da93d982424 100644 --- a/api/server/services/Endpoints/agents/backgroundCompletion.js +++ b/api/server/services/Endpoints/agents/backgroundCompletion.js @@ -11,6 +11,7 @@ const { releaseAgentBackgroundToolResultClaims, renewAgentTriggerProducerLease, retireAgentTrigger, + expediteCompletionWakeups, } = require('../../Agents/triggers'); const preregisterBackgroundToolCompletion = createBackgroundToolCompletionWakeupHandler( @@ -19,6 +20,7 @@ const preregisterBackgroundToolCompletion = createBackgroundToolCompletionWakeup renewAgentTriggerProducerLease, (deliveryKey, sourceId, result) => persistAgentBackgroundToolResult({ deliveryKey, sourceId, result }), + (deliveryKey) => expediteCompletionWakeups({ deliveryKeys: [deliveryKey] }), ); function createBackgroundToolResultPersistence({ req, updateToolCallResult }) { diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.js b/api/server/services/Endpoints/agents/subagentThreadStore.js index b33b7df5415..d98f0f275f6 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -12,7 +12,7 @@ const { SubagentActivityStream, } = require('@librechat/api'); const db = require('~/models'); -const { enqueueAgentTrigger } = require('../../Agents/triggers'); +const { enqueueAgentTrigger, expediteCompletionWakeups } = require('../../Agents/triggers'); const GENERATION_DRAIN_TIMEOUT_MS = 45_000; const GENERATION_DRAIN_POLL_MS = 100; @@ -78,6 +78,8 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( releaseOwnerAdmission: db.releaseSubagentAdmission, cancelUnroutedTask: cancelUnroutedGeneration, onTaskPrepared: completionWakeupHandler, + onTaskSettled: (userId, conversationId, taskIds) => + expediteCompletionWakeups({ user: userId, conversationId, taskIds }), }, ); diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.spec.js b/api/server/services/Endpoints/agents/subagentThreadStore.spec.js index bae8a2cdebd..a7d26094680 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.spec.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.spec.js @@ -45,6 +45,7 @@ jest.mock('~/models', () => ({ jest.mock('../../Agents/triggers', () => ({ enqueueAgentTrigger: jest.fn(), + expediteCompletionWakeups: jest.fn(), })); const { @@ -80,6 +81,16 @@ describe('subagent thread Redis lifecycle', () => { expect(mockCompletionWakeupHandler).toHaveBeenCalledWith({ taskId: 'task-1' }); }); + it('expedites only the settled task identities in their parent conversation', () => { + const { expediteCompletionWakeups } = require('../../Agents/triggers'); + taskStoreOptions.onTaskSettled('user-1', 'parent-1', ['task-1', 'recovered-task']); + expect(expediteCompletionWakeups).toHaveBeenCalledWith({ + user: 'user-1', + conversationId: 'parent-1', + taskIds: ['task-1', 'recovered-task'], + }); + }); + it('registers local task-store quiescence independently of optional Redis setup', () => { expect(taskStoreShutdownRegistration).toEqual([ 'subagent task store', diff --git a/packages/api/src/agents/backgroundCompletion.ts b/packages/api/src/agents/backgroundCompletion.ts index 8ef1d74cca5..02d57cb7d8e 100644 --- a/packages/api/src/agents/backgroundCompletion.ts +++ b/packages/api/src/agents/backgroundCompletion.ts @@ -42,6 +42,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 diff --git a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts index ceb95fe5409..a30d5a4b499 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts @@ -112,6 +112,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, diff --git a/packages/api/src/agents/backgroundCompletionWakeup.ts b/packages/api/src/agents/backgroundCompletionWakeup.ts index 1ca6fca26c9..e5ef7ca6312 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.ts @@ -550,6 +550,7 @@ export function createBackgroundToolCompletionWakeupHandler( retire: RetireBackgroundToolCompletion, renewProducerLease: RenewBackgroundToolCompletionProducerLease, persistResult?: PersistBackgroundToolCompletionResult, + expedite?: (deliveryKey: string) => void, ): ( registration: BackgroundToolWakeupRegistration, ) => Promise { @@ -613,6 +614,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/handlers.background.spec.ts b/packages/api/src/agents/handlers.background.spec.ts index 645a473cf9e..37ad6535d64 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); diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index f991418273d..edbab15caa5 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -6130,6 +6130,8 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand 'background tool result was not persisted', 'definite', ); + } else if (deliveryReady && !durableReceiptReady) { + completionAdmission?.expedite?.(); } } catch (persistError) { if (!durableReceiptReady) { diff --git a/packages/api/src/agents/subagentCompletionWakeup.spec.ts b/packages/api/src/agents/subagentCompletionWakeup.spec.ts index 16c54710c9b..0f096275573 100644 --- a/packages/api/src/agents/subagentCompletionWakeup.spec.ts +++ b/packages/api/src/agents/subagentCompletionWakeup.spec.ts @@ -60,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]!; @@ -112,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(); }); diff --git a/packages/api/src/agents/subagentCompletionWakeup.ts b/packages/api/src/agents/subagentCompletionWakeup.ts index 93285187bd1..0b70b205a21 100644 --- a/packages/api/src/agents/subagentCompletionWakeup.ts +++ b/packages/api/src/agents/subagentCompletionWakeup.ts @@ -810,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({ @@ -850,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 875d797218a..5fe25e09aa1 100644 --- a/packages/api/src/agents/triggers/README.md +++ b/packages/api/src/agents/triggers/README.md @@ -103,7 +103,13 @@ experimental clustered startup: `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, so a long wait costs about one re-check per cap instead of one every five seconds. +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. diff --git a/packages/api/src/agents/triggers/engine.spec.ts b/packages/api/src/agents/triggers/engine.spec.ts index c5532179a60..697834d4f8c 100644 --- a/packages/api/src/agents/triggers/engine.spec.ts +++ b/packages/api/src/agents/triggers/engine.spec.ts @@ -123,6 +123,92 @@ describe('createAgentTriggerDeliveryEngine', () => { 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) }); diff --git a/packages/api/src/agents/triggers/engine.ts b/packages/api/src/agents/triggers/engine.ts index ac5bf350e2e..e6d6f9e8fce 100644 --- a/packages/api/src/agents/triggers/engine.ts +++ b/packages/api/src/agents/triggers/engine.ts @@ -158,7 +158,7 @@ export interface AgentTriggerDeliveryStore { claimToken: string; attempt: number; availableAt: Date; - }) => Promise; + }) => Promise; complete: (input: { id: string; workerId: string; diff --git a/packages/api/src/agents/triggers/service.delivery.spec.ts b/packages/api/src/agents/triggers/service.delivery.spec.ts index 1130b7f6050..0b1278141aa 100644 --- a/packages/api/src/agents/triggers/service.delivery.spec.ts +++ b/packages/api/src/agents/triggers/service.delivery.spec.ts @@ -209,6 +209,143 @@ 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 }; diff --git a/packages/api/src/agents/triggers/service.ts b/packages/api/src/agents/triggers/service.ts index f208336e576..fc2125bef36 100644 --- a/packages/api/src/agents/triggers/service.ts +++ b/packages/api/src/agents/triggers/service.ts @@ -26,7 +26,9 @@ 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'; @@ -35,6 +37,9 @@ 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; @@ -53,6 +58,11 @@ export interface AgentTriggerServiceOptions { }; } +/** 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']; @@ -69,6 +79,16 @@ export interface AgentTriggerServiceDeps { 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 { @@ -124,6 +144,7 @@ export interface AgentTriggerDeliveryPersistence { retireAgentTriggerDelivery: AgentTriggerDeliveryMethods['retireAgentTriggerDelivery']; renewAgentTriggerDeliveryProducerLease: AgentTriggerDeliveryMethods['renewAgentTriggerDeliveryProducerLease']; persistAgentBackgroundToolResult?: AgentTriggerDeliveryMethods['persistAgentBackgroundToolResult']; + expediteAgentTriggerDeliveries?: AgentTriggerDeliveryMethods['expediteAgentTriggerDeliveries']; getAgentBackgroundToolResultClaim?: AgentTriggerDeliveryMethods['getAgentBackgroundToolResultClaim']; releaseAgentBackgroundToolResultClaims?: AgentTriggerDeliveryMethods['releaseAgentBackgroundToolResultClaims']; retryAgentTriggerDelivery: AgentTriggerDeliveryStore['retry']; @@ -210,6 +231,8 @@ export interface AgentTriggerService { 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; @@ -539,9 +562,45 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A 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 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; + unsubscribeGenerationSettled?.(); + unsubscribeGenerationSettled = undefined; await purgeRecoveryLoop?.stop(); await initializePromise?.catch(() => undefined); await deliveryEngine?.stop(); @@ -611,6 +670,11 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A ); 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(() => { @@ -727,7 +791,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 () => { @@ -736,6 +804,7 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A }), getBackgroundCompletionResultBatchSize: () => backgroundCompletionResultBatchSize, getCompletionWaitMaxIntervalMs: () => completionWaitMaxIntervalMs, + expediteCompletionWakeups: (input) => expediteCompletions(input), releaseBackgroundToolResultClaims: (input) => runAsSystem(async () => { const release = requireMethods().releaseAgentBackgroundToolResultClaims; diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 991ee1b882c..a92318e29f2 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -568,10 +568,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 +798,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 +1091,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 +2270,7 @@ class GenerationJobManagerClass { streamId: string, job: Pick< SerializableJobData, - 'createdAt' | 'conversationId' | 'providerExecutionId' | 'agentEventDeliveryKey' + 'createdAt' | 'conversationId' | 'providerExecutionId' | 'agentEventDeliveryKey' | 'userId' >, message: string, ): Promise { @@ -2288,6 +2302,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 +2319,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 || @@ -3901,6 +3931,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 +4102,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 +4325,7 @@ class GenerationJobManagerClass { metricStatus = 'error'; } recordGenerationJob(this.storeLabel, metricStatus); + this.notifyGenerationSettled(claim); } if (cleanupError != null) { @@ -4616,6 +4682,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 +8932,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 +9133,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/__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/index.ts b/packages/api/src/stream/index.ts index b6c97e60988..ef72b19ee5f 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'; diff --git a/packages/data-schemas/src/methods/triggerDelivery.spec.ts b/packages/data-schemas/src/methods/triggerDelivery.spec.ts index 3c6850d7b38..7498b932c49 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.spec.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.spec.ts @@ -453,6 +453,348 @@ 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); + }); + }); + 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( @@ -2468,7 +2810,7 @@ describe('agent trigger delivery methods', () => { ); }); - it("indexes a user's waiting deliveries for user-scoped readiness reads", async () => { + it("indexes a user's waiting deliveries for settle-time expediting", async () => { const deliveryIndexes = await Delivery.collection.indexes(); const userIndex = deliveryIndexes.find( (index) => diff --git a/packages/data-schemas/src/methods/triggerDelivery.ts b/packages/data-schemas/src/methods/triggerDelivery.ts index e3fe4c2d0fa..1e2b6ead6fb 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.ts @@ -246,6 +246,27 @@ export interface AgentEventActorReceiptStorageMetrics { deadDeliveries: number; } +/** 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 +291,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; @@ -307,6 +329,9 @@ export interface AgentTriggerDeliveryMethods { sourceId: string; now: Date; }) => Promise; + expediteAgentTriggerDeliveries: ( + input: ExpediteAgentTriggerDeliveriesInput, + ) => Promise; persistAgentBackgroundToolResult: ( input: PersistAgentBackgroundToolResultInput, ) => Promise; @@ -1846,57 +1871,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( @@ -1926,64 +1955,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( @@ -2287,6 +2263,75 @@ export function createAgentTriggerDeliveryMethods( : { status: 'expired', leaseUntil: delivery.producerLeaseUntil }; } + /** 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. */ @@ -4049,6 +4094,7 @@ export function createAgentTriggerDeliveryMethods( retireAgentTriggerDelivery, renewAgentTriggerDeliveryProducerLease, getAgentTriggerDeliveryProducerLease, + expediteAgentTriggerDeliveries, persistAgentBackgroundToolResult, getAgentBackgroundToolResult, getAgentBackgroundToolResultClaim, diff --git a/packages/data-schemas/src/schema/triggerDelivery.ts b/packages/data-schemas/src/schema/triggerDelivery.ts index 2d08f6d14c3..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,7 +254,7 @@ triggerDeliverySchema.index( { sparse: true }, ); triggerDeliverySchema.index({ status: 1, updatedAt: -1 }); -/** One user's waiting deliveries, read when their readiness changes. */ +/** 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 }); 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; } From 5862b3b7d1b41f400f08cd021725c0ef39b1923b Mon Sep 17 00:00:00 2001 From: "lia-by-librechat[bot]" <328778573+lia-by-librechat[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 22:48:56 -0400 Subject: [PATCH 27/73] =?UTF-8?q?=F0=9F=9B=8D=EF=B8=8F=20fix:=20Discover?= =?UTF-8?q?=20Managed=20Agents=20in=20Marketplace=20Search=20(#16324)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🛍️ fix: Discover Managed Agents in Marketplace Search * 🛍️ fix: Reject Invalid Agent List Permissions Before Bypass * 🛍️ fix: Scope Managed Agent Discovery and Avatar Refresh * 🛍️ fix: Bound Managed Avatar Cache to Signed URL Lifetime * 🛍️ fix: Invalidate Tenant Avatar Cache After Upload --------- Co-authored-by: Lia --- api/server/controllers/agents/v1.js | 68 ++++--- api/server/controllers/agents/v1.spec.js | 156 +++++++++++++++- .../controllers/avatar.contentFilter.spec.js | 24 +++ packages/api/src/agents/index.ts | 2 + packages/api/src/agents/listing.ts | 52 ++++++ .../api/src/agents/listingAvatars.spec.ts | 172 ++++++++++++++++++ packages/api/src/agents/listingAvatars.ts | 145 +++++++++++++++ packages/api/src/agents/reads.spec.ts | 114 ++++++++++++ packages/api/src/agents/reads.ts | 5 +- .../data-schemas/src/methods/agent.spec.ts | 115 ++++++++++++ packages/data-schemas/src/methods/agent.ts | 19 +- 11 files changed, 829 insertions(+), 43 deletions(-) create mode 100644 packages/api/src/agents/listing.ts create mode 100644 packages/api/src/agents/listingAvatars.spec.ts create mode 100644 packages/api/src/agents/listingAvatars.ts diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index e975d867461..d6137f281b9 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -15,6 +15,11 @@ const { collectEdgeAgentIds, replaceEdgeSourceId, mergeDeploymentSkillIds, + getAgentListAccess, + isFullAgentListAvatarCacheEntry, + getAgentListAvatarRefreshKey, + refreshAgentListAvatarsBeforePage, + refreshManagedAgentListPageAvatars, mergeAgentOcrConversion, sanitizeModelParameters, MAX_AVATAR_REFRESH_AGENTS, @@ -84,6 +89,7 @@ const { resolveConfigServers, userCanUseMCPServers, } = require('~/server/services/MCP'); +const { hasCapability } = require('~/server/middleware/roles/capabilities'); const { attachOwnerContacts } = require('~/server/services/Agents/ownerContact'); const { getMCPServersRegistry } = require('~/config'); const { getLogStores } = require('~/cache'); @@ -1706,7 +1712,7 @@ const deleteAgentHandler = async (req, res) => { }; /** - * Lists agents using ACL-aware permissions (ownership + explicit shares). + * Lists agents using ACL permissions or the manage:agents capability. * @route GET /Agents * @param {object} req - Express Request * @param {object} req.query - Request query @@ -1727,12 +1733,6 @@ const getListAgentsHandler = async (req, res) => { requiredPermission = PermissionBits.VIEW; } const canReturnSkillConfig = hasEditBit(requiredPermission); - /** - * Derived from the same bit as `canReturnSkillConfig` but answering a different question: - * skill-config exposure versus edit-permission reporting. An EDIT-scoped request matches - * only editable agents, so it needs no second lookup to know which ones those are. - */ - const needsEditableLookup = !hasEditBit(requiredPermission); // Base filter const filter = {}; @@ -1756,7 +1756,7 @@ const getListAgentsHandler = async (req, res) => { } const cache = getLogStores(CacheKeys.S3_EXPIRY_INTERVAL); - const refreshKey = `${userId}:agents_avatar_refresh`; + const refreshKey = getAgentListAvatarRefreshKey(req.user); /** * These reads share no inputs, so they resolve together rather than chaining round @@ -1766,8 +1766,8 @@ const getListAgentsHandler = async (req, res) => { * * `editableIds` lets a VIEW-scoped response mark which agents the caller may also edit, * so consumers wanting just the editable subset can filter one shared VIEW fetch rather - * than issuing a second full paginated walk under an EDIT-scoped cache key. Requests - * that already ask for EDIT get it for free: everything they match is editable. + * than issuing a second full paginated walk under an EDIT-scoped cache key. Managers and + * EDIT-scoped requests need no separate edit lookup: everything they match is editable. * * `idOnTheSource` is forwarded so `getUserPrincipals` resolves identity without reading * the user document; the auth strategies already normalize it to a value or null. Each @@ -1775,19 +1775,12 @@ const getListAgentsHandler = async (req, res) => { */ const { idOnTheSource } = req.user; const [ - accessibleIds, + { accessibleIds, editableIds }, publiclyAccessibleIds, cachedRefreshEntry, accessibleSkillIds, - editableIds, ] = await Promise.all([ - findAccessibleResources({ - userId, - role: req.user.role, - idOnTheSource, - resourceType: ResourceType.AGENT, - requiredPermissions: requiredPermission, - }), + getAgentListAccess(req.user, requiredPermission, { hasCapability, findAccessibleResources }), findPubliclyAccessibleResources({ resourceType: ResourceType.AGENT, requiredPermissions: PermissionBits.VIEW, @@ -1802,21 +1795,9 @@ const getListAgentsHandler = async (req, res) => { resourceType: ResourceType.SKILL, requiredPermissions: PermissionBits.VIEW, }), - needsEditableLookup - ? findAccessibleResources({ - userId, - role: req.user.role, - idOnTheSource, - resourceType: ResourceType.AGENT, - requiredPermissions: PermissionBits.EDIT, - }) - : null, ]); - const isValidCachedRefresh = - cachedRefreshEntry != null && - typeof cachedRefreshEntry === 'object' && - cachedRefreshEntry.urlCache != null; + const isValidCachedRefresh = isFullAgentListAvatarCacheEntry(cachedRefreshEntry); /** * Refresh all S3 avatars for this user's accessible agent set (not only the current page) @@ -1866,11 +1847,16 @@ const getListAgentsHandler = async (req, res) => { } }; - const cachedRefresh = await resolveAvatarRefresh(); + const cachedRefreshBeforePage = await refreshAgentListAvatarsBeforePage( + accessibleIds, + cachedRefreshEntry, + resolveAvatarRefresh, + ); - // Use the new ACL-aware function + // Use the ACL-scoped or explicitly tenant-scoped list query. const data = await db.getListAgentsByAccess({ accessibleIds, + tenantId: req.user.tenantId ?? null, otherParams: filter, limit, after: cursor, @@ -1883,13 +1869,23 @@ const getListAgentsHandler = async (req, res) => { return res.json(data); } + const cachedRefresh = await refreshManagedAgentListPageAvatars({ + accessibleIds, + agents, + cachedEntry: cachedRefreshBeforePage, + refreshS3Url, + cacheSet: cache.set.bind(cache), + cacheKey: refreshKey, + ttl: Time.THIRTY_MINUTES, + }); + const accessibleSkillSet = canReturnSkillConfig ? null : new Set(mergeDeploymentSkillIds(accessibleSkillIds).map((oid) => oid.toString())); const publicSet = new Set(publiclyAccessibleIds.map((oid) => oid.toString())); /** Null for EDIT-scoped requests, where every matched agent is editable by definition. */ - const editableSet = editableIds ? new Set(editableIds.map((oid) => oid.toString())) : null; + const editableSet = editableIds ? new Set(editableIds) : null; const agentsWithContacts = await attachOwnerContacts(agents); const urlCache = cachedRefresh?.urlCache; @@ -2016,7 +2012,7 @@ const uploadAgentAvatarHandler = async (req, res) => { try { const avatarCache = getLogStores(CacheKeys.S3_EXPIRY_INTERVAL); - await avatarCache.delete(`${req.user.id}:agents_avatar_refresh`); + await avatarCache.delete(getAgentListAvatarRefreshKey(req.user)); } catch (cacheErr) { logger.error('[/:agent_id/avatar] Error invalidating avatar refresh cache', cacheErr); } diff --git a/api/server/controllers/agents/v1.spec.js b/api/server/controllers/agents/v1.spec.js index d96af401fc1..bf704940ba7 100644 --- a/api/server/controllers/agents/v1.spec.js +++ b/api/server/controllers/agents/v1.spec.js @@ -3,7 +3,7 @@ const express = require('express'); const request = require('supertest'); const { nanoid } = require('nanoid'); const { v4: uuidv4 } = require('uuid'); -const { createModels, tenantStorage } = require('@librechat/data-schemas'); +const { createModels, tenantStorage, SystemCapabilities } = require('@librechat/data-schemas'); const { Tools, SkillsScope, @@ -111,6 +111,7 @@ const { refreshS3Url, } = require('@librechat/api'); const { grantPermission } = require('~/server/services/PermissionService'); +const { hasCapability } = require('~/server/middleware/roles/capabilities'); const db = require('~/models'); /** @@ -3391,6 +3392,81 @@ describe('Agent Controllers - Mass Assignment Protection', () => { }); }); + test('lets a manage:agents role discover unshared agents by search', async () => { + await db.grantCapability({ + principalType: PrincipalType.ROLE, + principalId: 'LIST_MANAGER_TEST', + capability: SystemCapabilities.MANAGE_AGENTS, + }); + mockReq.user = { + id: userB.toString(), + role: 'LIST_MANAGER_TEST', + idOnTheSource: null, + }; + mockReq.query.search = 'A2'; + findAccessibleResources.mockResolvedValue([]); + findPubliclyAccessibleResources.mockResolvedValue([]); + + expect(await hasCapability(mockReq.user, SystemCapabilities.MANAGE_AGENTS)).toBe(true); + await getListAgentsHandler(mockReq, mockRes); + + const response = mockRes.json.mock.calls[0][0]; + expect(response.data.map((agent) => agent.id)).toEqual([agentA2.id]); + expect(response.data[0].isEditable).toBe(true); + expect(findAccessibleResources).toHaveBeenCalledWith( + expect.objectContaining({ resourceType: ResourceType.AGENT }), + ); + }); + + test('restricts a manager to the authenticated tenant even without ambient tenant context', async () => { + const tenantA = `tenant-a-${uuidv4()}`; + const tenantB = `tenant-b-${uuidv4()}`; + const name = 'Tenant-Scoped Discovery'; + const agentInA = await tenantStorage.run({ tenantId: tenantA }, () => + Agent.create({ + id: `agent_${nanoid(12)}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userA, + }), + ); + await tenantStorage.run({ tenantId: tenantB }, () => + Agent.create({ + id: `agent_${nanoid(12)}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userB, + }), + ); + await db.grantCapability({ + principalType: PrincipalType.ROLE, + principalId: 'LIST_TENANT_MANAGER', + capability: SystemCapabilities.MANAGE_AGENTS, + }); + mockReq.user = { + id: userB.toString(), + role: 'LIST_TENANT_MANAGER', + idOnTheSource: null, + tenantId: tenantA, + }; + mockReq.query.search = name; + findAccessibleResources.mockResolvedValue([]); + findPubliclyAccessibleResources.mockResolvedValue([]); + + await getListAgentsHandler(mockReq, mockRes); + expect(mockCache.get).toHaveBeenCalledWith( + `${userB.toString()}:${tenantA}:agents_avatar_refresh`, + ); + expect(mockRes.json.mock.calls[0][0].data.map((agent) => agent.id)).toEqual([agentInA.id]); + + mockRes.json.mockClear(); + mockReq.user.tenantId = undefined; + await getListAgentsHandler(mockReq, mockRes); + expect(mockRes.json.mock.calls[0][0].data).toHaveLength(0); + }); + test('should return empty list when user has no accessible agents', async () => { // User B has no permissions and no owned agents mockReq.user.id = userB.toString(); @@ -4111,6 +4187,68 @@ describe('Agent Controllers - Mass Assignment Protection', () => { }); }); + test('refreshes only manager search results without writing or losing later cursor pages', async () => { + const db = require('~/models'); + const name = 'Paged Manager Avatar'; + const firstAgent = await Agent.create({ + id: `agent_${nanoid(12)}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userA, + avatar: { source: FileSources.s3, filepath: 'first.jpg' }, + }); + const secondAgent = await Agent.create({ + id: `agent_${nanoid(12)}`, + name, + provider: 'openai', + model: 'gpt-4', + author: userA, + avatar: { source: FileSources.s3, filepath: 'second.jpg' }, + }); + await db.grantCapability({ + principalType: PrincipalType.ROLE, + principalId: 'LIST_AVATAR_MANAGER', + capability: SystemCapabilities.MANAGE_AGENTS, + }); + const mockReq = { + user: { id: userB.toString(), role: 'LIST_AVATAR_MANAGER', idOnTheSource: null }, + query: { search: name, limit: '1' }, + }; + const mockRes = { status: jest.fn().mockReturnThis(), json: jest.fn().mockReturnThis() }; + const listSpy = jest.spyOn(db, 'getListAgentsByAccess'); + const updateSpy = jest.spyOn(db, 'updateAgent'); + findAccessibleResources.mockResolvedValue([]); + findPubliclyAccessibleResources.mockResolvedValue([]); + refreshS3Url.mockImplementation(async ({ filepath }) => `signed:${filepath}`); + + try { + await getListAgentsHandler(mockReq, mockRes); + const pageOne = mockRes.json.mock.calls[0][0]; + expect(pageOne.data).toHaveLength(1); + expect(pageOne.after).toBeTruthy(); + expect(pageOne.data[0].avatar.filepath).toMatch(/^signed:/); + expect(listSpy).toHaveBeenCalledTimes(1); + expect(refreshS3Url).toHaveBeenCalledTimes(1); + expect(updateSpy).not.toHaveBeenCalled(); + + mockRes.json.mockClear(); + mockReq.query.cursor = pageOne.after; + await getListAgentsHandler(mockReq, mockRes); + const pageTwo = mockRes.json.mock.calls[0][0]; + expect(pageTwo.data).toHaveLength(1); + expect([pageOne.data[0].id, pageTwo.data[0].id].sort()).toEqual( + [firstAgent.id, secondAgent.id].sort(), + ); + expect(pageTwo.data[0].avatar.filepath).toMatch(/^signed:/); + expect(refreshS3Url).toHaveBeenCalledTimes(2); + expect(updateSpy).not.toHaveBeenCalled(); + } finally { + listSpy.mockRestore(); + updateSpy.mockRestore(); + } + }); + test('should skip avatar refresh if cache hit', async () => { mockCache.get.mockResolvedValue({ urlCache: {} }); findAccessibleResources.mockResolvedValue([agentWithS3Avatar._id]); @@ -4131,6 +4269,22 @@ describe('Agent Controllers - Mass Assignment Protection', () => { expect(refreshS3Url).not.toHaveBeenCalled(); }); + test('does not treat a manager page cache as an ACL-wide refresh after role downgrade', async () => { + mockCache.get.mockResolvedValue({ urlCache: {}, scope: 'page' }); + findAccessibleResources.mockResolvedValue([agentWithS3Avatar._id]); + findPubliclyAccessibleResources.mockResolvedValue([]); + refreshS3Url.mockResolvedValue('refreshed-after-downgrade.jpg'); + const mockReq = { user: { id: userA.toString(), role: 'USER' }, query: {} }; + const mockRes = { status: jest.fn().mockReturnThis(), json: jest.fn().mockReturnThis() }; + + await getListAgentsHandler(mockReq, mockRes); + + expect(refreshS3Url).toHaveBeenCalledTimes(1); + expect(mockRes.json.mock.calls[0][0].data[0].avatar.filepath).toBe( + 'refreshed-after-downgrade.jpg', + ); + }); + test('should refresh and persist S3 avatars on cache miss', async () => { mockCache.get.mockResolvedValue(false); findAccessibleResources.mockResolvedValue([agentWithS3Avatar._id]); diff --git a/api/server/controllers/avatar.contentFilter.spec.js b/api/server/controllers/avatar.contentFilter.spec.js index f32675d4f2d..1e7afc7c0d4 100644 --- a/api/server/controllers/avatar.contentFilter.spec.js +++ b/api/server/controllers/avatar.contentFilter.spec.js @@ -131,6 +131,30 @@ describe('entity avatar filename content filtering', () => { unlinkSpy.mockRestore(); }); + it('invalidates the tenant-scoped list avatar cache after an accepted upload', async () => { + const { getLogStores } = require('~/cache'); + const { getStrategyFunctions } = require('~/server/services/Files/strategies'); + const { resizeAvatar } = require('~/server/services/Files/images/avatar'); + const avatarCache = { delete: jest.fn().mockResolvedValue(undefined) }; + getLogStores.mockReturnValue(avatarCache); + getStrategyFunctions.mockReturnValue({ + processAvatar: jest.fn().mockResolvedValue('new-avatar.jpg'), + }); + resizeAvatar.mockResolvedValue(Buffer.from('resized')); + db.getAgent.mockResolvedValue({ id: 'agent-1', avatar: null }); + db.updateAgent.mockResolvedValue({ id: 'agent-1', avatar: { filepath: 'new-avatar.jpg' } }); + inspectContent.mockReturnValue(null); + + const req = createRequest({ agent_id: 'agent-1' }); + req.config.fileStrategy = 's3'; + const res = createResponse(); + await uploadAgentAvatar(req, res); + + expect(res.status).toHaveBeenCalledWith(201); + expect(db.updateAgent).toHaveBeenCalled(); + expect(avatarCache.delete).toHaveBeenCalledWith('user-1:tenant-1:agents_avatar_refresh'); + }); + it('blocks an agent avatar filename before database or storage work', async () => { const req = createRequest({ agent_id: 'agent-1' }); const res = createResponse(); 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/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/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) { From dd21e7fb6d53c6f1a398b1c45004ff4313480f9d Mon Sep 17 00:00:00 2001 From: "lia-by-librechat[bot]" <328778573+lia-by-librechat[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 22:53:31 -0400 Subject: [PATCH 28/73] =?UTF-8?q?=F0=9F=93=A3=20fix:=20Surface=20MCP=20Cre?= =?UTF-8?q?dential=20Failures=20During=20Tool=20Loading=20(#16323)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 📣 fix: Surface MCP Credential Failures During Tool Loading * fix: address MCP credential propagation review --------- Co-authored-by: Lia --- api/app/clients/tools/util/handleTools.js | 109 ++++---------- .../clients/tools/util/handleTools.test.js | 139 ++++++++++++++++++ .../agents/__tests__/openai.spec.js | 63 +++++++- .../request.partialDisconnect.spec.js | 2 + .../__tests__/request.resumeMetadata.spec.js | 67 +++++++++ .../agents/__tests__/responses.unit.spec.js | 49 ++++++ api/server/controllers/agents/openai.js | 9 +- api/server/controllers/agents/request.js | 8 +- api/server/controllers/agents/responses.js | 9 +- .../Endpoints/agents/initialize.spec.js | 37 ++++- packages/api/src/agents/errors.spec.ts | 52 +++++++ packages/api/src/agents/errors.ts | 31 +++- packages/api/src/index.ts | 1 + packages/api/src/mcp/loading.spec.ts | 121 +++++++++++++++ packages/api/src/mcp/loading.ts | 131 +++++++++++++++++ 15 files changed, 722 insertions(+), 106 deletions(-) create mode 100644 packages/api/src/mcp/loading.spec.ts create mode 100644 packages/api/src/mcp/loading.ts diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index d2a5ee02e03..463402e489a 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -6,6 +6,7 @@ const { toolRolePermissions, checkToolRolePermission, createSafeUser, + loadMCPTools, createAuthIdentityContext, selectMCPUpstreamTokenProvider, mcpToolPattern, @@ -621,10 +622,6 @@ const loadTools = async ({ } const loadedTools = (await Promise.all(toolPromises)).flatMap((plugin) => plugin || []); - const mcpToolPromises = []; - /** MCP server tools are initialized sequentially by server */ - let index = -1; - const failedMCPServers = new Set(); const safeUser = createSafeUser(options.req?.user); const requestScopedConnections = options.requestScopedConnections ?? getMCPRequestContext(options.req, options.res); @@ -653,83 +650,33 @@ const loadTools = async ({ }), }); - for (const [serverName, toolConfigs] of Object.entries(requestedMCPTools)) { - index++; - /** @type {LCAvailableTools} */ - let availableTools = options.mcpAvailableTools?.[serverName]; - for (const config of toolConfigs) { - try { - if (failedMCPServers.has(serverName)) { - continue; - } - const mcpParams = { - mcpPermissionContext, - index, - signal, - user: safeUser, - userMCPAuthMap, - configServers, - requestBody: options.requestBody ?? options.req?.body, - requestScopedConnections, - res: options.res, - upstreamTokenProvider, - upstreamTokenProviderResolver, - oboIdentityContext, - streamId: options.req?._resumableStreamId || null, - jobCreatedAt: options.jobCreatedAt, - model: agent?.model ?? model, - serverName: config.serverName, - provider: agent?.provider ?? endpoint, - config: config.config, - }; - - if (config.type === 'all' && toolConfigs.length === 1) { - /** Handle async loading for single 'all' tool config */ - mcpToolPromises.push( - createMCPTools(mcpParams).catch((error) => { - logger.error(`Error loading ${serverName} tools:`, error); - return null; - }), - ); - continue; - } - if (!availableTools) { - try { - availableTools = await getMCPServerTools(safeUser.id, serverName, config.config); - } catch (error) { - logger.error(`Error fetching available tools for MCP server ${serverName}:`, error); - } - } - - /** Handle synchronous loading */ - const mcpTool = - config.type === 'all' - ? await createMCPTools(mcpParams) - : await createMCPTool({ - ...mcpParams, - availableTools, - toolKey: config.toolKey, - onAvailableTools: (tools) => { - availableTools = tools; - }, - }); - - if (Array.isArray(mcpTool)) { - loadedTools.push(...mcpTool); - } else if (mcpTool) { - loadedTools.push(mcpTool); - } else { - failedMCPServers.add(serverName); - logger.warn( - `MCP tool creation failed for "${config.toolKey}", server may be unavailable or unauthenticated.`, - ); - } - } catch (error) { - logger.error(`Error loading MCP tool for server ${serverName}:`, error); - } - } - } - loadedTools.push(...(await Promise.all(mcpToolPromises)).flatMap((plugin) => plugin || [])); + loadedTools.push( + ...(await loadMCPTools({ + userId: user, + requestedTools: requestedMCPTools, + availableTools: options.mcpAvailableTools, + createTools: createMCPTools, + createTool: createMCPTool, + getAvailableTools: getMCPServerTools, + context: { + mcpPermissionContext, + signal, + user: safeUser, + userMCPAuthMap, + configServers, + requestBody: options.requestBody ?? options.req?.body, + requestScopedConnections, + res: options.res, + upstreamTokenProvider, + upstreamTokenProviderResolver, + oboIdentityContext, + streamId: options.req?._resumableStreamId || null, + jobCreatedAt: options.jobCreatedAt, + model: agent?.model ?? model, + provider: agent?.provider ?? endpoint, + }, + })), + ); return { loadedTools, toolContextMap, dynamicToolContextMap, primedCodeFiles }; }; diff --git a/api/app/clients/tools/util/handleTools.test.js b/api/app/clients/tools/util/handleTools.test.js index 8c6c1ecc402..f1c885690c7 100644 --- a/api/app/clients/tools/util/handleTools.test.js +++ b/api/app/clients/tools/util/handleTools.test.js @@ -401,6 +401,145 @@ describe('Tool Handlers', () => { ); }); + describe('MCP credential failure propagation', () => { + const { + OpenIDReauthRequiredError, + MCPAuthenticationRejectedError, + MCPAuthenticationRefreshError, + OboTokenResolutionError, + } = require('@librechat/api'); + const failures = [ + 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), + ]; + const toolKey = (name, server = 'private-mcp') => + `${name}${Constants.mcp_delimiter}${server}`; + const load = (tools, signal) => + loadTools({ + user: fakeUser._id.toString(), + tools, + signal, + options: { req: { user: { id: fakeUser._id.toString(), role: 'USER' }, body: {} } }, + }); + + beforeEach(() => { + mockGetServerConfig.mockResolvedValue({ + type: 'streamable-http', + url: 'https://example.com/mcp', + source: 'yaml', + }); + mockGetMCPServerTools.mockResolvedValue({}); + mockCreateMCPTool.mockReset(); + mockCreateMCPTools.mockReset(); + }); + + it.each(failures)('preserves an all-tools credential failure: %s', async (error) => { + mockCreateMCPTools.mockRejectedValueOnce(error); + await expect(load([toolKey(Constants.mcp_all)])).rejects.toBe(error); + }); + + it.each(failures)('preserves a selected-tool credential failure: %s', async (error) => { + mockCreateMCPTool.mockRejectedValueOnce(error); + await expect(load([toolKey('search')])).rejects.toBe(error); + }); + + it('keeps ordinary unavailable servers optional and loads healthy tools', async () => { + mockCreateMCPTools.mockRejectedValueOnce(new Error('server offline')); + const healthyTool = { name: toolKey('search', 'healthy') }; + mockCreateMCPTool.mockResolvedValueOnce(healthyTool); + await expect(load([toolKey(Constants.mcp_all), healthyTool.name])).resolves.toMatchObject({ + loadedTools: [healthyTool], + }); + }); + + it('preserves owned cancellation but not a dependency abort of a live run', async () => { + const controller = new AbortController(); + const error = new DOMException('Stopped', 'AbortError'); + mockCreateMCPTool.mockRejectedValueOnce(error); + await expect(load([toolKey('search')], controller.signal)).resolves.toMatchObject({ + loadedTools: [], + }); + mockCreateMCPTool.mockRejectedValueOnce(error); + controller.abort(error); + await expect(load([toolKey('search')], controller.signal)).rejects.toBe(error); + }); + + it('observes early bulk rejection while a selected tool is still loading', async () => { + const error = failures[0]; + mockCreateMCPTools.mockRejectedValueOnce(error); + mockCreateMCPTool.mockImplementationOnce( + () => new Promise((resolve) => setImmediate(() => resolve({ name: 'healthy' }))), + ); + await expect(load([toolKey(Constants.mcp_all), toolKey('search', 'healthy')])).rejects.toBe( + error, + ); + }); + + it('settles outstanding bulk loads before surfacing a selected-tool failure', async () => { + const error = failures[0]; + let bulkSettled = false; + mockCreateMCPTools.mockImplementationOnce( + () => + new Promise((resolve) => + setImmediate(() => { + bulkSettled = true; + resolve([{ name: 'healthy' }]); + }), + ), + ); + mockCreateMCPTool.mockRejectedValueOnce(error); + await expect(load([toolKey(Constants.mcp_all, 'healthy'), toolKey('search')])).rejects.toBe( + error, + ); + expect(bulkSettled).toBe(true); + }); + + it.each([false, true])( + 'lets cancellation supersede an auth failure after sibling settlement: authFailure=%s', + async (authFailure) => { + const controller = new AbortController(); + const stopped = new Error('request cancelled'); + let siblingSettled = false; + mockCreateMCPTools.mockImplementationOnce(async () => { + if (authFailure) { + throw failures[0]; + } + return [{ name: 'first' }]; + }); + mockCreateMCPTools.mockImplementationOnce( + () => + new Promise((resolve) => + setImmediate(() => { + controller.abort(stopped); + siblingSettled = true; + resolve([{ name: 'sibling' }]); + }), + ), + ); + await expect( + load( + [toolKey(Constants.mcp_all), toolKey(Constants.mcp_all, 'healthy')], + controller.signal, + ), + ).rejects.toBe(stopped); + expect(siblingSettled).toBe(true); + }, + ); + + it('allows a later request to recover after re-authentication', async () => { + mockCreateMCPTools.mockRejectedValueOnce(failures[0]); + await expect(load([toolKey(Constants.mcp_all)])).rejects.toBe(failures[0]); + mockCreateMCPTools.mockResolvedValueOnce([{ name: 'search' }]); + await expect(load([toolKey(Constants.mcp_all)])).resolves.toMatchObject({ + loadedTools: [{ name: 'search' }], + }); + }); + }); + it('passes request body to chat MCP tool creation and skips stale cache for BODY-scoped servers', async () => { const serverName = 'body-scoped'; const toolKey = `search${Constants.mcp_delimiter}${serverName}`; diff --git a/api/server/controllers/agents/__tests__/openai.spec.js b/api/server/controllers/agents/__tests__/openai.spec.js index 1ba61904821..8659082215b 100644 --- a/api/server/controllers/agents/__tests__/openai.spec.js +++ b/api/server/controllers/agents/__tests__/openai.spec.js @@ -174,6 +174,8 @@ jest.mock('@librechat/agents', () => ({ })); jest.mock('@librechat/api', () => ({ + getAgentErrorMetadata: (...args) => + jest.requireActual('@librechat/api').getAgentErrorMetadata(...args), /* Provisioning moved into this package; the controllers build the callback from it. */ createProvisionFilesCallback: () => async () => {}, createAgentExecutionContext: (context) => context, @@ -428,8 +430,8 @@ jest.mock('~/cache', () => ({ jest.mock('~/server/services/ToolService', () => ({ loadAgentTools: jest.fn().mockResolvedValue([]), loadToolsForExecution: jest.fn().mockResolvedValue([]), - isFatalAgentInitializationError: jest.fn((error) => - ['AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE', 'resource_recovery_required'].includes(error?.code), + isFatalAgentInitializationError: jest.fn((...args) => + jest.requireActual('@librechat/api').isFatalAgentInitializationError(...args), ), })); @@ -1674,6 +1676,63 @@ describe('OpenAIChatCompletionController', () => { ); }); + const credentialCases = () => { + const { + OpenIDReauthRequiredError, + MCPAuthenticationRejectedError, + MCPAuthenticationRefreshError, + OboTokenResolutionError, + } = jest.requireActual('@librechat/api'); + return [ + [new OpenIDReauthRequiredError('Please sign in again'), 401, undefined], + [ + new MCPAuthenticationRejectedError('private-mcp', false), + 403, + 'MCP_AUTHENTICATION_REJECTED', + ], + [ + new MCPAuthenticationRefreshError(new Error('temporary failure')), + 503, + 'MCP_AUTHENTICATION_REFRESH_FAILED', + ], + [ + new OboTokenResolutionError('session_refresh_failed', 'Please sign in again', false), + 403, + 'MCP_AUTHENTICATION_REJECTED', + ], + [ + new OboTokenResolutionError('exchange_failed', 'Temporary exchange failure', true), + 503, + 'MCP_AUTHENTICATION_REFRESH_FAILED', + ], + ]; + }; + it.each(credentialCases())( + 'preserves remote chat credential response metadata: %s', + async (error, status, code) => { + const { initializeAgent, createErrorResponse } = require('@librechat/api'); + const { loadAgentTools } = require('~/server/services/ToolService'); + loadAgentTools.mockRejectedValueOnce(error); + initializeAgent.mockImplementationOnce(async ({ req, res, loadTools, agent }) => { + await loadTools({ + req, + res, + tools: ['search_mcp_private'], + model: agent.model, + agentId: agent.id, + provider: agent.provider, + }); + }); + await OpenAIChatCompletionController(req, res); + expect(res.status).toHaveBeenCalledWith(status); + expect(createErrorResponse).toHaveBeenCalledWith( + error.message, + status < 500 ? 'invalid_request_error' : 'server_error', + code ?? null, + ); + }, + ); + it('returns 503 when an agent expects MCP tools but resolves none', async () => { const { initializeAgent } = require('@librechat/api'); const { loadAgentTools } = require('~/server/services/ToolService'); diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index 437e74c45d7..43e0f23daa6 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -47,6 +47,8 @@ jest.mock('@librechat/data-schemas', () => ({ })); jest.mock('@librechat/api', () => ({ + getAgentErrorMetadata: (...args) => + jest.requireActual('@librechat/api').getAgentErrorMetadata(...args), sendEvent: jest.fn(), isScheduleFireRequest: jest.fn(() => false), exemptFromConcurrencyLimiter: jest.fn(() => false), diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index 058d3302dc8..b3408f34cd5 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -252,6 +252,8 @@ jest.mock('@librechat/data-schemas', () => ({ })); jest.mock('@librechat/api', () => ({ + getAgentErrorMetadata: (...args) => + jest.requireActual('@librechat/api').getAgentErrorMetadata(...args), sendEvent: jest.fn(), logAgentMemorySnapshot: jest.fn(), isScheduleFireRequest: (...args) => mockIsScheduleFireRequest(...args), @@ -4070,6 +4072,71 @@ describe('ResumableAgentController resume metadata', () => { }); }); + it.each([ + new (jest.requireActual('@librechat/api').OpenIDReauthRequiredError)( + 'Please sign in again to continue using this MCP server.', + ), + new (jest.requireActual('@librechat/api').MCPAuthenticationRejectedError)('private-mcp', false), + new (jest.requireActual('@librechat/api').MCPAuthenticationRefreshError)( + new Error('upstream refresh unavailable'), + ), + new (jest.requireActual('@librechat/api').OboTokenResolutionError)( + 'session_refresh_failed', + 'Please sign in again', + false, + ), + new (jest.requireActual('@librechat/api').OboTokenResolutionError)( + 'exchange_failed', + 'Temporary exchange failure', + true, + ), + ])( + 'publishes an actionable MCP initialization failure without an HTTP 401: %s', + async (error) => { + const initializeClient = jest.fn().mockRejectedValue(error); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Use the private tool.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(res.status).toHaveBeenCalledWith(200); + expect(res.status).not.toHaveBeenCalledWith(401); + expect(initializeClient).toHaveBeenCalledTimes(1); + let metadata = { + status: error.status ?? error.statusCode, + ...(error.code ? { code: error.code } : {}), + }; + if (error.name === 'OboTokenResolutionError') { + metadata = { + status: error.retryable ? 503 : 403, + code: error.retryable + ? 'MCP_AUTHENTICATION_REFRESH_FAILED' + : 'MCP_AUTHENTICATION_REJECTED', + retryable: error.retryable, + }; + } + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( + 'conversation-123', + JSON.stringify({ + ...metadata, + error: error.message, + }), + 1000, + expect.objectContaining({ beforeErrorPublication: expect.any(Function) }), + ); + }, + ); + it('finalizes the failed job before releasing the idempotency claim', async () => { mockGenerationJobManager.claimGeneration.mockResolvedValue(wonGenerationClaim()); const initializeClient = jest.fn().mockRejectedValue(new Error('init boom after res.json')); diff --git a/api/server/controllers/agents/__tests__/responses.unit.spec.js b/api/server/controllers/agents/__tests__/responses.unit.spec.js index 069ad64c292..1eb14ff9740 100644 --- a/api/server/controllers/agents/__tests__/responses.unit.spec.js +++ b/api/server/controllers/agents/__tests__/responses.unit.spec.js @@ -197,6 +197,8 @@ jest.mock('@librechat/agents', () => ({ })); jest.mock('@librechat/api', () => ({ + getAgentErrorMetadata: (...args) => + jest.requireActual('@librechat/api').getAgentErrorMetadata(...args), /* Provisioning moved into this package; the controllers build the callback from it. */ createProvisionFilesCallback: () => async () => {}, createAgentExecutionContext: (context) => context, @@ -1820,6 +1822,53 @@ describe('createResponse controller', () => { }); describe('safe error logging', () => { + const credentialCases = () => { + const { + OpenIDReauthRequiredError, + MCPAuthenticationRejectedError, + MCPAuthenticationRefreshError, + OboTokenResolutionError, + } = jest.requireActual('@librechat/api'); + return [ + [new OpenIDReauthRequiredError('Please sign in again'), 401, undefined], + [ + new MCPAuthenticationRejectedError('private-mcp', false), + 403, + 'MCP_AUTHENTICATION_REJECTED', + ], + [ + new MCPAuthenticationRefreshError(new Error('temporary failure')), + 503, + 'MCP_AUTHENTICATION_REFRESH_FAILED', + ], + [ + new OboTokenResolutionError('session_refresh_failed', 'Please sign in again', false), + 403, + 'MCP_AUTHENTICATION_REJECTED', + ], + [ + new OboTokenResolutionError('exchange_failed', 'Temporary exchange failure', true), + 503, + 'MCP_AUTHENTICATION_REFRESH_FAILED', + ], + ]; + }; + it.each(credentialCases())( + 'preserves remote Responses credential metadata: %s', + async (error, status, code) => { + const api = require('@librechat/api'); + api.initializeAgent.mockRejectedValueOnce(error); + await createResponse(req, res); + expect(api.sendResponsesErrorResponse).toHaveBeenCalledWith( + res, + status, + error.message, + status < 500 ? 'invalid_request' : 'server_error', + ...(code ? [code] : []), + ); + }, + ); + it('does not classify a client disconnect as an upstream model error', async () => { const api = require('@librechat/api'); const { logger } = require('@librechat/data-schemas'); diff --git a/api/server/controllers/agents/openai.js b/api/server/controllers/agents/openai.js index 76d60a7c458..23fa533a1db 100644 --- a/api/server/controllers/agents/openai.js +++ b/api/server/controllers/agents/openai.js @@ -55,6 +55,7 @@ const { isContentFilterError, getSafeErrorMetadata, getUserFacingProviderError, + getAgentErrorMetadata, getRemoteAgentPermissions, createToolExecuteHandler, createOwnedToolEndHandler, @@ -260,13 +261,11 @@ function handleExecutionError({ error, res, context, appConfig }) { error.body.error, ); } - const statusCode = - typeof error?.status === 'number' && error.status >= 400 && error.status < 600 - ? error.status - : 500; + const errorMetadata = getAgentErrorMetadata(error); + const statusCode = errorMetadata?.status ?? 500; const errorType = statusCode >= 400 && statusCode < 500 ? 'invalid_request_error' : 'server_error'; - const errorCode = !protectionEnabled && typeof error?.code === 'string' ? error.code : null; + const errorCode = !protectionEnabled ? (errorMetadata?.code ?? null) : null; sendErrorResponse(res, statusCode, errorMessage, errorType, errorCode); } diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 3cfa9d71aed..79dd2d2fbef 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -46,6 +46,7 @@ const { resolveAgentTurnExecutionPlan, logAgentMemorySnapshot, getCodeWorkspaceSelectionErrorDetails, + getAgentErrorMetadata, shouldPersistCodeWorkspaceInitializationError, resolvePersistableCodeEnvironmentDecision, getFailedTurnTraceFields, @@ -113,13 +114,12 @@ function getInitializationFailure(error) { }; } - const candidateStatus = error?.status ?? error?.statusCode; - if (!Number.isInteger(candidateStatus) || candidateStatus < 400 || candidateStatus >= 600) { + const metadata = getAgentErrorMetadata(error); + if (!metadata?.status) { return null; } return { - status: candidateStatus, - ...(typeof error?.code === 'string' ? { code: error.code } : {}), + ...metadata, ...getCodeWorkspaceSelectionErrorDetails(error), error: error?.message || 'Failed to start generation', }; diff --git a/api/server/controllers/agents/responses.js b/api/server/controllers/agents/responses.js index f44a9c37c50..f09f292dd6e 100644 --- a/api/server/controllers/agents/responses.js +++ b/api/server/controllers/agents/responses.js @@ -53,6 +53,7 @@ const { isContentFilterError, getSafeErrorMetadata, getUserFacingProviderError, + getAgentErrorMetadata, createToolExecuteHandler, createOwnedToolEndHandler, resolveRecursionLimit, @@ -147,12 +148,10 @@ function handleExecutionError({ error, res, appConfig }) { error.body.error, ); } - const statusCode = - typeof error?.status === 'number' && error.status >= 400 && error.status < 600 - ? error.status - : 500; + const errorMetadata = getAgentErrorMetadata(error); + const statusCode = errorMetadata?.status ?? 500; const errorType = statusCode >= 400 && statusCode < 500 ? 'invalid_request' : 'server_error'; - const errorCode = !protectionEnabled && typeof error?.code === 'string' ? error.code : undefined; + const errorCode = !protectionEnabled ? errorMetadata?.code : undefined; if (errorCode === undefined) { sendResponsesErrorResponse(res, statusCode, errorMessage, errorType); } else { diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index fe65c2493ab..62b05309f01 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -69,13 +69,8 @@ jest.mock('~/server/services/ToolService', () => ({ loadAgentTools: jest.fn(), loadToolsForExecution: (...args) => mockLoadToolsForExecution(...args), getAccessibleMcpServerNames: (...args) => mockGetAccessibleMcpServerNames(...args), - isFatalAgentInitializationError: (error) => - [ - 'AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE', - 'resource_recovery_required', - 'stateful_code_environment_not_allowed', - 'code_workspace_unavailable', - ].includes(error?.code), + isFatalAgentInitializationError: + jest.requireActual('@librechat/api').isFatalAgentInitializationError, })); jest.mock('~/server/controllers/ModelController', () => ({ @@ -452,6 +447,34 @@ describe('initializeClient — processAgent ACL gate', () => { ).rejects.toBe(toolError); }); + it.each([ + new (require('@librechat/api').OpenIDReauthRequiredError)('Please sign in again'), + new (require('@librechat/api').MCPAuthenticationRejectedError)('private-mcp', false), + new (require('@librechat/api').MCPAuthenticationRefreshError)(new Error('Retry later')), + ])('preserves credential failure through the runtime agent loader: %s', async (toolError) => { + loadAgentTools.mockRejectedValueOnce(toolError); + mockInitializeAgent.mockImplementationOnce(async ({ req, res, loadTools, agent }) => { + await loadTools({ + req, + res, + tools: ['run_query_mcp_warehouse'], + model: agent.model, + agentId: agent.id, + provider: agent.provider, + }); + return makePrimaryConfig([]); + }); + + await expect( + initializeClient({ + req: makeReq(), + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }), + ).rejects.toBe(toolError); + }); + it('aborts the run when a handoff target resolves none of its expected MCP tools', async () => { const target = await createAgent({ id: AUTHORIZED_ID, 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/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/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; +} From c6282fa888f000b72a8eb94ab4c560dc8803cd85 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 24 Sep 2026 23:15:50 -0400 Subject: [PATCH 29/73] =?UTF-8?q?=F0=9F=A7=B7=20fix:=20Wait=20for=20Long?= =?UTF-8?q?=20Dispatch=20Turns=20Before=20Giving=20Up=20Background=20Code?= =?UTF-8?q?=20Results=20(#16345)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐢 fix: Back Off Waiting Completion Wake-ups and Deliver Them When Ready * fix: Signal every readiness path and make the completion wait cap configurable * 🧷 fix: Wait for Long Dispatch Turns Before Giving Up Background Code Results * fix: Mark held completion deliveries instead of pipeline expedite, and announce store-won approval expiry * fix: Keep code tasks protected and wait without a fixed cutoff * fix: Scope settle expedites to the resumed conversation and close the remaining signal gaps * style: Format the settlement waiter * fix: Announce a won approval expiry once and release subagent wake-up registrations * fix: Outlast terminal persistence before giving up an unanchored code result * fix: Anchor code results the moment their dispatch turn settles * fix: Scope harvest epoch pinning to the owning stream --------- Co-authored-by: Lia --- api/server/controllers/agents/callbacks.js | 9 +- .../services/Endpoints/agents/initialize.js | 2 + .../Endpoints/agents/initialize.spec.js | 3 + packages/api/src/agents/background.spec.ts | 48 +++ packages/api/src/agents/background.ts | 23 +- .../src/agents/handlers.background.spec.ts | 53 ++++ packages/api/src/agents/handlers.ts | 14 +- .../api/src/agents/harvest.lifecycle.spec.ts | 156 ++++++++++ packages/api/src/agents/harvest.spec.ts | 215 ++++++++++++++ packages/api/src/agents/harvest.ts | 163 +++++++++-- .../api/src/stream/GenerationJobManager.ts | 19 +- .../api/src/stream/__tests__/settled.spec.ts | 276 ++++++++++++++++++ packages/api/src/stream/index.ts | 4 + .../api/src/stream/interfaces/IJobStore.ts | 10 +- packages/api/src/stream/settled.ts | 114 ++++++++ 15 files changed, 1065 insertions(+), 44 deletions(-) create mode 100644 packages/api/src/agents/harvest.lifecycle.spec.ts create mode 100644 packages/api/src/stream/__tests__/settled.spec.ts create mode 100644 packages/api/src/stream/settled.ts diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 4b27bde7fed..a42afb36257 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -20,6 +20,7 @@ const { sendEvent, computeUsageCostUSD, GenerationJobManager, + waitForGenerationSettled, writeAttachmentEvent, createToolExecuteHandler, createOwnedToolEndHandler, @@ -1286,14 +1287,20 @@ function createPtcProgressEmitter({ res, streamId = null, jobCreatedAt }) { * output?: string; * attachments?: Object[]; * }) => Promise} params.updateToolCallResult + * @param {number} [params.jobCreatedAt] - Immutable dispatch generation epoch. + * @param {string} [params.streamId] - The stream owning that epoch. */ -function createBackgroundCodeResultHandler({ req, updateToolCallResult }) { +function createBackgroundCodeResultHandler({ req, updateToolCallResult, jobCreatedAt, streamId }) { return createCodeHarvestHandler({ req, updateToolCallResult, preflightCodeOutputBatch, processCodeOutput, runPreviewFinalize, + generationCreatedAt: jobCreatedAt, + generationStreamId: streamId, + waitForGenerationSettled: (conversationId, options) => + waitForGenerationSettled(GenerationJobManager, conversationId, options), }); } diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 875d085951b..f6f0f3ebd0d 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -511,6 +511,8 @@ const initializeClientWithProvider = async ({ }, persistBackgroundCodeResult: createBackgroundCodeResultHandler({ req, + streamId, + jobCreatedAt, updateToolCallResult: db.updateToolCallResult, }), backgroundToolCompletion: { diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 62b05309f01..b93442123f3 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -304,6 +304,9 @@ describe('initializeClient — processAgent ACL gate', () => { jobCreatedAt: 1234, }), ); + expect( + require('~/server/controllers/agents/callbacks').createBackgroundCodeResultHandler, + ).toHaveBeenCalledWith(expect.objectContaining({ jobCreatedAt: 1234 })); expect(createToolEndCallback).toHaveBeenCalledWith( expect.objectContaining({ streamId: 'conv_1', diff --git a/packages/api/src/agents/background.spec.ts b/packages/api/src/agents/background.spec.ts index 5418fdeb3f1..70dca239679 100644 --- a/packages/api/src/agents/background.spec.ts +++ b/packages/api/src/agents/background.spec.ts @@ -1369,6 +1369,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; diff --git a/packages/api/src/agents/background.ts b/packages/api/src/agents/background.ts index 2c9b744ee47..4caf96f09a6 100644 --- a/packages/api/src/agents/background.ts +++ b/packages/api/src/agents/background.ts @@ -834,7 +834,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 +863,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 +1625,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 { diff --git a/packages/api/src/agents/handlers.background.spec.ts b/packages/api/src/agents/handlers.background.spec.ts index 37ad6535d64..6eac3eedb1d 100644 --- a/packages/api/src/agents/handlers.background.spec.ts +++ b/packages/api/src/agents/handlers.background.spec.ts @@ -2621,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 edbab15caa5..601d2dd077b 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -315,6 +315,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. */ @@ -6169,6 +6170,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 @@ -6257,10 +6265,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, 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/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index a92318e29f2..7a2120965ee 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, @@ -3772,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; 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/index.ts b/packages/api/src/stream/index.ts index ef72b19ee5f..9c8180e108e 100644 --- a/packages/api/src/stream/index.ts +++ b/packages/api/src/stream/index.ts @@ -85,3 +85,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(); + }); +} From 34938ebcbad520be862fb3862fbfd757b9b026ad Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 24 Sep 2026 23:16:17 -0400 Subject: [PATCH 30/73] =?UTF-8?q?=F0=9F=93=AE=20fix:=20Count=20Undelivered?= =?UTF-8?q?=20Background=20Results=20as=20Outstanding=20(#16343)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 📮 fix: Count Undelivered Background Results as Outstanding * fix: Reconcile delivered tasks, discard exactly, and count pending subagent results * fix: Settle manually claimed remote deliveries and keep pending controls while wake-ups drain * fix: Report dead-lettered deliveries as failed and read subagent wake-ups from the durable store * fix: List restored dead letters and retire deliveries on local durable claims --- .../Endpoints/agents/backgroundCompletion.js | 12 + .../services/Endpoints/agents/initialize.js | 3 + packages/api/src/agents/background.spec.ts | 517 ++++++++++++++++++ packages/api/src/agents/background.ts | 261 ++++++++- .../api/src/agents/backgroundCompletion.ts | 53 ++ .../agents/backgroundCompletionWakeup.spec.ts | 141 +++++ .../src/agents/backgroundCompletionWakeup.ts | 99 ++++ packages/api/src/agents/handlers.ts | 15 +- packages/api/src/agents/triggers/service.ts | 3 +- .../src/methods/triggerDelivery.spec.ts | 246 +++++++++ .../src/methods/triggerDelivery.ts | 196 +++++++ 11 files changed, 1527 insertions(+), 19 deletions(-) diff --git a/api/server/services/Endpoints/agents/backgroundCompletion.js b/api/server/services/Endpoints/agents/backgroundCompletion.js index da93d982424..94d1f52eec9 100644 --- a/api/server/services/Endpoints/agents/backgroundCompletion.js +++ b/api/server/services/Endpoints/agents/backgroundCompletion.js @@ -1,9 +1,14 @@ const { createBackgroundToolCompletionWakeupHandler, createBackgroundToolDeadClaimRecovery, + createPendingBackgroundCompletions, createBackgroundToolResultHandler, claimBackgroundToolResult: claimResult, } = require('@librechat/api'); +const { + listPendingAgentBackgroundToolCompletions, + listUndeliveredAgentTriggerTaskIds, +} = require('~/models'); const { enqueueAgentTrigger, persistAgentBackgroundToolResult, @@ -23,6 +28,12 @@ const preregisterBackgroundToolCompletion = createBackgroundToolCompletionWakeup (deliveryKey) => expediteCompletionWakeups({ deliveryKeys: [deliveryKey] }), ); +const pendingBackgroundToolCompletions = createPendingBackgroundCompletions({ + list: listPendingAgentBackgroundToolCompletions, + listTaskIds: listUndeliveredAgentTriggerTaskIds, + retire: retireAgentTrigger, +}); + function createBackgroundToolResultPersistence({ req, updateToolCallResult }) { return createBackgroundToolResultHandler({ req, updateToolCallResult }); } @@ -46,6 +57,7 @@ function createDeadBackgroundToolClaimRecovery( module.exports = { preregisterBackgroundToolCompletion, + pendingBackgroundToolCompletions, createBackgroundToolResultPersistence, claimBackgroundToolResult, createDeadBackgroundToolClaimRecovery, diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index f6f0f3ebd0d..31782af1b0d 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -94,6 +94,7 @@ const { processAddedConvo } = require('./addedConvo'); const subagentThreadTaskStore = require('./subagentThreadStore'); const { preregisterBackgroundToolCompletion, + pendingBackgroundToolCompletions, createBackgroundToolResultPersistence, claimBackgroundToolResult, createDeadBackgroundToolClaimRecovery, @@ -517,6 +518,8 @@ const initializeClientWithProvider = async ({ }), backgroundToolCompletion: { ...(completionWakeupsEnabled ? { preregister: preregisterBackgroundToolCompletion } : {}), + /** Deliveries admitted before wake-ups were disabled still drain and still count. */ + pending: pendingBackgroundToolCompletions, persist: createBackgroundToolResultPersistence({ req, updateToolCallResult: db.updateToolCallResult, diff --git a/packages/api/src/agents/background.spec.ts b/packages/api/src/agents/background.spec.ts index 70dca239679..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', + ); }); }); @@ -3105,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 = { @@ -3459,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 4caf96f09a6..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 @@ -1857,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. @@ -1917,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 }, @@ -1931,6 +2009,7 @@ function serializeTask( : {}), ...taskTimings(task), ...resultFields(task, includeResult), + ...taskDelivery(task), ...taskNote(task), ...(task.error !== undefined ? { error: task.error } : {}), }; @@ -1970,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( @@ -2130,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) ?? {}; @@ -2149,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, @@ -2268,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; @@ -2416,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; @@ -2497,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({ @@ -2587,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, @@ -2603,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, @@ -2624,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 02d57cb7d8e..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. */ @@ -65,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 a30d5a4b499..652ba4df253 100644 --- a/packages/api/src/agents/backgroundCompletionWakeup.spec.ts +++ b/packages/api/src/agents/backgroundCompletionWakeup.spec.ts @@ -4,6 +4,7 @@ import type { CodeApprovalMode } from 'librechat-data-provider'; import type { EnqueueBackgroundToolCompletion } from './backgroundCompletionWakeup'; import { BACKGROUND_TOOL_WAKEUP_INPUT_MAX_CHARS, + createPendingBackgroundCompletions, createBackgroundToolCompletionWakeupHandler, createBackgroundToolCompletionWakeupResolver, createBackgroundToolDeadClaimRecovery, @@ -891,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 e5ef7ca6312..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, @@ -23,6 +25,7 @@ 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'; @@ -544,6 +547,102 @@ export function createBackgroundToolCompletionWakeupResolver({ }; } +/** 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, diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 601d2dd077b..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'; @@ -346,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; @@ -6677,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/triggers/service.ts b/packages/api/src/agents/triggers/service.ts index fc2125bef36..63aa707b22c 100644 --- a/packages/api/src/agents/triggers/service.ts +++ b/packages/api/src/agents/triggers/service.ts @@ -213,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: { @@ -770,6 +770,7 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A settledAt: new Date(), ...(options?.onlyIfUnclaimed === true ? { onlyIfUnclaimed: true } : {}), ...(options?.onlyIfDead === true ? { onlyIfDead: true } : {}), + ...(options?.requireTransition === true ? { requireTransition: true } : {}), }, recovery, ), diff --git a/packages/data-schemas/src/methods/triggerDelivery.spec.ts b/packages/data-schemas/src/methods/triggerDelivery.spec.ts index 7498b932c49..a6ad23cb447 100644 --- a/packages/data-schemas/src/methods/triggerDelivery.spec.ts +++ b/packages/data-schemas/src/methods/triggerDelivery.spec.ts @@ -795,6 +795,252 @@ describe('agent trigger delivery methods', () => { }); }); + 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( diff --git a/packages/data-schemas/src/methods/triggerDelivery.ts b/packages/data-schemas/src/methods/triggerDelivery.ts index 1e2b6ead6fb..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,32 @@ 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[]; @@ -316,6 +360,9 @@ export interface AgentTriggerDeliveryMethods { /** 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; @@ -329,6 +376,20 @@ 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; @@ -2077,6 +2138,9 @@ export function createAgentTriggerDeliveryMethods( /** 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 { @@ -2164,6 +2228,9 @@ export function createAgentTriggerDeliveryMethods( } return true; } + if (input.requireTransition === true) { + return false; + } return ( (await Delivery().exists({ deliveryKey: input.deliveryKey, @@ -2263,6 +2330,133 @@ 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 @@ -4094,6 +4288,8 @@ export function createAgentTriggerDeliveryMethods( retireAgentTriggerDelivery, renewAgentTriggerDeliveryProducerLease, getAgentTriggerDeliveryProducerLease, + listPendingAgentBackgroundToolCompletions, + listUndeliveredAgentTriggerTaskIds, expediteAgentTriggerDeliveries, persistAgentBackgroundToolResult, getAgentBackgroundToolResult, From fe8a993997c0a8854faf8c4d0e555b919535f1de Mon Sep 17 00:00:00 2001 From: "lia-by-librechat[bot]" <328778573+lia-by-librechat[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:35:04 -0400 Subject: [PATCH 31/73] =?UTF-8?q?=F0=9F=9A=A5=20fix:=20Pause=20Rejected=20?= =?UTF-8?q?Steer=20Recoveries=20(#16306)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🚥 fix: Pause Rejected Steer Recoveries * 🚥 fix: Preserve Held Recovery Boundaries and Refused Copy Feedback --------- Co-authored-by: Lia --- .../__tests__/request.resumeMetadata.spec.js | 73 +++--- api/server/controllers/agents/request.js | 6 +- .../Chat/Input/PendingSteerChips.tsx | 62 ++++- .../__tests__/PendingSteerChips.test.tsx | 91 ++++++++ .../components/Chat/Steering/recovery.spec.ts | 66 ++++++ .../src/components/Chat/Steering/recovery.ts | 64 ++++++ .../Chat/__tests__/useQueueDrain.spec.tsx | 75 ++++++ .../Chat/__tests__/useSteerConvert.spec.tsx | 37 +++ .../hooks/Chat/__tests__/useSteering.spec.tsx | 215 ++++++++++++++++-- client/src/hooks/Chat/useQueueDrain.ts | 44 +++- client/src/hooks/Chat/useSteerConvert.ts | 64 ++++-- client/src/hooks/Chat/useSteering.ts | 97 +++++--- .../SSE/__tests__/useResumableSSE.spec.ts | 54 +++++ client/src/hooks/SSE/useResumableSSE.ts | 48 +++- client/src/locales/en/translation.json | 5 + .../api/src/stream/GenerationJobManager.ts | 2 +- packages/api/src/stream/SteerRecovery.ts | 34 ++- .../api/src/stream/__tests__/recovery.spec.ts | 152 +++++++++++++ .../recovery.stream_integration.spec.ts | 133 +++++++++++ .../implementations/InMemoryJobStore.ts | 26 ++- .../stream/implementations/RedisJobStore.ts | 40 ++-- packages/api/src/stream/index.ts | 1 + 22 files changed, 1235 insertions(+), 154 deletions(-) create mode 100644 client/src/components/Chat/Steering/recovery.spec.ts create mode 100644 client/src/components/Chat/Steering/recovery.ts create mode 100644 packages/api/src/stream/__tests__/recovery.spec.ts create mode 100644 packages/api/src/stream/__tests__/recovery.stream_integration.spec.ts diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index b3408f34cd5..ef482d04aec 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -252,6 +252,9 @@ jest.mock('@librechat/data-schemas', () => ({ })); jest.mock('@librechat/api', () => ({ + getSteerRecoveryFailure: jest.requireActual( + '../../../../../packages/api/src/stream/SteerRecovery', + ).getSteerRecoveryFailure, getAgentErrorMetadata: (...args) => jest.requireActual('@librechat/api').getAgentErrorMetadata(...args), sendEvent: jest.fn(), @@ -2943,41 +2946,45 @@ describe('ResumableAgentController resume metadata', () => { ); }); - it('returns a recovery conflict when the atomic store rejects changed source content', async () => { - const mismatch = new Error('recovery mismatch'); - mismatch.code = 'RECOVERY_PAYLOAD_MISMATCH'; - mockGenerationJobManager.createJob.mockRejectedValue(mismatch); - const req = { - user: { id: 'user-123' }, - body: { - text: 'Changed words', - messageId: 'recovered-user-msg', - clientRequestId: 'steer-recovery:server-steer-1', - conversationId: 'conversation-123', - endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, - }, - config: {}, - }; - const res = createResumableResponse(); + it.each(['source_missing', 'protocol_mismatch', 'payload_mismatch'])( + 'returns a recovery conflict with its actual reason (%s)', + async (reason) => { + const mismatch = new Error('recovery mismatch'); + mismatch.code = 'RECOVERY_PAYLOAD_MISMATCH'; + mismatch.reason = reason; + mockGenerationJobManager.createJob.mockRejectedValue(mismatch); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Changed words', + messageId: 'recovered-user-msg', + clientRequestId: 'steer-recovery:server-steer-1', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); - await AgentController(req, res, jest.fn(), jest.fn(), null); + await AgentController(req, res, jest.fn(), jest.fn(), null); - expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( - 'conversation-123', - 'user-123', - 'conversation-123', - expect.objectContaining({ - recoveredSteerId: 'server-steer-1', - recoveredSteerPayload: { text: 'Changed words', fileIds: [] }, - }), - ); - expect(res.status).toHaveBeenCalledWith(409); - expect(res.json).toHaveBeenCalledWith( - expect.objectContaining({ code: 'RECOVERY_PAYLOAD_MISMATCH' }), - ); - expect(mockGenerationJobManager.completeJob).not.toHaveBeenCalled(); - expect(mockGenerationJobManager.steering.consumeRecovered).not.toHaveBeenCalled(); - }); + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( + 'conversation-123', + 'user-123', + 'conversation-123', + expect.objectContaining({ + recoveredSteerId: 'server-steer-1', + recoveredSteerPayload: { text: 'Changed words', fileIds: [] }, + }), + ); + expect(res.status).toHaveBeenCalledWith(409); + expect(res.json).toHaveBeenCalledWith( + expect.objectContaining({ code: 'RECOVERY_PAYLOAD_MISMATCH', reason }), + ); + expect(mockGenerationJobManager.completeJob).not.toHaveBeenCalled(); + expect(mockGenerationJobManager.steering.consumeRecovered).not.toHaveBeenCalled(); + }, + ); it('restores a conditional queued send when a newer generation wins the create CAS', async () => { mockGenerationJobManager.claimGeneration.mockResolvedValue( diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 79dd2d2fbef..7135ca4113d 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -28,6 +28,7 @@ const { acceptAgentStartupTelemetry, isSteerPreemptSupported, buildRecoveredSteerPayload, + getSteerRecoveryFailure, deleteAgentCheckpoint, getAttachmentTitleText, createMCPRuntimeRequestBody, @@ -3439,10 +3440,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit sendGenerationJson( res, 409, - { - code: 'RECOVERY_PAYLOAD_MISMATCH', - error: 'The queued message changed before it could be recovered. Please retry.', - }, + getSteerRecoveryFailure(error, { conversationId, streamId, recoveredSteerId }), generationProtocolVersion, ); } else if (initializationFailure) { diff --git a/client/src/components/Chat/Input/PendingSteerChips.tsx b/client/src/components/Chat/Input/PendingSteerChips.tsx index b18f5270694..2adcc371b1f 100644 --- a/client/src/components/Chat/Input/PendingSteerChips.tsx +++ b/client/src/components/Chat/Input/PendingSteerChips.tsx @@ -29,6 +29,10 @@ import { useDefaultToggleEntry, useInterruptToggleEntry, } from './SteerMenu'; +import { + recoveryDispositionsFamily, + recoveryDisposition, +} from '~/components/Chat/Steering/recovery'; import { useQueuedTurnPortal } from '~/components/Chat/Steering/QueuedTurnPortal'; import { escalatingSteerFamily, revealedQueuedTurnFamily } from '~/store/steer'; import { QUEUE_ICON, STEER_ICON } from '~/components/Chat/Steering/identity'; @@ -141,6 +145,12 @@ function QueuedRow({ const interruptToggle = useInterruptToggleEntry(); const fileCount = message.files?.length ?? 0; const quoteCount = message.quotes?.length ?? 0; + const dispositions = useAtomValue(recoveryDispositionsFamily(steering.queueKey)); + const disposition = recoveryDisposition(dispositions, message); + const recoveryHeld = disposition != null; + const recoveryBlocked = disposition === 'blocked' || disposition === 'cancelled'; + const recoveryPending = disposition === 'cancelling'; + actionPending = actionPending || recoveryPending; const isRecovered = message.recoverySteerId != null; const isRejected = message.server?.status === 'rejected'; const isIndeterminate = message.server?.status === 'indeterminate'; @@ -149,8 +159,11 @@ function QueuedRow({ let statusLabel: | 'com_ui_queued_turn_reconciliation_required' | 'com_ui_steer_delivery_unconfirmed' - | 'com_ui_queued_turn_failed' = 'com_ui_queued_turn_failed'; - if (isIndeterminate) { + | 'com_ui_queued_turn_failed' + | 'com_ui_steer_recovery_held' = 'com_ui_queued_turn_failed'; + if (recoveryHeld) { + statusLabel = 'com_ui_steer_recovery_held'; + } else if (isIndeterminate) { statusLabel = 'com_ui_queued_turn_reconciliation_required'; } else if (isUnconfirmed) { statusLabel = 'com_ui_steer_delivery_unconfirmed'; @@ -166,6 +179,7 @@ function QueuedRow({ const canSteerNow = steering.duringRunActive && steering.canSteer && !isRecovered; const showPrimary = !starting && + !recoveryHeld && serverActionable && (canSteerNow || (!steering.duringRunActive && steering.canSendQueuedNow)); /** `canSteer` is defined as false while paused on approval, but the @@ -258,12 +272,45 @@ function QueuedRow({ onClick: edit, }, ]; + if (recoveryBlocked) { + entries.push( + { + key: 'copy-recovery', + label: localize('com_ui_steer_copy_to_composer'), + icon: