From 2e2c29e10dbde0e54b368977d246cfa4c84200a0 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 31 Aug 2026 14:13:00 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=AB=80=20fix:=20Expire=20Aborted=20Backgr?= =?UTF-8?q?ound=20Tool=20Leases=20(#15428)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: Bound Background Tool Producer Heartbeats * fix: fence timed-out delivery retirement --- .../api/src/agents/backgroundCompletion.ts | 3 + .../src/agents/handlers.background.spec.ts | 22 +++++-- packages/api/src/agents/handlers.ts | 62 ++++++++++++++++--- 3 files changed, 75 insertions(+), 12 deletions(-) diff --git a/packages/api/src/agents/backgroundCompletion.ts b/packages/api/src/agents/backgroundCompletion.ts index a6bd7dc205..914b9d3822 100644 --- a/packages/api/src/agents/backgroundCompletion.ts +++ b/packages/api/src/agents/backgroundCompletion.ts @@ -1,5 +1,8 @@ /** Deadline after which an invocation owner requests cancellation. */ export const BACKGROUND_TASK_TIMEOUT_MS: number = 30 * 60 * 1000; +/** Gives a cooperative tool a short window to settle after cancellation before + * its automatic completion delivery is retired. */ +export const BACKGROUND_TASK_ABORT_GRACE_MS: number = 60 * 1000; /** Three missed heartbeats prove the process-local executor has been lost. */ export const BACKGROUND_TOOL_PRODUCER_LEASE_MS: number = 30_000; export const BACKGROUND_TOOL_PRODUCER_HEARTBEAT_MS: number = 10_000; diff --git a/packages/api/src/agents/handlers.background.spec.ts b/packages/api/src/agents/handlers.background.spec.ts index d57a1c964b..b56f795ae6 100644 --- a/packages/api/src/agents/handlers.background.spec.ts +++ b/packages/api/src/agents/handlers.background.spec.ts @@ -2,6 +2,7 @@ import { z } from 'zod'; import { logger } from '@librechat/data-schemas'; import type { StructuredToolInterface } from '@librechat/agents/langchain/tools'; import type { FiltersConfig } from 'librechat-data-provider'; +import { BACKGROUND_TASK_ABORT_GRACE_MS, BACKGROUND_TASK_TIMEOUT_MS } from './backgroundCompletion'; import { backgroundTaskRegistry, CHECK_BACKGROUND_TASK_NAME } from './background'; import { ContentFilterError } from '../middleware/contentFilter'; import { createToolExecuteHandler } from './handlers'; @@ -540,7 +541,7 @@ describe('createToolExecuteHandler — background tool calls', () => { } }); - it('keeps an abort-resistant invocation running without a false timeout receipt', async () => { + it('retires the automatic wakeup after an abort-resistant invocation exceeds its grace period', async () => { jest.useFakeTimers({ doNotFake: ['setImmediate'] }); try { const tool = { @@ -550,12 +551,14 @@ describe('createToolExecuteHandler — background tool calls', () => { invoke: jest.fn(() => new Promise(() => undefined)), } as unknown as StructuredToolInterface; const persist = jest.fn(async () => true); + const renew = jest.fn(async () => true); + const retire = jest.fn(async () => true); const handler = createToolExecuteHandler({ loadTools: async () => ({ loadedTools: [tool] }), backgroundToolCompletion: { preregister: jest.fn(async () => ({ - renew: jest.fn(async () => true), - retire: jest.fn(async () => true), + renew, + retire, })), persist, claim: jest.fn(async () => ({ status: 'acquired' as const })), @@ -576,12 +579,23 @@ describe('createToolExecuteHandler — background tool calls', () => { metadata: { thread_id: 'exec_convo', run_id: 'response-timeout-resistant' }, }); - jest.advanceTimersByTime(31 * 60 * 1000); + jest.advanceTimersByTime(BACKGROUND_TASK_TIMEOUT_MS + BACKGROUND_TASK_ABORT_GRACE_MS); + await flushMicrotasks(); await flushMicrotasks(); const taskId = JSON.parse(dispatch.content).background_task_id as string; expect(backgroundTaskRegistry.get('exec_user', 'exec_convo', taskId)?.status).toBe('running'); expect(persist).not.toHaveBeenCalled(); + expect(renew).toHaveBeenCalled(); + expect(retire).toHaveBeenCalledWith( + 'background task did not settle after its abort grace period', + { onlyIfUnclaimed: true }, + ); + + const renewalsBefore = renew.mock.calls.length; + jest.advanceTimersByTime(3 * 60 * 1000); + await flushMicrotasks(); + expect(renew).toHaveBeenCalledTimes(renewalsBefore); } finally { jest.useRealTimers(); } diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 2004e1f830..3e8ede7dc4 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -67,6 +67,10 @@ import { contentFilterModelBoundBlockResponse, isContentFilterError, } from '~/middleware/contentFilter'; +import { + BACKGROUND_TASK_ABORT_GRACE_MS, + BACKGROUND_TOOL_PRODUCER_HEARTBEAT_MS, +} from './backgroundCompletion'; import { hasIntentArg, stripIntentArg, @@ -75,7 +79,6 @@ import { } from './intent'; import { getSafeErrorMetadata, logAxiosError, runOutsideTracing, truncateMiddle } from '~/utils'; import { resolveCallerCapabilityProjectionSnapshot } from './callerCapabilities'; -import { BACKGROUND_TOOL_PRODUCER_HEARTBEAT_MS } from './backgroundCompletion'; import { buildSkillPrimeMessage, SKILL_FILE_PREFIX } from './skills'; import { createSkillContentDigest } from './compatibility'; import { parseFrontmatter } from '../skills/import'; @@ -4918,11 +4921,15 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand } }; let producerHeartbeatInFlight: Promise | undefined; + let producerHeartbeatStopped = false; const producerAdmission = completionAdmission; const producerHeartbeat = producerAdmission == null ? undefined : setInterval(() => { + if (producerHeartbeatStopped) { + return; + } if (producerHeartbeatInFlight != null) { return; } @@ -4946,12 +4953,51 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand }); }, BACKGROUND_TOOL_PRODUCER_HEARTBEAT_MS); (producerHeartbeat as { unref?: () => void } | undefined)?.unref?.(); + const stopProducerHeartbeat = async (retireReason?: string): Promise => { + if (!producerHeartbeatStopped) { + producerHeartbeatStopped = true; + if (producerHeartbeat != null) { + clearInterval(producerHeartbeat); + } + } + await producerHeartbeatInFlight; + if (retireReason == null || producerAdmission == null) { + return; + } + try { + const retired = await producerAdmission.retire(retireReason, { + onlyIfUnclaimed: true, + }); + if (!retired) { + logger.warn( + `[background] Could not retire timed-out completion delivery for task ${task.id}.`, + ); + } + } catch (retireError) { + logger.warn( + `[background] Failed to retire timed-out completion delivery for task ${task.id}:`, + retireError, + ); + } + }; + let producerRetirementTimeout: ReturnType | undefined; + const requestBackgroundAbort = (): void => { + backgroundAbortController.abort( + new DOMException('Background task timed out', 'AbortError'), + ); + producerRetirementTimeout = setTimeout(() => { + producerRetirementTimeout = undefined; + void stopProducerHeartbeat( + 'background task did not settle after its abort grace period', + ); + }, BACKGROUND_TASK_ABORT_GRACE_MS); + producerRetirementTimeout.unref?.(); + }; void (async () => { try { - const result = await withBackgroundTaskTimeout(invokePromise, () => - backgroundAbortController.abort( - new DOMException('Background task timed out', 'AbortError'), - ), + const result = await withBackgroundTaskTimeout( + invokePromise, + requestBackgroundAbort, ); if (isCodeCall) { markCodeSandboxWarm(); @@ -5093,10 +5139,10 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand await persistBackgroundResult({ output: deliveredError, status: 'error' }); await wakeDetachedActor(); } finally { - if (producerHeartbeat != null) { - clearInterval(producerHeartbeat); + if (producerRetirementTimeout != null) { + clearTimeout(producerRetirementTimeout); } - await producerHeartbeatInFlight; + await stopProducerHeartbeat(); } })(); if (