From 6daafda86fa58eb6c3c0cfbb5f3f93f8017b61fe Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 18 Aug 2026 08:57:37 -0400 Subject: [PATCH 01/15] =?UTF-8?q?=F0=9F=A7=AF=20ci:=20Disarm=20the=20Grace?= =?UTF-8?q?ful-Shutdown=20Force-Exit=20Timer=20on=20Test=20Reset=20(#14972?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(shutdown): disarm the force-exit timer when test state is reset `shutdown()` arms a 60s timer that calls `process.exit(1)` as a safety net for drains that never finish. It is cleared only when the drain runs to completion, so a drain that never settles — an HTTP server whose close callback never fires, a task that hangs — leaves it armed. `__resetShutdownStateForTests()` clears the task list, the shutting-down flag and the server reference, but not that timer. A suite that triggers a signal therefore leaves a live self-destruct behind: `unref` keeps it from holding the process open, but it still fires if anything else keeps the process alive to the timeout, and `process.exit(1)` then takes down whatever is running a minute later. Jest reports that as a bare `process.exit called with "1"` with no failing test, because the run dies before it can print a summary. Track the timer at module scope, clear it from the reset helper, and clear it from a `finally` so a throwing drain step cannot leak it either. The new test fails without the reset change: it starts a drain that never settles, resets state, advances 120s, and asserts the process was not exited. * Scope the force-exit timer to the shutdown that armed it Hoisting the timer to module scope introduced an aliasing hazard: a drain that settles late runs its `finally` against whatever `forceExitTimer` points at by then. If state was reset and a second shutdown armed its own timer in the meantime, the late `finally` cleared the second shutdown's safety net instead of its own. Keep a local handle per shutdown, always clear that, and null the module reference only while it still identifies the same timer. The added test fails without this: it starts a drain whose close callback is withheld, resets state, starts a second shutdown, then releases the first callback and asserts the second net still force-exits. --- packages/api/src/app/shutdown.spec.ts | 59 +++++++++++++++++++++++++++ packages/api/src/app/shutdown.ts | 38 +++++++++++++---- 2 files changed, 89 insertions(+), 8 deletions(-) diff --git a/packages/api/src/app/shutdown.spec.ts b/packages/api/src/app/shutdown.spec.ts index 16f2de5c514..9e0e76f1999 100644 --- a/packages/api/src/app/shutdown.spec.ts +++ b/packages/api/src/app/shutdown.spec.ts @@ -347,4 +347,63 @@ describe('setupGracefulShutdown', () => { expect(calls).toEqual(['async-done']); expect(exitSpy).toHaveBeenCalledWith(0); }); + it('disarms the force-exit timer when shutdown state is reset', async () => { + jest.useFakeTimers(); + try { + // A drain that never settles: the server close callback is never invoked, + // so `shutdown` stays awaiting and never reaches its own `clearTimeout`. + jest.spyOn(server, 'close').mockImplementation(() => server); + setupGracefulShutdown(server); + triggerSignal('SIGTERM'); + await Promise.resolve(); + + // The safety net is armed and would exit the process on its own. + __resetShutdownStateForTests(); + jest.advanceTimersByTime(120_000); + + // Without the reset clearing it, this timer fires long after the suite + // that armed it has finished, killing the run with code 1. + expect(exitSpy).not.toHaveBeenCalledWith(1); + } finally { + jest.useRealTimers(); + } + }); + + it("keeps a later shutdown's safety net when an earlier drain settles late", async () => { + // `setImmediate` stays real so the first shutdown's continuation can actually + // reach its `finally`; only the force-exit timer is faked. + jest.useFakeTimers({ doNotFake: ['setImmediate'] }); + try { + let releaseFirstClose: (() => void) | undefined; + jest.spyOn(server, 'close').mockImplementation((cb?: (err?: Error) => void) => { + if (cb) { + releaseFirstClose = () => cb(); + } + return server; + }); + setupGracefulShutdown(server); + triggerSignal('SIGTERM'); + await flush(); + + // A second shutdown arms its own net after the first is reset away. + __resetShutdownStateForTests(); + const secondServer = http.createServer(); + Object.defineProperty(secondServer, 'listening', { value: true, configurable: true }); + jest.spyOn(secondServer, 'close').mockImplementation(() => secondServer); + setupGracefulShutdown(secondServer); + triggerSignal('SIGTERM'); + await flush(); + + // The first drain settles only now; its `finally` must not disarm the second. + releaseFirstClose?.(); + await flush(); + await flush(); + await flush(); + + jest.advanceTimersByTime(120_000); + expect(exitSpy).toHaveBeenCalledWith(1); + } finally { + jest.useRealTimers(); + } + }); }); diff --git a/packages/api/src/app/shutdown.ts b/packages/api/src/app/shutdown.ts index e8662296d50..656936af3fe 100644 --- a/packages/api/src/app/shutdown.ts +++ b/packages/api/src/app/shutdown.ts @@ -23,6 +23,7 @@ const tasks: ShutdownTask[] = []; let nextRegistrationOrder = 0; let isShuttingDown = false; let httpServer: Server | null = null; +let forceExitTimer: NodeJS.Timeout | null = null; /** * Register a cleanup task for graceful shutdown. Post-drain is the default phase. @@ -73,6 +74,11 @@ export function __resetShutdownStateForTests(): void { nextRegistrationOrder = 0; isShuttingDown = false; httpServer = null; + /** A drain that never settles leaves this armed. It is `unref`'d, so it does + * not hold the process open — but it does fire if anything else keeps the + * process alive past the timeout, exiting a suite that had long since moved + * on with code 1 and no attributable failure. */ + clearForceExitTimer(); } async function runShutdownTasks(phase: ShutdownPhase): Promise { @@ -93,6 +99,13 @@ async function runShutdownTasks(phase: ShutdownPhase): Promise { } } +function clearForceExitTimer(): void { + if (forceExitTimer) { + clearTimeout(forceExitTimer); + forceExitTimer = null; + } +} + async function shutdown(signal: NodeJS.Signals): Promise { if (isShuttingDown) { return; @@ -100,24 +113,33 @@ async function shutdown(signal: NodeJS.Signals): Promise { isShuttingDown = true; logger.info(`Received ${signal}, draining HTTP server...`); + /** Owned locally so a late `finally` from a superseded drain cannot clear the + * safety net belonging to a shutdown that started after it. */ const forceExit = setTimeout(() => { logger.warn(`Graceful shutdown exceeded ${SHUTDOWN_TIMEOUT_MS}ms, forcing exit`); process.exit(1); }, SHUTDOWN_TIMEOUT_MS); forceExit.unref(); + forceExitTimer = forceExit; let exitCode = 0; - const serverClosePromise = closeHttpServer().catch((err) => { - logger.error('Error closing HTTP server during graceful shutdown:', err); - exitCode = 1; - }); + try { + const serverClosePromise = closeHttpServer().catch((err) => { + logger.error('Error closing HTTP server during graceful shutdown:', err); + exitCode = 1; + }); - await runShutdownTasks('pre-drain'); - await serverClosePromise; - await runShutdownTasks('post-drain'); + await runShutdownTasks('pre-drain'); + await serverClosePromise; + await runShutdownTasks('post-drain'); + } finally { + clearTimeout(forceExit); + if (forceExitTimer === forceExit) { + forceExitTimer = null; + } + } - clearTimeout(forceExit); logger.info('Graceful shutdown complete, exiting'); process.exit(exitCode); } From da0491d5dbec567f8e35f38c29fff5e7c206043d Mon Sep 17 00:00:00 2001 From: Ravi Kumar L Date: Tue, 18 Aug 2026 17:16:58 +0200 Subject: [PATCH 02/15] =?UTF-8?q?=F0=9F=92=BB=20fix(agents):=20require=20C?= =?UTF-8?q?ode=20Interpreter=20for=20programmatic=20MCP=20tools=20(#14977)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(agents): require code interpreter for programmatic MCP tools * test(data-provider): fix tool options fixture type * fix(agents): address programmatic tool review feedback * fix(agents): avoid no-op update on version revert --- api/server/controllers/agents/v1.js | 53 ++++- api/server/controllers/agents/v1.spec.js | 224 ++++++++++++++++++ .../SidePanel/Agents/AgentPanel.tsx | 5 +- .../SidePanel/Agents/MCPToolItem.tsx | 9 +- .../ItemDialog/__tests__/McpSection.spec.tsx | 53 ++++- .../Tools/ItemDialog/sections/McpSection.tsx | 21 +- .../Agents/Tools/ToolsMarketplaceDialog.tsx | 6 + .../SidePanel/Agents/Tools/ToolsSection.tsx | 12 +- .../__tests__/ToolsMarketplaceDialog.spec.tsx | 31 ++- .../Tools/__tests__/ToolsSection.spec.tsx | 31 +++ .../__tests__/AgentPanel.helpers.spec.ts | 26 ++ .../Agents/__tests__/MCPToolItem.spec.tsx | 25 +- client/src/locales/en/translation.json | 1 + .../src/agentToolOptions.spec.ts | 38 +++ .../data-provider/src/agentToolOptions.ts | 36 +++ packages/data-provider/src/index.ts | 1 + 16 files changed, 557 insertions(+), 15 deletions(-) create mode 100644 packages/data-provider/src/agentToolOptions.spec.ts create mode 100644 packages/data-provider/src/agentToolOptions.ts diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index 296a36f9087..4ccbdb44a17 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -37,6 +37,7 @@ const { AgentCapabilities, EModelEndpoint, resolveAllowedStatefulCodeEnvironments, + removeCodeExecutionCaller, removeNullishValues, } = require('librechat-data-provider'); const { @@ -271,6 +272,12 @@ const isSubagentsCapabilityEnabled = (req) => { return capabilities.includes(AgentCapabilities.subagents); }; +const isCodeInterpreterCapabilityEnabled = (req) => { + const capabilities = req.config?.endpoints?.[EModelEndpoint.agents]?.capabilities; + if (!Array.isArray(capabilities)) return false; + return capabilities.includes(AgentCapabilities.execute_code); +}; + /** Reject a newly selected stateful workspace scope that the deployment owner * has excluded. Disabled sessions and unrelated edits remain saveable so an * allowlist tightening never silently rewrites or strands an existing agent. */ @@ -533,6 +540,13 @@ const createAgentHandler = async (req, res) => { const validatedData = agentCreateSchema.parse(req.body); const { tools = [], ...agentData } = removeNullishValues(validatedData); + if ( + (!isCodeInterpreterCapabilityEnabled(req) || !tools.includes(Tools.execute_code)) && + agentData.tool_options != null + ) { + agentData.tool_options = removeCodeExecutionCaller(agentData.tool_options); + } + if ( !validateStatefulCodeEnvironment( req, @@ -818,7 +832,12 @@ const updateAgentHandler = async (req, res) => { updateData.stateful_code_sessions !== undefined || updateData.stateful_code_environment !== undefined; const includesToolsConfiguration = Array.isArray(updateData.tools); - if (includesStatefulConfiguration || includesToolsConfiguration) { + const includesToolOptionsConfiguration = updateData.tool_options !== undefined; + if ( + includesStatefulConfiguration || + includesToolsConfiguration || + includesToolOptionsConfiguration + ) { existingAgent = await db.getAgent({ id }); if (!existingAgent) { return res.status(404).json({ error: 'Agent not found' }); @@ -851,6 +870,18 @@ const updateAgentHandler = async (req, res) => { return; } } + + if (includesToolsConfiguration || includesToolOptionsConfiguration) { + const effectiveTools = updateData.tools ?? existingAgent.tools; + const effectiveToolOptions = updateData.tool_options ?? existingAgent.tool_options; + if ( + (!isCodeInterpreterCapabilityEnabled(req) || + !effectiveTools?.includes(Tools.execute_code)) && + effectiveToolOptions != null + ) { + updateData.tool_options = removeCodeExecutionCaller(effectiveToolOptions); + } + } } if (updateData.model_parameters && typeof updateData.model_parameters === 'object') { @@ -1265,6 +1296,14 @@ const duplicateAgentHandler = async (req, res) => { }); } + if ( + (!isCodeInterpreterCapabilityEnabled(req) || + !newAgentData.tools?.includes(Tools.execute_code)) && + newAgentData.tool_options != null + ) { + newAgentData.tool_options = removeCodeExecutionCaller(newAgentData.tool_options); + } + const newAgent = await db.createAgent(newAgentData); try { @@ -1757,6 +1796,18 @@ const revertAgentVersionHandler = async (req, res) => { } } + const effectiveRevertTools = revertUpdates.tools ?? updatedAgent.tools; + const hasCodeExecutionCaller = Object.values(updatedAgent.tool_options ?? {}).some((options) => + options.allowed_callers?.includes('code_execution'), + ); + if ( + (!isCodeInterpreterCapabilityEnabled(req) || + !effectiveRevertTools?.includes(Tools.execute_code)) && + hasCodeExecutionCaller + ) { + revertUpdates.tool_options = removeCodeExecutionCaller(updatedAgent.tool_options); + } + if (updatedAgent.tool_resources) { const removedCount = await pruneToolResourceFileIdsForAgent({ tool_resources: updatedAgent.tool_resources, diff --git a/api/server/controllers/agents/v1.spec.js b/api/server/controllers/agents/v1.spec.js index 435b0da43f0..e0ed01c94f1 100644 --- a/api/server/controllers/agents/v1.spec.js +++ b/api/server/controllers/agents/v1.spec.js @@ -182,6 +182,23 @@ describe('Agent Controllers - Mass Assignment Protection', () => { }); describe('createAgentHandler', () => { + test('removes programmatic tool options when Code Interpreter capability is disabled', async () => { + mockReq.body = { + name: 'Invalid Programmatic Agent', + provider: 'openai', + model: 'gpt-4', + tools: [Tools.execute_code, 'search_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }; + + await createAgentHandler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(201); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({}); + }); + test('rejects a stateful environment excluded by deployment policy', async () => { mockReq.config = { endpoints: { @@ -824,6 +841,136 @@ describe('Agent Controllers - Mass Assignment Protection', () => { expect(agentInDb.name).toBe('Updated Agent'); }); + test('removes newly added programmatic options when Code Interpreter capability is disabled', async () => { + await Agent.updateOne( + { id: existingAgentId }, + { tools: [Tools.execute_code, 'search_mcp_example'] }, + ); + mockReq.user.id = existingAgentAuthorId.toString(); + mockReq.params.id = existingAgentId; + mockReq.body = { + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }; + + await updateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).not.toHaveBeenCalledWith(400); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({}); + }); + + test('removes programmatic callers when Code Interpreter is disabled', async () => { + await Agent.updateOne( + { id: existingAgentId }, + { + tools: [Tools.execute_code, 'search_mcp_example'], + tool_options: { + search_mcp_example: { + allowed_callers: ['code_execution'], + defer_loading: true, + }, + }, + }, + ); + mockReq.user.id = existingAgentAuthorId.toString(); + mockReq.params.id = existingAgentId; + mockReq.body = { tools: ['search_mcp_example'] }; + + await updateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).not.toHaveBeenCalledWith(400); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({ + search_mcp_example: { defer_loading: true }, + }); + }); + + test('allows unrelated edits to a legacy inconsistent agent', async () => { + await Agent.updateOne( + { id: existingAgentId }, + { + tools: ['search_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }, + ); + mockReq.user.id = existingAgentAuthorId.toString(); + mockReq.params.id = existingAgentId; + mockReq.body = { description: 'Still saveable' }; + + await updateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).not.toHaveBeenCalledWith(400); + expect(mockRes.json.mock.calls[0][0].description).toBe('Still saveable'); + }); + + test('allows detaching a programmatic tool from a legacy inconsistent agent', async () => { + await Agent.updateOne( + { id: existingAgentId }, + { + tools: ['search_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }, + ); + mockReq.user.id = existingAgentAuthorId.toString(); + mockReq.params.id = existingAgentId; + mockReq.body = { tools: [] }; + + await updateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).not.toHaveBeenCalledWith(400); + expect(mockRes.json.mock.calls[0][0].tools).toEqual([]); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({}); + }); + + test('allows clearing programmatic options from a legacy inconsistent agent', async () => { + await Agent.updateOne( + { id: existingAgentId }, + { + tools: ['search_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }, + ); + mockReq.user.id = existingAgentAuthorId.toString(); + mockReq.params.id = existingAgentId; + mockReq.body = { tool_options: {} }; + + await updateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).not.toHaveBeenCalledWith(400); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({}); + }); + + test('removes all newly submitted programmatic options from a legacy agent', async () => { + await Agent.updateOne( + { id: existingAgentId }, + { + tools: ['search_mcp_example', 'lookup_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }, + ); + mockReq.user.id = existingAgentAuthorId.toString(); + mockReq.params.id = existingAgentId; + mockReq.body = { + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + lookup_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }; + + await updateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).not.toHaveBeenCalledWith(400); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({}); + }); + test('rejects selecting a stateful environment excluded by deployment policy', async () => { mockReq.user.id = existingAgentAuthorId.toString(); mockReq.params.id = existingAgentId; @@ -1495,6 +1642,83 @@ describe('Agent Controllers - Mass Assignment Protection', () => { const agentInDb = await Agent.findOne({ id: agent.id }).lean(); expect(agentInDb.tool_resources.file_search.file_ids).toEqual([ownedFileId, otherFileId]); }); + + test('duplicateAgentHandler removes programmatic options without Code Interpreter', async () => { + const sourceAgent = await Agent.create({ + id: `agent_${uuidv4()}`, + name: 'Legacy Programmatic Agent', + provider: 'openai', + model: 'gpt-4', + author: mockReq.user.id, + tools: ['search_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }); + const db = require('~/models'); + jest.spyOn(db, 'getActions').mockResolvedValueOnce([]); + mockReq.params.id = sourceAgent.id; + + await duplicateAgentHandler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(201); + expect(mockRes.json.mock.calls[0][0].agent.tool_options).toEqual({}); + }); + + test('revertAgentVersionHandler removes restored programmatic options without Code Interpreter', async () => { + const agent = await Agent.create({ + id: `agent_${uuidv4()}`, + name: 'Current Agent', + provider: 'openai', + model: 'gpt-4', + author: mockReq.user.id, + versions: [ + { + name: 'Legacy Programmatic Agent', + provider: 'openai', + model: 'gpt-4', + tools: ['search_mcp_example'], + tool_options: { + search_mcp_example: { allowed_callers: ['code_execution'] }, + }, + }, + ], + }); + mockReq.params.id = agent.id; + mockReq.body = { version_index: 0 }; + + await revertAgentVersionHandler(mockReq, mockRes); + + expect(mockRes.json).toHaveBeenCalled(); + expect(mockRes.json.mock.calls[0][0].tool_options).toEqual({}); + }); + + test('revertAgentVersionHandler does not update unchanged tool options', async () => { + const agent = await Agent.create({ + id: `agent_${uuidv4()}`, + name: 'Current Agent', + provider: 'openai', + model: 'gpt-4', + author: mockReq.user.id, + versions: [ + { + name: 'Historical Agent', + provider: 'openai', + model: 'gpt-4', + tool_options: {}, + }, + ], + }); + const db = require('~/models'); + const updateAgentSpy = jest.spyOn(db, 'updateAgent'); + mockReq.params.id = agent.id; + mockReq.body = { version_index: 0 }; + + await revertAgentVersionHandler(mockReq, mockRes); + + expect(mockRes.json).toHaveBeenCalled(); + expect(updateAgentSpy).not.toHaveBeenCalled(); + }); }); describe('Mass Assignment Attack Scenarios', () => { diff --git a/client/src/components/SidePanel/Agents/AgentPanel.tsx b/client/src/components/SidePanel/Agents/AgentPanel.tsx index 4dc92260151..7f2517c278e 100644 --- a/client/src/components/SidePanel/Agents/AgentPanel.tsx +++ b/client/src/components/SidePanel/Agents/AgentPanel.tsx @@ -11,6 +11,7 @@ import { ResourceType, EModelEndpoint, PermissionBits, + removeCodeExecutionCaller, resolveStatefulCodeEnvironment, isAssistantsEndpoint, } from 'librechat-data-provider'; @@ -91,6 +92,8 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n * execute_code is disabled so a stale opt-in can't silently reactivate later. */ const normalizedStatefulCodeSessions = data.execute_code === true ? stateful_code_sessions : false; + const normalizedToolOptions = + data.execute_code === true ? tool_options : removeCodeExecutionCaller(tool_options); const normalizedStatefulCodeEnvironment = stateful_code_environment ?? 'user'; const shouldResetAvatar = @@ -118,7 +121,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n recursion_limit, category, support_contact, - tool_options, + tool_options: normalizedToolOptions, skills, skills_enabled, /** A hidden stale 'agent' scope must not survive disabling memory — diff --git a/client/src/components/SidePanel/Agents/MCPToolItem.tsx b/client/src/components/SidePanel/Agents/MCPToolItem.tsx index d77caf4e7c1..b8d971de064 100644 --- a/client/src/components/SidePanel/Agents/MCPToolItem.tsx +++ b/client/src/components/SidePanel/Agents/MCPToolItem.tsx @@ -18,6 +18,7 @@ interface MCPToolItemProps { intentDisabled: boolean; deferredToolsEnabled: boolean; programmaticToolsEnabled: boolean; + programmaticToolsAvailable: boolean; backgroundToolsEnabled: boolean; toolIntentsEnabled: boolean; onToggleSelect: () => void; @@ -44,6 +45,7 @@ export default function MCPToolItem({ onToggleIntent, deferredToolsEnabled, programmaticToolsEnabled, + programmaticToolsAvailable, backgroundToolsEnabled, toolIntentsEnabled, }: MCPToolItemProps) { @@ -95,8 +97,13 @@ export default function MCPToolItem({ icon={Code2} pressed={isProgrammatic} label={localize('com_ui_mcp_programmatic')} - tooltip={localize('com_ui_mcp_click_to_programmatic')} + tooltip={localize( + programmaticToolsAvailable + ? 'com_ui_mcp_click_to_programmatic' + : 'com_ui_mcp_programmatic_requires_code', + )} activeClass="text-violet-500" + disabled={!programmaticToolsAvailable && !isProgrammatic} onToggle={onToggleProgrammatic} /> )} diff --git a/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx b/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx index 0a15331a123..eac90450d2b 100644 --- a/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx @@ -12,7 +12,9 @@ const mockInitializeServer = jest.fn(); const mockIsConnectionDeferred = jest.fn((): boolean => false); const mockToggleIntentAll = jest.fn(); const mockIsToolProgrammaticOnly = jest.fn((_toolId: string): boolean => false); +const mockAreAllToolsProgrammatic = jest.fn((): boolean => false); const mockCapabilities = { + codeEnabled: false, deferredToolsEnabled: false, programmaticToolsEnabled: false, backgroundToolsEnabled: false, @@ -21,10 +23,19 @@ const mockCapabilities = { jest.mock('react-hook-form', () => ({ useFormContext: () => ({ control: {}, setValue: mockSetValue, getValues: mockGetValues }), - useWatch: ({ name }: { name: string }) => - name === 'tool_options' ? mockGetToolOptions() : mockGetValues(), + useWatch: ({ name }: { name: string }) => { + if (name === 'tool_options') { + return mockGetToolOptions(); + } + if (name === 'execute_code') { + return mockCodeInterpreterSelected(); + } + return mockGetValues(); + }, })); +const mockCodeInterpreterSelected = jest.fn((): boolean => false); + jest.mock('~/Providers', () => ({ useAgentPanelContext: () => ({ mcpServersMap: mockMcpServersMap() }), })); @@ -60,7 +71,7 @@ jest.mock('~/hooks', () => ({ toggleToolBackground: jest.fn(), toggleToolIntent: jest.fn(), areAllToolsDeferred: () => false, - areAllToolsProgrammatic: () => false, + areAllToolsProgrammatic: mockAreAllToolsProgrammatic, areAllToolsBackground: () => false, areAllToolsIntent: () => false, toggleDeferAll: jest.fn(), @@ -164,10 +175,18 @@ describe('McpSection', () => { mockToggleIntentAll.mockClear(); mockIsToolProgrammaticOnly.mockReset(); mockIsToolProgrammaticOnly.mockReturnValue(false); + mockAreAllToolsProgrammatic.mockReset(); + mockAreAllToolsProgrammatic.mockReturnValue(false); mockGetToolOptions.mockReset(); mockGetToolOptions.mockReturnValue(undefined); mockMcpServersMap.mockReset(); mockMcpServersMap.mockReturnValue(new Map()); + mockCodeInterpreterSelected.mockReset(); + mockCodeInterpreterSelected.mockReturnValue(false); + mockCapabilities.codeEnabled = false; + mockCapabilities.deferredToolsEnabled = false; + mockCapabilities.programmaticToolsEnabled = false; + mockCapabilities.backgroundToolsEnabled = false; mockCapabilities.toolIntentsEnabled = false; }); @@ -405,6 +424,34 @@ describe('McpSection', () => { expect(mockToggleIntentAll).toHaveBeenCalledWith(item.server.tools); }); + test('bulk programmatic toggle requires Code Interpreter to be available and selected', () => { + mockCapabilities.programmaticToolsEnabled = true; + const { unmount } = render(); + expect(screen.getByRole('button', { name: 'com_ui_mcp_programmatic_all' })).toHaveAttribute( + 'aria-disabled', + 'true', + ); + unmount(); + + mockCapabilities.codeEnabled = true; + mockCodeInterpreterSelected.mockReturnValue(true); + render(); + expect(screen.getByRole('button', { name: 'com_ui_mcp_programmatic_all' })).not.toHaveAttribute( + 'aria-disabled', + ); + }); + + test('bulk programmatic toggle can clear a legacy programmatic configuration', () => { + mockCapabilities.programmaticToolsEnabled = true; + mockAreAllToolsProgrammatic.mockReturnValue(true); + + render(); + + expect( + screen.getByRole('button', { name: 'com_ui_mcp_unprogrammatic_all' }), + ).not.toHaveAttribute('aria-disabled'); + }); + test('bulk intent skips programmatic-only tools (label can never reach them)', () => { mockCapabilities.toolIntentsEnabled = true; mockIsToolProgrammaticOnly.mockImplementation((toolId: string) => toolId === 'mcp:srv:a'); diff --git a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx index ee792c06f79..d03ed996303 100644 --- a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx @@ -3,6 +3,7 @@ import { Clock, Code2, Captions, Zap } from 'lucide-react'; import { useFormContext, useWatch } from 'react-hook-form'; import { Button, Spinner, Checkbox, Skeleton } from '@librechat/client'; import { + AgentCapabilities, Constants, splitMCPToolKey, normalizeServerName, @@ -83,11 +84,15 @@ export default function McpSection({ item }: Props) { const { mcpServersMap, mcpToolsLoading } = useAgentPanelContext(); const { agentsConfig } = useGetAgentsConfig(); const { + codeEnabled, deferredToolsEnabled, programmaticToolsEnabled, backgroundToolsEnabled, toolIntentsEnabled, } = useAgentCapabilities(agentsConfig?.capabilities); + const codeInterpreterSelected = useWatch({ control, name: AgentCapabilities.execute_code }); + const programmaticToolsAvailable = + codeEnabled && programmaticToolsEnabled && codeInterpreterSelected === true; const { isToolDeferred, isToolProgrammatic, @@ -274,6 +279,13 @@ export default function McpSection({ item }: Props) { const allSelected = hasTools && selectedTools.length === tools.length; const allDeferred = areAllToolsDeferred(tools); const allProgrammatic = areAllToolsProgrammatic(tools); + const programmaticBulkLabel = localize( + allProgrammatic ? 'com_ui_mcp_unprogrammatic_all' : 'com_ui_mcp_programmatic_all', + ); + const programmaticBulkTooltip = + programmaticToolsAvailable || allProgrammatic + ? programmaticBulkLabel + : localize('com_ui_mcp_programmatic_requires_code'); const allBackground = areAllToolsBackground(tools); /** Programmatic-only tools can never carry an intent label (the backend's * `canInjectIntentParam` skips non-direct tools), so both the bulk toggle @@ -450,12 +462,10 @@ export default function McpSection({ item }: Props) { icon={Code2} size="md" pressed={allProgrammatic} - label={localize( - allProgrammatic - ? 'com_ui_mcp_unprogrammatic_all' - : 'com_ui_mcp_programmatic_all', - )} + label={programmaticBulkLabel} activeClass="text-violet-600 dark:text-violet-500" + tooltip={programmaticBulkTooltip} + disabled={!programmaticToolsAvailable && !allProgrammatic} onToggle={() => toggleProgrammaticAll(tools)} /> )} @@ -533,6 +543,7 @@ export default function McpSection({ item }: Props) { intentDisabled={isToolProgrammaticOnly(tool.tool_id)} deferredToolsEnabled={deferredToolsEnabled} programmaticToolsEnabled={programmaticToolsEnabled} + programmaticToolsAvailable={programmaticToolsAvailable} backgroundToolsEnabled={backgroundToolsEnabled} toolIntentsEnabled={toolIntentsEnabled} onToggleSelect={() => toggleToolSelect(tool.tool_id)} diff --git a/client/src/components/SidePanel/Agents/Tools/ToolsMarketplaceDialog.tsx b/client/src/components/SidePanel/Agents/Tools/ToolsMarketplaceDialog.tsx index 9b95ec79e0c..4f214167f96 100644 --- a/client/src/components/SidePanel/Agents/Tools/ToolsMarketplaceDialog.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ToolsMarketplaceDialog.tsx @@ -1,6 +1,7 @@ import { useState, useMemo, useCallback } from 'react'; import { Search } from 'lucide-react'; import { useFormContext } from 'react-hook-form'; +import { AgentCapabilities, removeCodeExecutionCaller } from 'librechat-data-provider'; import { Input, OGDialog, @@ -109,6 +110,11 @@ export default function ToolsMarketplaceDialog({ switch (patch.type) { case 'builtin': setValue(patch.field as keyof AgentForm, patch.value as never, { shouldDirty: true }); + if (patch.field === AgentCapabilities.execute_code && patch.value === false) { + setValue('tool_options', removeCodeExecutionCaller(getValues('tool_options')), { + shouldDirty: true, + }); + } break; case 'tool-add': { const current = (getValues('tools') ?? []) as string[]; diff --git a/client/src/components/SidePanel/Agents/Tools/ToolsSection.tsx b/client/src/components/SidePanel/Agents/Tools/ToolsSection.tsx index bf106b9bb4e..882bb7a4759 100644 --- a/client/src/components/SidePanel/Agents/Tools/ToolsSection.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ToolsSection.tsx @@ -1,7 +1,12 @@ import { useState, useMemo, useCallback, useRef } from 'react'; import { Plus } from 'lucide-react'; import { useFormContext, useWatch } from 'react-hook-form'; -import { PermissionTypes, Permissions, AgentCapabilities } from 'librechat-data-provider'; +import { + PermissionTypes, + Permissions, + AgentCapabilities, + removeCodeExecutionCaller, +} from 'librechat-data-provider'; import { Label, Switch, @@ -153,6 +158,11 @@ export default function ToolsSection({ agentId }: Props) { switch (patch.type) { case 'builtin': setValue(patch.field as keyof AgentForm, patch.value as never, { shouldDirty: true }); + if (patch.field === AgentCapabilities.execute_code && patch.value === false) { + setValue('tool_options', removeCodeExecutionCaller(getValues('tool_options')), { + shouldDirty: true, + }); + } break; case 'tool-remove': { const current = (getValues('tools') ?? []) as string[]; diff --git a/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsMarketplaceDialog.spec.tsx b/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsMarketplaceDialog.spec.tsx index 946c4223434..80a257d2f34 100644 --- a/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsMarketplaceDialog.spec.tsx +++ b/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsMarketplaceDialog.spec.tsx @@ -3,8 +3,9 @@ import { fireEvent, render, screen } from '@testing-library/react'; import ToolsMarketplaceDialog from '../ToolsMarketplaceDialog'; const mockSetValue = jest.fn(); -const mockGetValues = jest.fn((): string[] => []); +const mockGetValues = jest.fn((_: string): unknown => []); let mockWatchedTools: string[] = []; +let mockExecuteCode = false; let mockMcpServersMap = new Map(); jest.mock('react-hook-form', () => ({ @@ -17,7 +18,7 @@ jest.mock('react-hook-form', () => ({ const map: Record = { tools: mockWatchedTools, skills: [], - execute_code: false, + execute_code: mockExecuteCode, web_search: false, file_search: false, artifacts: '', @@ -156,6 +157,7 @@ describe('ToolsMarketplaceDialog', () => { mockGetValues.mockClear(); mockGetValues.mockReturnValue([]); mockWatchedTools = []; + mockExecuteCode = false; mockMcpServersMap = new Map(); mockToggleFavorite.mockClear(); mockFavoriteKeys = new Set(); @@ -183,6 +185,31 @@ describe('ToolsMarketplaceDialog', () => { ); }); + test('disabling Code Interpreter clears programmatic MCP callers immediately', () => { + mockExecuteCode = true; + mockGetValues.mockImplementation((name: string) => + name === 'tool_options' + ? { + search: { allowed_callers: ['code_execution'], defer_loading: true }, + direct: { allowed_callers: ['direct'] }, + } + : [], + ); + + render(); + fireEvent.click(screen.getByRole('button', { name: /com_ui_run_code/ })); + + expect(mockSetValue).toHaveBeenCalledWith('execute_code', false, { shouldDirty: true }); + expect(mockSetValue).toHaveBeenCalledWith( + 'tool_options', + { + search: { defer_loading: true }, + direct: { allowed_callers: ['direct'] }, + }, + { shouldDirty: true }, + ); + }); + test('typing in search input filters the catalog', () => { render(); const input = screen.getByPlaceholderText('com_ui_tools_marketplace_search'); diff --git a/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsSection.spec.tsx b/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsSection.spec.tsx index c9959ff2fc1..38672bb5e45 100644 --- a/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsSection.spec.tsx +++ b/client/src/components/SidePanel/Agents/Tools/__tests__/ToolsSection.spec.tsx @@ -198,6 +198,37 @@ describe('ToolsSection', () => { expect(screen.queryByTestId('item-dialog')).not.toBeInTheDocument(); expect(mockSetValue).toHaveBeenCalledWith('file_search', false, { shouldDirty: true }); }); + + test('clears programmatic MCP callers when Code Interpreter is removed', () => { + mockSelected = [ + { + kind: 'builtin', + id: 'execute_code', + name: 'Run Code', + description: '', + iconKey: 'execute_code', + }, + ]; + mockFormValues = { + tool_options: { + search: { allowed_callers: ['code_execution'], defer_loading: true }, + direct: { allowed_callers: ['direct'] }, + }, + }; + + render(); + fireEvent.click(screen.getByRole('button', { name: 'remove-execute_code' })); + + expect(mockSetValue).toHaveBeenCalledWith('execute_code', false, { shouldDirty: true }); + expect(mockSetValue).toHaveBeenCalledWith( + 'tool_options', + { + search: { defer_loading: true }, + direct: { allowed_callers: ['direct'] }, + }, + { shouldDirty: true }, + ); + }); }); describe('use all skills toggle', () => { diff --git a/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts b/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts index da3471824e8..78f26921e26 100644 --- a/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts +++ b/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts @@ -85,6 +85,32 @@ describe('composeAgentUpdatePayload', () => { expect(payload.stateful_code_sessions).toBe(false); }); + it('removes programmatic callers when execute_code is disabled', () => { + const form = createForm(); + form.execute_code = false; + form.tool_options = { + search: { allowed_callers: ['code_execution'], defer_loading: true }, + }; + + const { payload } = composeAgentUpdatePayload(form, 'agent_123'); + + expect(payload.tool_options).toEqual({ search: { defer_loading: true } }); + }); + + it('preserves programmatic callers when execute_code is enabled', () => { + const form = createForm(); + form.execute_code = true; + form.tool_options = { + search: { allowed_callers: ['code_execution'] }, + }; + + const { payload } = composeAgentUpdatePayload(form, 'agent_123'); + + expect(payload.tool_options).toEqual({ + search: { allowed_callers: ['code_execution'] }, + }); + }); + it('preserves stateful_code_sessions when execute_code is enabled', () => { const form = createForm(); form.execute_code = true; diff --git a/client/src/components/SidePanel/Agents/__tests__/MCPToolItem.spec.tsx b/client/src/components/SidePanel/Agents/__tests__/MCPToolItem.spec.tsx index 14863747444..deedff10a72 100644 --- a/client/src/components/SidePanel/Agents/__tests__/MCPToolItem.spec.tsx +++ b/client/src/components/SidePanel/Agents/__tests__/MCPToolItem.spec.tsx @@ -35,6 +35,7 @@ function setup(overrides: Partial> = {} intentDisabled: false, deferredToolsEnabled: false, programmaticToolsEnabled: false, + programmaticToolsAvailable: false, backgroundToolsEnabled: false, toolIntentsEnabled: false, onToggleSelect: jest.fn(), @@ -108,12 +109,34 @@ describe('MCPToolItem', () => { }); test('programmatic is an inline button rendered only when enabled', () => { - const props = setup({ programmaticToolsEnabled: true }); + const props = setup({ programmaticToolsEnabled: true, programmaticToolsAvailable: true }); const programmaticButton = screen.getByRole('button', { name: 'com_ui_mcp_programmatic' }); fireEvent.click(programmaticButton); expect(props.onToggleProgrammatic).toHaveBeenCalledTimes(1); }); + test('programmatic is inert until Code Interpreter is selected', () => { + const props = setup({ programmaticToolsEnabled: true, programmaticToolsAvailable: false }); + const programmaticButton = screen.getByRole('button', { name: 'com_ui_mcp_programmatic' }); + + expect(programmaticButton).toHaveAttribute('aria-disabled', 'true'); + fireEvent.click(programmaticButton); + expect(props.onToggleProgrammatic).not.toHaveBeenCalled(); + }); + + test('an existing programmatic setting can be cleared after Code Interpreter is disabled', () => { + const props = setup({ + programmaticToolsEnabled: true, + programmaticToolsAvailable: false, + isProgrammatic: true, + }); + const programmaticButton = screen.getByRole('button', { name: 'com_ui_mcp_programmatic' }); + + expect(programmaticButton).not.toHaveAttribute('aria-disabled'); + fireEvent.click(programmaticButton); + expect(props.onToggleProgrammatic).toHaveBeenCalledTimes(1); + }); + test('background is an inline button rendered only when enabled', () => { const props = setup({ backgroundToolsEnabled: true }); const backgroundButton = screen.getByRole('button', { name: 'com_ui_mcp_background' }); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index a9499bb3bc7..f4fa8311205 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1505,6 +1505,7 @@ "com_ui_mcp_oauth_secret_reentry_required": "OAuth settings changed. Re-enter the client secret to save this MCP server.", "com_ui_mcp_oauth_timeout": "OAuth login timed out for {{0}}", "com_ui_mcp_programmatic": "Programmatic", + "com_ui_mcp_programmatic_requires_code": "Enable Code Interpreter before making MCP tools programmatic.", "com_ui_mcp_programmatic_all": "Mark all as programmatic", "com_ui_mcp_reauthentication_required": "MCP server '{{0}}' needs authentication. Reconnect to continue; if that fails, revoke its OAuth access and try again.", "com_ui_mcp_server": "MCP Server", diff --git a/packages/data-provider/src/agentToolOptions.spec.ts b/packages/data-provider/src/agentToolOptions.spec.ts new file mode 100644 index 00000000000..f0ce65814dd --- /dev/null +++ b/packages/data-provider/src/agentToolOptions.spec.ts @@ -0,0 +1,38 @@ +import type { AgentToolOptions } from './types/assistants'; +import { removeCodeExecutionCaller } from './agentToolOptions'; + +describe('removeCodeExecutionCaller', () => { + it('removes a programmatic-only entry that has no other options', () => { + expect( + removeCodeExecutionCaller({ + search: { allowed_callers: ['code_execution'] }, + }), + ).toEqual({}); + }); + + it('preserves direct calling and unrelated options', () => { + expect( + removeCodeExecutionCaller({ + search: { + allowed_callers: ['direct', 'code_execution'], + defer_loading: true, + }, + }), + ).toEqual({ + search: { + allowed_callers: ['direct'], + defer_loading: true, + }, + }); + }); + + it('does not mutate its input', () => { + const input: AgentToolOptions = { + search: { allowed_callers: ['code_execution'], run_in_background: true }, + }; + + removeCodeExecutionCaller(input); + + expect(input.search.allowed_callers).toEqual(['code_execution']); + }); +}); diff --git a/packages/data-provider/src/agentToolOptions.ts b/packages/data-provider/src/agentToolOptions.ts new file mode 100644 index 00000000000..a611b497046 --- /dev/null +++ b/packages/data-provider/src/agentToolOptions.ts @@ -0,0 +1,36 @@ +import type { AgentToolOptions, AllowedCaller } from './types/assistants'; + +/** + * Removes Code Interpreter as an allowed caller without mutating the input. + * Tool entries and unrelated options are preserved; an empty entry is removed. + */ +export function removeCodeExecutionCaller( + toolOptions: AgentToolOptions | undefined, +): AgentToolOptions | undefined { + if (toolOptions == null) { + return toolOptions; + } + + const normalized: AgentToolOptions = {}; + for (const [toolName, options] of Object.entries(toolOptions)) { + const callers = options.allowed_callers; + if (callers?.includes('code_execution') !== true) { + normalized[toolName] = options; + continue; + } + + const allowedCallers = callers.filter( + (caller): caller is AllowedCaller => caller !== 'code_execution', + ); + const { allowed_callers: _removed, ...remainingOptions } = options; + const nextOptions = + allowedCallers.length > 0 + ? { ...remainingOptions, allowed_callers: allowedCallers } + : remainingOptions; + if (Object.keys(nextOptions).length > 0) { + normalized[toolName] = nextOptions; + } + } + + return normalized; +} diff --git a/packages/data-provider/src/index.ts b/packages/data-provider/src/index.ts index 73766b315fd..8c98967cf32 100644 --- a/packages/data-provider/src/index.ts +++ b/packages/data-provider/src/index.ts @@ -58,5 +58,6 @@ export { default as createPayload } from './createPayload'; /* feedback */ export * from './feedback'; export * from './parameterSettings'; +export * from './agentToolOptions'; /* code-execution sandbox */ export * from './codeEnvRef'; From a33b128c47e334a6864a88a80ffbfa7baedbd245 Mon Sep 17 00:00:00 2001 From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:02:33 -0700 Subject: [PATCH 03/15] =?UTF-8?q?=F0=9F=AA=AA=20fix:=20Preserve=20Stored?= =?UTF-8?q?=20Access=20Token=20Expiry=20Over=20ID=20Token=20Exp=20(#14982)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🪪 fix: Preserve Stored Access Token Expiry Over ID Token Exp extractOpenIDTokenInfo let the ID token exp claim overwrite the token set's stored expires_at. The ID token is minted at login and never refreshed, so once a session outlives the ID token TTL, isOpenIDTokenValid reports the access token as expired even when expires_at is hours in the future, and OpenID placeholder substitution silently stops: MCP headers configured with {{LIBRECHAT_OPENID_ACCESS_TOKEN}} ship the literal placeholder string as the bearer credential and the receiving server rejects every connection with an unparseable JWT until the user fully logs out and back in. The ID token exp now only fills a missing expiresAt instead of overriding a stored one. Identity claim enrichment from the ID token is unchanged, and the exp fallback for token sets without expires_at is preserved. * 🪪 fix: Validate ID Token Expiry Before ID Token Placeholder Substitution The precedence fix made isOpenIDTokenValid track only the access token expiry, so an MCP header using {{LIBRECHAT_OPENID_ID_TOKEN}} could substitute an ID token that had already expired. The ID token exp is now preserved separately as idTokenExpiresAt and checked at the ID token substitution site, so an expired ID token substitutes empty rather than a stale credential while access token substitution is unaffected. * 🪪 fix: Address OpenID Expiry Review Round Fix expires_at at the source in the OpenID JWT strategy. The stored value described the incoming bearer's exp even when access_token came from the session or a cookie, so it could describe a different credential entirely. A new decodeJwtExpiry helper reads the exp of the token actually stored, and payload.exp is kept only when the raw bearer is the resolved access token. Opaque session or cookie tokens now store no expiry rather than a wrong one. Apply a 30 second clock skew buffer in isOpenIDTokenValid and isIdTokenCurrent via a new exported OPENID_EXPIRY_BUFFER_SECONDS, mirroring OPENID_REUSE_EXPIRY_BUFFER_SECONDS in AuthController. Tokens that would expire in transit are treated as already expired. Make isIdTokenCurrent fail closed when idTokenExpiresAt is absent. exp is REQUIRED in an ID token, so a missing value means the token is malformed or the claims parse threw. The check uses == null so an exp of 0 counts as present and therefore expired. Read the ID token exp with a numeric type check so an exp of 0 records idTokenExpiresAt and fails closed downstream while a non-numeric exp is ignored, and compare the stored expiry with != null so a gap filled expiry of 0 reads as expired instead of as no expiry at all. Raise an actionable re authentication error for the ID token placeholder instead of substituting an empty string. An empty substitution produced a malformed Authorization header and a 400 downstream rather than a clean signal that the user must re authenticate. Raise the same re authentication error from processSingleValue when a user has an OpenID identity, the stored token set is no longer valid, and the value still contains a credential bearing OpenID placeholder, so the expired access token case that motivated this PR signals re auth instead of silently shipping or stripping the placeholder. Only the access token, ID token, and generic token names raise: identity metadata resolves from the user document and an expiry hint never needed a token, so those keep their existing literal then strip behaviour. Unknown placeholder names also stay literal and diagnosable, matching the existing resolvable placeholder policy. Add the comments the review asked for on the exp fallback heuristic, the EXPIRES_AT placeholder semantics, why stale ID token claims stay usable for identity fields, and the advisory nature of the freshness check. * 🪪 fix: Honour Opaque Access Tokens And Type The OpenID Re-Auth Error Drop the ID token exp fallback in extractOpenIDTokenInfo. Storing the access token expiry honestly means an opaque access token now records no expiry, and the fallback then handed the ID token exp authority over a credential it does not describe. A deployment issuing opaque access tokens alongside a short lived ID token saw isOpenIDTokenValid go false and the credential guard reject a perfectly good access token, which worked before this branch. An unknown access token expiry is now treated as no expiry, and the ID token exp only ever gates ID token substitution through idTokenExpiresAt. Give the re-authentication signal a type. OpenIDReauthRequiredError is raised at both the ID token placeholder and the credential placeholder guard, ErrorController maps it to a 401 carrying the actionable message, and the class exposes statusCode so the agent generation path answers 401 instead of a bare 500 for the same condition. Omit rather than blank a header whose credential placeholder is still unresolved on a final resolution pass, since an empty bearer credential is malformed under RFC 6750 while an absent header lets the upstream answer its own challenge. Identity placeholders keep stripping to an empty string. Move the resolvable placeholder docblock onto the pattern it describes, resolve an EXPIRES_AT of 0 as the string 0 for consistency with the neighbouring null checks, and let AuthController consume the exported OPENID_EXPIRY_BUFFER_SECONDS so the 30 second skew allowance has a single definition. --- api/server/controllers/AuthController.js | 4 +- api/server/controllers/AuthController.spec.js | 1 + api/strategies/openIdJwtStrategy.js | 15 +- api/strategies/openIdJwtStrategy.spec.js | 52 +++- packages/api/src/middleware/error.spec.ts | 29 ++ packages/api/src/middleware/error.ts | 6 + packages/api/src/utils/env.spec.ts | 218 +++++++++++++- packages/api/src/utils/env.ts | 44 ++- packages/api/src/utils/oidc.spec.ts | 283 +++++++++++++++++- packages/api/src/utils/oidc.ts | 48 ++- 10 files changed, 670 insertions(+), 30 deletions(-) diff --git a/api/server/controllers/AuthController.js b/api/server/controllers/AuthController.js index b3743df8280..47c7f2df0e6 100644 --- a/api/server/controllers/AuthController.js +++ b/api/server/controllers/AuthController.js @@ -8,6 +8,7 @@ const { findOpenIDUser, getOpenIdIssuer, buildOpenIDRefreshParams, + OPENID_EXPIRY_BUFFER_SECONDS, } = require('@librechat/api'); const { requestPasswordReset, @@ -28,7 +29,6 @@ const { getGraphApiToken } = require('~/server/services/GraphTokenService'); const { getOpenIdConfig, getOpenIdEmail } = require('~/strategies'); const AUTH_REFRESH_USER_PROJECTION = '-password -__v -totpSecret -backupCodes -federatedTokens'; -const OPENID_REUSE_EXPIRY_BUFFER_SECONDS = 30; /** * Max age (ms) LibreChat reuses a cached OpenID session token before forcing an IdP refresh. * Env-overridable (accepts an arithmetic expression, e.g. `60 * 60 * 24 * 1000`, like @@ -110,7 +110,7 @@ const getReusableOpenIDSessionToken = (openidTokens) => { if ( decoded && typeof decoded === 'object' && - decoded.exp > now + OPENID_REUSE_EXPIRY_BUFFER_SECONDS + decoded.exp > now + OPENID_EXPIRY_BUFFER_SECONDS ) { return candidate; } diff --git a/api/server/controllers/AuthController.spec.js b/api/server/controllers/AuthController.spec.js index 40c20bbbe18..607d4b1b7db 100644 --- a/api/server/controllers/AuthController.spec.js +++ b/api/server/controllers/AuthController.spec.js @@ -22,6 +22,7 @@ jest.mock('~/models', () => ({ findUser: jest.fn(), })); jest.mock('@librechat/api', () => ({ + OPENID_EXPIRY_BUFFER_SECONDS: 30, math: jest.fn((value, fallback) => fallback), isEnabled: jest.fn(), findOpenIDUser: jest.fn(), diff --git a/api/strategies/openIdJwtStrategy.js b/api/strategies/openIdJwtStrategy.js index 5cd55020f1e..3aefa03cbb1 100644 --- a/api/strategies/openIdJwtStrategy.js +++ b/api/strategies/openIdJwtStrategy.js @@ -21,6 +21,15 @@ const { const { updateUser, findUser, isAgentTriggerPrincipalActive } = require('~/models'); const getLogStores = require('~/cache/getLogStores'); +function decodeJwtExpiry(token) { + try { + const payload = JSON.parse(Buffer.from(token.split('.')[1], 'base64').toString()); + return typeof payload.exp === 'number' ? payload.exp : undefined; + } catch { + return undefined; + } +} + const getOpenIdJwtAudience = () => { const parsedAudience = (process.env.OPENID_AUDIENCE ?? '') .split(',') @@ -224,11 +233,13 @@ const openIdJwtLogin = (openIdConfig) => { refreshToken = refreshToken || parsedCookies.refreshToken; } + const resolvedAccessToken = accessToken || rawToken; user.federatedTokens = { - access_token: accessToken || rawToken, + access_token: resolvedAccessToken, id_token: idToken, refresh_token: refreshToken, - expires_at: payload.exp, + expires_at: + resolvedAccessToken === rawToken ? payload.exp : decodeJwtExpiry(resolvedAccessToken), }; done(null, user); diff --git a/api/strategies/openIdJwtStrategy.spec.js b/api/strategies/openIdJwtStrategy.spec.js index e893628edbe..a9ade4b1608 100644 --- a/api/strategies/openIdJwtStrategy.spec.js +++ b/api/strategies/openIdJwtStrategy.spec.js @@ -296,7 +296,7 @@ describe('openIdJwtStrategy – token source handling', () => { access_token: 'session-access', id_token: 'session-id', refresh_token: 'session-refresh', - expires_at: payload.exp, + expires_at: undefined, }); }); @@ -315,7 +315,7 @@ describe('openIdJwtStrategy – token source handling', () => { access_token: 'cookie-access', id_token: 'cookie-id', refresh_token: 'cookie-refresh', - expires_at: payload.exp, + expires_at: undefined, }); }); @@ -340,7 +340,7 @@ describe('openIdJwtStrategy – token source handling', () => { access_token: 'session-access', id_token: 'cookie-id', refresh_token: 'session-refresh', - expires_at: payload.exp, + expires_at: undefined, }); }); @@ -357,6 +357,52 @@ describe('openIdJwtStrategy – token source handling', () => { expect(user.federatedTokens.access_token).toBe('raw-bearer-token'); expect(user.federatedTokens.id_token).toBe('cookie-id'); expect(user.federatedTokens.refresh_token).toBe('cookie-refresh'); + expect(user.federatedTokens.expires_at).toBe(payload.exp); + }); + + it('should decode expires_at from a session access token that is itself a JWT', async () => { + const sessionAccessExp = 1234567890; + const sessionAccessToken = `header.${Buffer.from( + JSON.stringify({ sub: 'oidc-123', exp: sessionAccessExp }), + ).toString('base64')}.signature`; + const req = { + headers: { authorization: 'Bearer raw-bearer-token' }, + session: { + openidTokens: { + accessToken: sessionAccessToken, + idToken: 'session-id', + refreshToken: 'session-refresh', + }, + }, + }; + + const { user } = await invokeVerify(req, payload); + + expect(user.federatedTokens.access_token).toBe(sessionAccessToken); + expect(user.federatedTokens.expires_at).toBe(sessionAccessExp); + expect(user.federatedTokens.expires_at).not.toBe(payload.exp); + }); + + it('should store an opaque session access token with no expiry alongside a decodable stale ID token', async () => { + const staleIdToken = `header.${Buffer.from( + JSON.stringify({ sub: 'oidc-123', exp: Math.floor(Date.now() / 1000) - 3600 }), + ).toString('base64')}.signature`; + const req = { + headers: { authorization: 'Bearer raw-bearer-token' }, + session: { + openidTokens: { + accessToken: 'opaque-session-access', + idToken: staleIdToken, + refreshToken: 'session-refresh', + }, + }, + }; + + const { user } = await invokeVerify(req, payload); + + expect(user.federatedTokens.access_token).toBe('opaque-session-access'); + expect(user.federatedTokens.id_token).toBe(staleIdToken); + expect(user.federatedTokens.expires_at).toBeUndefined(); }); it('should set id_token to undefined when not available in session or cookies', async () => { diff --git a/packages/api/src/middleware/error.spec.ts b/packages/api/src/middleware/error.spec.ts index eb636007dbd..99a13771f31 100644 --- a/packages/api/src/middleware/error.spec.ts +++ b/packages/api/src/middleware/error.spec.ts @@ -2,6 +2,7 @@ import { logger, tenantStorage } from '@librechat/data-schemas'; import type { Request, Response } from 'express'; import type { ValidationError, MongoServerError, CustomError } from '~/types'; import { ErrorController, createCustomError } from './error'; +import { OpenIDReauthRequiredError } from '~/utils/oidc'; // Mock the logger jest.mock('@librechat/data-schemas', () => ({ @@ -223,6 +224,34 @@ describe('ErrorController', () => { }); }); + describe('OpenIDReauthRequiredError handling', () => { + it('should map a re-auth error to a 401 carrying the actionable message', () => { + const error = new OpenIDReauthRequiredError( + 'OpenID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + ); + + ErrorController(error, mockReq, mockRes, mockNext); + + expect(mockRes.status).toHaveBeenCalledWith(401); + expect(mockRes.send).toHaveBeenCalledWith({ + error: 'invalid_token', + message: + 'OpenID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + }); + }); + + it('should carry a 401 statusCode for callers that read the status directly', () => { + expect(new OpenIDReauthRequiredError('re-auth').statusCode).toBe(401); + }); + + it('should not fall through to the bare 500 path', () => { + ErrorController(new OpenIDReauthRequiredError('re-auth'), mockReq, mockRes, mockNext); + + expect(mockRes.status).not.toHaveBeenCalledWith(500); + expect(mockRes.send).not.toHaveBeenCalledWith('An unknown error occurred.'); + }); + }); + describe('Unknown error handling', () => { it('should handle unknown errors', () => { const unknownError = new Error('Some unknown error'); diff --git a/packages/api/src/middleware/error.ts b/packages/api/src/middleware/error.ts index b39e5e8a121..ad8afedf222 100644 --- a/packages/api/src/middleware/error.ts +++ b/packages/api/src/middleware/error.ts @@ -3,6 +3,7 @@ import { logger, tenantStorage } from '@librechat/data-schemas'; import type { NextFunction, Request, Response } from 'express'; import type { MongoServerError, ValidationError, CustomError } from '~/types'; import { buildTenantIsolationErrorLogContext } from './auth'; +import { OpenIDReauthRequiredError } from '~/utils/oidc'; const handleDuplicateKeyError = (err: MongoServerError, res: Response) => { logger.warn('Duplicate key error: ' + (err.errmsg || err.message)); @@ -83,6 +84,11 @@ export const ErrorController = ( return handleDuplicateKeyError(error, res); } + if (err instanceof OpenIDReauthRequiredError) { + logger.warn('OpenID re-authentication required: ' + err.message); + return res.status(401).send({ error: 'invalid_token', message: err.message }); + } + if (isCustomError(error) && error.statusCode && error.body) { return res.status(error.statusCode).send(error.body); } diff --git a/packages/api/src/utils/env.spec.ts b/packages/api/src/utils/env.spec.ts index 395f5e70ed3..a34adb0fcb8 100644 --- a/packages/api/src/utils/env.spec.ts +++ b/packages/api/src/utils/env.spec.ts @@ -2257,7 +2257,7 @@ describe('resolveHeaders stripUnresolved', () => { expect(result['X-Convo']).toBe(''); }); - it('strips OpenID token placeholders when no valid token is available', () => { + it('omits OpenID credential headers when no valid token is available', () => { const result = resolveHeaders({ headers: { 'X-Access': '{{LIBRECHAT_OPENID_ACCESS_TOKEN}}', @@ -2267,8 +2267,35 @@ describe('resolveHeaders stripUnresolved', () => { stripUnresolved: true, }); - expect(result['X-Access']).toBe(''); - expect(result['X-Token']).toBe(''); + expect(result).not.toHaveProperty('X-Access'); + expect(result).not.toHaveProperty('X-Token'); + }); + + it('omits the credential header but strips identity placeholders to empty', () => { + const result = resolveHeaders({ + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + 'X-Org': '{{LIBRECHAT_OPENID_USER_ID}}', + }, + user: createTestUser({ id: 'user-123' }), + stripUnresolved: true, + }); + + expect(result).not.toHaveProperty('Authorization'); + expect(result['X-Org']).toBe(''); + }); + + it('preserves credential placeholders literally when stripUnresolved is false', () => { + const result = resolveHeaders({ + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + 'X-Org': '{{LIBRECHAT_OPENID_USER_ID}}', + }, + user: createTestUser({ id: 'user-123' }), + }); + + expect(result.Authorization).toBe('Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}'); + expect(result['X-Org']).toBe('{{LIBRECHAT_OPENID_USER_ID}}'); }); it('leaves unknown and non-resolvable placeholders untouched', () => { @@ -2293,3 +2320,188 @@ describe('resolveHeaders stripUnresolved', () => { expect(result['X-User-Id']).toBe('{{LIBRECHAT_USER_ID}}'); }); }); + +describe('processMCPEnv OpenID re-authentication signalling', () => { + function createOpenIDUser(expiresAt: number): IUser { + return { + ...createTestUser({ id: 'user-123', provider: 'openid' }), + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'stored-access-token', + id_token: 'stored-id-token', + refresh_token: 'stored-refresh-token', + expires_at: expiresAt, + }, + } as IUser; + } + + function tokenlessOpenIDUser(): IUser { + return { + ...createTestUser({ id: 'user-123', provider: 'openid' }), + openidId: 'oidc-sub-456', + } as IUser; + } + + const expiredSeconds = Math.floor(Date.now() / 1000) - 3600; + const validSeconds = Math.floor(Date.now() / 1000) + 3600; + + it('should throw an actionable re-auth error when the token set is expired and a credential placeholder is present', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + }, + }; + + expect(() => processMCPEnv({ options, user: createOpenIDUser(expiredSeconds) })).toThrow( + 'OpenID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + ); + }); + + it('should still resolve other placeholders when the token set is expired but no OpenID placeholder is present', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + 'X-User-Id': '{{LIBRECHAT_USER_ID}}', + }, + }; + + const result = processMCPEnv({ options, user: createOpenIDUser(expiredSeconds) }); + + if (isStreamableHTTPOptions(result)) { + expect(result.headers?.['X-User-Id']).toBe('user-123'); + } else { + throw new Error('Expected streamable-http options'); + } + }); + + it('should leave an unknown OpenID placeholder name literal instead of raising re-auth', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCES_TOKEN}}', + }, + }; + + const result = processMCPEnv({ options, user: createOpenIDUser(expiredSeconds) }); + + if (isStreamableHTTPOptions(result)) { + expect(result.headers?.Authorization).toBe('Bearer {{LIBRECHAT_OPENID_ACCES_TOKEN}}'); + } else { + throw new Error('Expected streamable-http options'); + } + }); + + it('should leave OpenID placeholders untouched for a user with no OpenID identity', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + }, + }; + + const result = processMCPEnv({ options, user: createTestUser({ id: 'user-123' }) }); + + if (isStreamableHTTPOptions(result)) { + expect(result.headers?.Authorization).toBe('Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}'); + } else { + throw new Error('Expected streamable-http options'); + } + }); + + it('should substitute the access token when the token set is still valid', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + }, + }; + + const result = processMCPEnv({ options, user: createOpenIDUser(validSeconds) }); + + if (isStreamableHTTPOptions(result)) { + expect(result.headers?.Authorization).toBe('Bearer stored-access-token'); + } else { + throw new Error('Expected streamable-http options'); + } + }); + + it('should raise re-auth from resolveHeaders when the token set is expired', () => { + expect(() => + resolveHeaders({ + headers: { Authorization: 'Bearer {{LIBRECHAT_OPENID_ID_TOKEN}}' }, + user: createOpenIDUser(expiredSeconds), + stripUnresolved: true, + }), + ).toThrow( + 'OpenID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ID_TOKEN}}', + ); + }); + + it('should leave identity metadata placeholders literal for an OpenID user with no stored tokens', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + 'X-User-Id': '{{LIBRECHAT_OPENID_USER_ID}}', + }, + }; + + const result = processMCPEnv({ options, user: tokenlessOpenIDUser() }); + + if (isStreamableHTTPOptions(result)) { + expect(result.headers?.['X-User-Id']).toBe('{{LIBRECHAT_OPENID_USER_ID}}'); + } else { + throw new Error('Expected streamable-http options'); + } + + expect( + resolveHeaders({ + headers: { 'X-User-Id': '{{LIBRECHAT_OPENID_USER_ID}}' }, + user: tokenlessOpenIDUser(), + stripUnresolved: true, + })['X-User-Id'], + ).toBe(''); + }); + + it('should still raise re-auth for a credential placeholder when no tokens are stored', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + }, + }; + + expect(() => processMCPEnv({ options, user: tokenlessOpenIDUser() })).toThrow( + 'OpenID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + ); + }); + + it('should not raise re-auth for a metadata-only template when the token set is expired', () => { + const options: MCPOptions = { + type: 'streamable-http', + url: 'https://api.example.com', + headers: { + 'X-User-Id': '{{LIBRECHAT_OPENID_USER_ID}}', + 'X-User-Email': '{{LIBRECHAT_OPENID_USER_EMAIL}}', + 'X-User-Name': '{{LIBRECHAT_OPENID_USER_NAME}}', + 'X-Expires': '{{LIBRECHAT_OPENID_EXPIRES_AT}}', + }, + }; + + const result = processMCPEnv({ options, user: createOpenIDUser(expiredSeconds) }); + + if (isStreamableHTTPOptions(result)) { + expect(result.headers?.['X-User-Id']).toBe('{{LIBRECHAT_OPENID_USER_ID}}'); + expect(result.headers?.['X-Expires']).toBe('{{LIBRECHAT_OPENID_EXPIRES_AT}}'); + } else { + throw new Error('Expected streamable-http options'); + } + }); +}); diff --git a/packages/api/src/utils/env.ts b/packages/api/src/utils/env.ts index 1c39d0b9dd9..39a044ed682 100644 --- a/packages/api/src/utils/env.ts +++ b/packages/api/src/utils/env.ts @@ -1,3 +1,4 @@ +import { logger } from '@librechat/data-schemas'; import { extractEnvVariable } from 'librechat-data-provider'; import type { MCPOptions } from 'librechat-data-provider'; import type { IUser } from '@librechat/data-schemas'; @@ -7,6 +8,7 @@ import { isOpenIDTokenValid, extractOpenIDTokenInfo, processOpenIDPlaceholders, + OpenIDReauthRequiredError, } from './oidc'; /** @@ -144,6 +146,8 @@ export function createSafeUser( */ export const ALLOWED_BODY_FIELDS = ['conversationId', 'parentMessageId', 'messageId'] as const; +const OPENID_PLACEHOLDER_NAMES = `LIBRECHAT_OPENID_(?:${OPENID_TOKEN_FIELDS.join('|')}|TOKEN)`; + /** * Matches every placeholder this module knows how to resolve: the enumerated * `{{LIBRECHAT_USER_*}}`, `{{LIBRECHAT_BODY_*}}`, and `{{LIBRECHAT_OPENID_*}}` @@ -155,13 +159,24 @@ const RESOLVABLE_PLACEHOLDER_PATTERN = new RegExp( [ `LIBRECHAT_USER_(?:${ALLOWED_USER_FIELDS.map((field) => field.toUpperCase()).join('|')})`, `LIBRECHAT_BODY_(?:${ALLOWED_BODY_FIELDS.map((field) => field.toUpperCase()).join('|')})`, - `LIBRECHAT_OPENID_(?:${OPENID_TOKEN_FIELDS.join('|')}|TOKEN)`, + OPENID_PLACEHOLDER_NAMES, ] .map((names) => `\\{\\{(?:${names})\\}\\}`) .join('|'), 'g', ); +/** + * The subset of OpenID placeholders that cannot resolve without a usable token + * set. Identity metadata (`USER_ID`, `USER_EMAIL`, `USER_NAME`) comes from the + * user document and `EXPIRES_AT` is only ever a hint, so those must keep their + * pre-existing literal-then-strip behaviour when the token set is invalid + * rather than raising re-auth. Non-global so `exec` stays stateless, and + * unknown names are excluded so a typo stays literal and diagnosable. + */ +const OPENID_CREDENTIAL_PLACEHOLDER_PATTERN = + /\{\{LIBRECHAT_OPENID_(?:ACCESS_TOKEN|ID_TOKEN|TOKEN)\}\}/; + /** * Replaces resolvable-but-unresolved placeholders with an empty string so * LibreChat's internal template syntax is never sent upstream as if it were @@ -320,6 +335,16 @@ function processSingleValue({ const openidTokenInfo = extractOpenIDTokenInfo(user); if (openidTokenInfo && isOpenIDTokenValid(openidTokenInfo)) { value = processOpenIDPlaceholders(value, openidTokenInfo); + } else if (openidTokenInfo) { + const unresolvable = OPENID_CREDENTIAL_PLACEHOLDER_PATTERN.exec(value); + if (unresolvable) { + logger.warn( + `OpenID token is expired or unavailable; cannot resolve ${unresolvable[0]} for the current request`, + ); + throw new OpenIDReauthRequiredError( + `OpenID token is expired or unavailable; re-authentication is required to resolve ${unresolvable[0]}`, + ); + } } if (body) { @@ -610,7 +635,22 @@ export function resolveHeaders(options?: { body, isHeader: true, // Important: Enable header encoding }); - resolvedHeaders[key] = stripUnresolved ? stripUnresolvedPlaceholders(processed) : processed; + if (!stripUnresolved) { + resolvedHeaders[key] = processed; + return; + } + + /** Reached only when the credential guard did not fire, i.e. the user has no OpenID identity at all: blanking the credential would emit `Authorization: Bearer `, which RFC 6750 rejects for a missing b64token */ + const unresolvedCredential = OPENID_CREDENTIAL_PLACEHOLDER_PATTERN.exec(processed); + if (unresolvedCredential) { + logger.warn( + `Omitting header "${key}": ${unresolvedCredential[0]} could not be resolved for the current request`, + ); + delete resolvedHeaders[key]; + return; + } + + resolvedHeaders[key] = stripUnresolvedPlaceholders(processed); }); } diff --git a/packages/api/src/utils/oidc.spec.ts b/packages/api/src/utils/oidc.spec.ts index e7088d9897f..0cefdf60684 100644 --- a/packages/api/src/utils/oidc.spec.ts +++ b/packages/api/src/utils/oidc.spec.ts @@ -1,5 +1,10 @@ -import { extractOpenIDTokenInfo, isOpenIDTokenValid, processOpenIDPlaceholders } from './oidc'; import type { IUser } from '@librechat/data-schemas'; +import { + OpenIDReauthRequiredError, + extractOpenIDTokenInfo, + isOpenIDTokenValid, + processOpenIDPlaceholders, +} from './oidc'; describe('OpenID Token Utilities', () => { describe('extractOpenIDTokenInfo', () => { @@ -116,6 +121,156 @@ describe('OpenID Token Utilities', () => { expect(result?.userId).toBe('user-123'); }); + + it('should keep the stored access token expiry when the ID token carries an older exp', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const staleIdTokenPayload = Buffer.from( + JSON.stringify({ sub: 'oidc-sub-456', exp: nowSeconds - 3600 }), + ).toString('base64'); + const user: Partial = { + id: 'user-123', + provider: 'openid', + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'fresh-access-token', + id_token: `header.${staleIdTokenPayload}.signature`, + expires_at: nowSeconds + 3600, + }, + }; + + const result = extractOpenIDTokenInfo(user); + + expect(result?.expiresAt).toBe(nowSeconds + 3600); + expect(isOpenIDTokenValid(result)).toBe(true); + expect(processOpenIDPlaceholders('Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', result)).toBe( + 'Bearer fresh-access-token', + ); + }); + + it('should leave expiresAt unset when no expires_at is stored, regardless of the ID token exp', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const idTokenPayload = Buffer.from( + JSON.stringify({ sub: 'oidc-sub-456', exp: nowSeconds + 1800 }), + ).toString('base64'); + const user: Partial = { + id: 'user-123', + provider: 'openid', + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'access-token-value', + id_token: `header.${idTokenPayload}.signature`, + }, + }; + + const result = extractOpenIDTokenInfo(user); + + expect(result?.expiresAt).toBeUndefined(); + expect(result?.idTokenExpiresAt).toBe(nowSeconds + 1800); + expect(isOpenIDTokenValid(result)).toBe(true); + }); + + it('should keep an opaque access token valid when the stored ID token is already expired', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const staleIdTokenPayload = Buffer.from( + JSON.stringify({ sub: 'oidc-sub-456', exp: nowSeconds - 3600 }), + ).toString('base64'); + const user: Partial = { + id: 'user-123', + provider: 'openid', + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'opaque-access-token', + id_token: `header.${staleIdTokenPayload}.signature`, + }, + }; + + const result = extractOpenIDTokenInfo(user); + + expect(result?.expiresAt).toBeUndefined(); + expect(result?.idTokenExpiresAt).toBe(nowSeconds - 3600); + expect(isOpenIDTokenValid(result)).toBe(true); + expect(processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ACCESS_TOKEN}}', result)).toBe( + 'opaque-access-token', + ); + expect(() => processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', result)).toThrow( + OpenIDReauthRequiredError, + ); + }); + + it('should gate only the ID token on an exp of 0, leaving access token validity untouched', () => { + const idTokenPayload = Buffer.from(JSON.stringify({ sub: 'oidc-sub-456', exp: 0 })).toString( + 'base64', + ); + const user: Partial = { + id: 'user-123', + provider: 'openid', + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'access-token-value', + id_token: `header.${idTokenPayload}.signature`, + }, + }; + + const result = extractOpenIDTokenInfo(user); + + expect(result?.idTokenExpiresAt).toBe(0); + expect(result?.expiresAt).toBeUndefined(); + expect(isOpenIDTokenValid(result)).toBe(true); + expect(() => processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', result)).toThrow( + /re-authentication is required/, + ); + }); + + it('should ignore a non-numeric ID token exp', () => { + const idTokenPayload = Buffer.from( + JSON.stringify({ sub: 'oidc-sub-456', exp: '1700000000' }), + ).toString('base64'); + const user: Partial = { + id: 'user-123', + provider: 'openid', + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'access-token-value', + id_token: `header.${idTokenPayload}.signature`, + }, + }; + + const result = extractOpenIDTokenInfo(user); + + expect(result?.idTokenExpiresAt).toBeUndefined(); + expect(result?.expiresAt).toBeUndefined(); + expect(() => processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', result)).toThrow( + /re-authentication is required/, + ); + }); + + it('should still enrich identity fields from ID token claims when expires_at is stored', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const idTokenPayload = Buffer.from( + JSON.stringify({ + sub: 'claims-sub', + email: 'claims@example.com', + name: 'Claims Name', + exp: nowSeconds - 3600, + }), + ).toString('base64'); + const user: Partial = { + id: 'user-123', + provider: 'openid', + openidId: 'oidc-sub-456', + federatedTokens: { + access_token: 'access-token-value', + id_token: `header.${idTokenPayload}.signature`, + expires_at: nowSeconds + 3600, + }, + }; + + const result = extractOpenIDTokenInfo(user); + + expect(result?.userId).toBe('claims-sub'); + expect(result?.userEmail).toBe('claims@example.com'); + expect(result?.userName).toBe('Claims Name'); + }); }); describe('isOpenIDTokenValid', () => { @@ -173,7 +328,7 @@ describe('OpenID Token Utilities', () => { expect(isOpenIDTokenValid(tokenInfo)).toBe(false); }); - it('should return true when token is just about to expire (within 1 second)', () => { + it('should return false when token expires within the expiry buffer', () => { const almostExpiredTimestamp = Math.floor(Date.now() / 1000) + 1; const tokenInfo = { accessToken: 'access-token-value', @@ -181,11 +336,85 @@ describe('OpenID Token Utilities', () => { userId: 'oidc-sub-456', }; + expect(isOpenIDTokenValid(tokenInfo)).toBe(false); + }); + + it('should return true when token expires beyond the expiry buffer', () => { + const beyondBufferTimestamp = Math.floor(Date.now() / 1000) + 120; + const tokenInfo = { + accessToken: 'access-token-value', + expiresAt: beyondBufferTimestamp, + userId: 'oidc-sub-456', + }; + expect(isOpenIDTokenValid(tokenInfo)).toBe(true); }); + + it('should pin the buffer boundary at 30 seconds', () => { + const nowMs = 1_700_000_000_000; + const nowSeconds = Math.floor(nowMs / 1000); + const nowSpy = jest.spyOn(Date, 'now').mockReturnValue(nowMs); + + try { + expect( + isOpenIDTokenValid({ + accessToken: 'access-token-value', + expiresAt: nowSeconds + 29, + userId: 'oidc-sub-456', + }), + ).toBe(false); + expect( + isOpenIDTokenValid({ + accessToken: 'access-token-value', + expiresAt: nowSeconds + 31, + userId: 'oidc-sub-456', + }), + ).toBe(true); + } finally { + nowSpy.mockRestore(); + } + }); + + it('should return false when expiresAt is 0', () => { + const tokenInfo = { + accessToken: 'access-token-value', + expiresAt: 0, + userId: 'oidc-sub-456', + }; + + expect(isOpenIDTokenValid(tokenInfo)).toBe(false); + }); }); describe('processOpenIDPlaceholders', () => { + it('should not substitute an expired ID token for the ID token placeholder', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const tokenInfo = { + accessToken: 'fresh-access-token', + idToken: 'stale-id-token-value', + expiresAt: nowSeconds + 3600, + idTokenExpiresAt: nowSeconds - 3600, + }; + + expect(() => processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', tokenInfo)).toThrow( + /re-authentication is required/, + ); + }); + + it('should substitute a current ID token for the ID token placeholder', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const tokenInfo = { + accessToken: 'fresh-access-token', + idToken: 'current-id-token-value', + expiresAt: nowSeconds + 3600, + idTokenExpiresAt: nowSeconds + 1800, + }; + + const result = processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', tokenInfo); + + expect(result).toBe('current-id-token-value'); + }); + it('should replace LIBRECHAT_OPENID_TOKEN with access token', () => { const tokenInfo = { accessToken: 'access-token-value', @@ -213,6 +442,7 @@ describe('OpenID Token Utilities', () => { it('should replace LIBRECHAT_OPENID_ID_TOKEN with id token', () => { const tokenInfo = { + idTokenExpiresAt: Math.floor(Date.now() / 1000) + 3600, idToken: 'id-token-value', userId: 'oidc-sub-456', }; @@ -262,6 +492,7 @@ describe('OpenID Token Utilities', () => { const tokenInfo = { accessToken: 'access-token-value', idToken: 'id-token-value', + idTokenExpiresAt: Math.floor(Date.now() / 1000) + 3600, userId: 'oidc-sub-456', userEmail: 'test@example.com', }; @@ -278,21 +509,43 @@ describe('OpenID Token Utilities', () => { it('should replace empty string when token field is undefined', () => { const tokenInfo = { accessToken: undefined, - idToken: undefined, userId: 'oidc-sub-456', }; - const input = - 'Access: {{LIBRECHAT_OPENID_TOKEN}}, ID: {{LIBRECHAT_OPENID_ID_TOKEN}}, User: {{LIBRECHAT_OPENID_USER_ID}}'; + const input = 'Access: {{LIBRECHAT_OPENID_TOKEN}}, User: {{LIBRECHAT_OPENID_USER_ID}}'; const result = processOpenIDPlaceholders(input, tokenInfo); - expect(result).toBe('Access: , ID: , User: oidc-sub-456'); + expect(result).toBe('Access: , User: oidc-sub-456'); + }); + + it('should throw for the ID token placeholder when no ID token is stored', () => { + const tokenInfo = { + accessToken: 'access-token-value', + userId: 'oidc-sub-456', + }; + + expect(() => processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', tokenInfo)).toThrow( + /re-authentication is required/, + ); + }); + + it('should throw for the ID token placeholder when the ID token has no decodable exp', () => { + const tokenInfo = { + accessToken: 'access-token-value', + idToken: 'malformed-id-token', + userId: 'oidc-sub-456', + }; + + expect(() => processOpenIDPlaceholders('{{LIBRECHAT_OPENID_ID_TOKEN}}', tokenInfo)).toThrow( + /re-authentication is required/, + ); }); it('should handle all placeholder types in one value', () => { const tokenInfo = { accessToken: 'access-token-value', idToken: 'id-token-value', + idTokenExpiresAt: Math.floor(Date.now() / 1000) + 3600, userId: 'oidc-sub-456', userEmail: 'test@example.com', userName: 'Test User', @@ -369,11 +622,10 @@ describe('OpenID Token Utilities', () => { userName: undefined, }; - const input = - 'Access: {{LIBRECHAT_OPENID_TOKEN}}, ID: {{LIBRECHAT_OPENID_ID_TOKEN}}, User: {{LIBRECHAT_OPENID_USER_ID}}'; + const input = 'Access: {{LIBRECHAT_OPENID_TOKEN}}, User: {{LIBRECHAT_OPENID_USER_ID}}'; const result = processOpenIDPlaceholders(input, tokenInfo); - expect(result).toBe('Access: , ID: , User: oidc-sub-456'); + expect(result).toBe('Access: , User: oidc-sub-456'); }); it('should return original value when tokenInfo is null', () => { @@ -428,6 +680,11 @@ describe('OpenID Token Utilities', () => { }); it('should resolve LIBRECHAT_OPENID_ID_TOKEN and LIBRECHAT_OPENID_ACCESS_TOKEN to different values', () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const idTokenPayload = Buffer.from( + JSON.stringify({ sub: 'oidc-sub-456', exp: nowSeconds + 3600 }), + ).toString('base64'); + const myIdToken = `header.${idTokenPayload}.signature`; const user: Partial = { id: 'user-123', provider: 'openid', @@ -436,22 +693,22 @@ describe('OpenID Token Utilities', () => { name: 'Test User', federatedTokens: { access_token: 'my-access-token', - id_token: 'my-id-token', + id_token: myIdToken, refresh_token: 'my-refresh-token', - expires_at: Math.floor(Date.now() / 1000) + 3600, + expires_at: nowSeconds + 3600, }, }; const tokenInfo = extractOpenIDTokenInfo(user); expect(tokenInfo).not.toBeNull(); expect(tokenInfo!.accessToken).toBe('my-access-token'); - expect(tokenInfo!.idToken).toBe('my-id-token'); + expect(tokenInfo!.idToken).toBe(myIdToken); expect(tokenInfo!.accessToken).not.toBe(tokenInfo!.idToken); const input = 'ACCESS={{LIBRECHAT_OPENID_ACCESS_TOKEN}}, ID={{LIBRECHAT_OPENID_ID_TOKEN}}'; const result = processOpenIDPlaceholders(input, tokenInfo!); - expect(result).toBe('ACCESS=my-access-token, ID=my-id-token'); + expect(result).toBe(`ACCESS=my-access-token, ID=${myIdToken}`); // Verify they are not the same value (the reported bug) expect(result).not.toBe('ACCESS=my-access-token, ID=my-access-token'); }); diff --git a/packages/api/src/utils/oidc.ts b/packages/api/src/utils/oidc.ts index fcf2db247ba..27b7dc5e4a5 100644 --- a/packages/api/src/utils/oidc.ts +++ b/packages/api/src/utils/oidc.ts @@ -5,6 +5,7 @@ export interface OpenIDTokenInfo { accessToken?: string; idToken?: string; expiresAt?: number; + idTokenExpiresAt?: number; userId?: string; userEmail?: string; userName?: string; @@ -40,6 +41,25 @@ export const GRAPH_TOKEN_PLACEHOLDER = '{{LIBRECHAT_GRAPH_ACCESS_TOKEN}}'; */ export const DEFAULT_GRAPH_SCOPES = 'https://graph.microsoft.com/.default'; +/** Shared with AuthController's OpenID session reuse check: a token within the buffer would expire in transit and 401 downstream */ +export const OPENID_EXPIRY_BUFFER_SECONDS = 30; + +/** + * Signals that the stored OpenID credentials cannot satisfy a placeholder, so the user must + * re-authenticate. `ErrorController` maps this to a 401, and `statusCode` additionally lets + * status-reading callers (the agent generation path's `getInitializationFailure`) answer 401 + * instead of a bare 500. Deliberately carries no `body`, so the structural `isCustomError` + * guard cannot capture it ahead of the explicit mapping. + */ +export class OpenIDReauthRequiredError extends Error { + readonly statusCode = 401; + + constructor(message: string) { + super(message); + this.name = 'OpenIDReauthRequiredError'; + } +} + export function extractOpenIDTokenInfo( user: Partial | null | undefined, ): OpenIDTokenInfo | null { @@ -85,10 +105,13 @@ export function extractOpenIDTokenInfo( ); tokenInfo.claims = payload; + /** Cached profile claims, not an authentication assertion: stale claims stay usable for identity fields even when the ID token itself is expired */ if (payload.sub) tokenInfo.userId = payload.sub; if (payload.email) tokenInfo.userEmail = payload.email; if (payload.name) tokenInfo.userName = payload.name; - if (payload.exp) tokenInfo.expiresAt = payload.exp; + if (typeof payload.exp === 'number') { + tokenInfo.idTokenExpiresAt = payload.exp; + } } catch (jwtError) { logger.warn('Could not parse ID token claims:', jwtError); } @@ -101,14 +124,22 @@ export function extractOpenIDTokenInfo( } } +/** Advisory freshness check, not a security boundary: the ID token signature is not verified here. `exp` is REQUIRED in an ID token, so a missing value means the token is malformed or unparseable and fails closed. */ +function isIdTokenCurrent(tokenInfo: OpenIDTokenInfo): boolean { + if (tokenInfo.idTokenExpiresAt == null) { + return false; + } + return Math.floor(Date.now() / 1000) < tokenInfo.idTokenExpiresAt - OPENID_EXPIRY_BUFFER_SECONDS; +} + export function isOpenIDTokenValid(tokenInfo: OpenIDTokenInfo | null): boolean { if (!tokenInfo || !tokenInfo.accessToken) { return false; } - if (tokenInfo.expiresAt) { + if (tokenInfo.expiresAt != null) { const now = Math.floor(Date.now() / 1000); - if (now >= tokenInfo.expiresAt) { + if (now >= tokenInfo.expiresAt - OPENID_EXPIRY_BUFFER_SECONDS) { logger.warn('OpenID token has expired'); return false; } @@ -140,7 +171,13 @@ export function processOpenIDPlaceholders( replacementValue = tokenInfo.accessToken || ''; break; case 'ID_TOKEN': - replacementValue = tokenInfo.idToken || ''; + if (!tokenInfo.idToken || !isIdTokenCurrent(tokenInfo)) { + logger.warn('OpenID ID token is expired or unavailable; re-authentication is required'); + throw new OpenIDReauthRequiredError( + 'OpenID ID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ID_TOKEN}}', + ); + } + replacementValue = tokenInfo.idToken; break; case 'USER_ID': replacementValue = tokenInfo.userId || ''; @@ -152,7 +189,8 @@ export function processOpenIDPlaceholders( replacementValue = tokenInfo.userName || ''; break; case 'EXPIRES_AT': - replacementValue = tokenInfo.expiresAt ? String(tokenInfo.expiresAt) : ''; + /** The stored token-set expires_at only: the ID token exp never stands in for it */ + replacementValue = tokenInfo.expiresAt != null ? String(tokenInfo.expiresAt) : ''; break; } From e4d6bb71f9d58aab8150781066cbdfeb67fe21c4 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 18 Aug 2026 22:05:13 -0400 Subject: [PATCH 04/15] =?UTF-8?q?=F0=9F=93=81=20feat:=20Surface=20Stateful?= =?UTF-8?q?=20Workspace=20Downloads=20(#14984)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: surface stateful workspace downloads * fix: sort workspace change imports * fix: reuse workspace button primitives * fix: hide collapsed workspace actions --- .../agents/__tests__/callbacks.spec.js | 149 +++++++++++++++++- api/server/controllers/agents/callbacks.js | 42 ++++- api/server/routes/files/files.js | 5 +- api/server/routes/files/files.test.js | 12 +- .../Chat/Messages/Content/ContentParts.tsx | 22 ++- .../Content/Parts/WorkspaceChanges.tsx | 142 +++++++++++++++++ .../Parts/__tests__/WorkspaceChanges.test.tsx | 107 +++++++++++++ .../Content/__tests__/ContentParts.test.tsx | 54 +++++++ client/src/locales/en/translation.json | 3 + packages/data-provider/src/schemas.ts | 7 + 10 files changed, 533 insertions(+), 10 deletions(-) create mode 100644 client/src/components/Chat/Messages/Content/Parts/WorkspaceChanges.tsx create mode 100644 client/src/components/Chat/Messages/Content/Parts/__tests__/WorkspaceChanges.test.tsx diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js index d37f3933670..b13eae0f619 100644 --- a/api/server/controllers/agents/__tests__/callbacks.spec.js +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -7,6 +7,7 @@ jest.mock('nanoid', () => ({ jest.mock('@librechat/api', () => ({ sendEvent: jest.fn(), + writeAttachmentEvent: jest.fn(), GenerationJobManager: { emitChunk: jest.fn(), }, @@ -444,6 +445,7 @@ describe('createToolEndCallback', () => { name, toolName = 'execute_code', hostFileAuthoring = false, + created, codeExecutionContext, }) { return { @@ -452,6 +454,8 @@ describe('createToolEndCallback', () => { tool_call_id: toolCallId, artifact: { ...(hostFileAuthoring ? { __librechat_file_authoring: true } : {}), + ...(created === undefined ? {} : { created }), + path: name, session_id: 'sess-1', files: [{ id: fileId, name, session_id: 'sess-1' }], }, @@ -667,8 +671,17 @@ describe('createToolEndCallback', () => { conversationId: 'thread789', messageId: 'run-create', toolCallId: 'tool-create', - status: 'ready', + status: 'pending', }, + finalize: jest.fn().mockResolvedValue({ + file_id: 'fid-created', + filename: 'created.txt', + filepath: '/uploads/created.txt', + type: 'text/plain', + conversationId: 'thread789', + messageId: 'run-create', + status: 'ready', + }), }); const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); @@ -680,6 +693,7 @@ describe('createToolEndCallback', () => { name: 'created.txt', toolName: 'create_file', hostFileAuthoring: true, + created: true, codeExecutionContext: { baseUrl: 'https://code-stateful.example.com', executionProfile: 'stateful', @@ -687,6 +701,7 @@ describe('createToolEndCallback', () => { }); await toolEndCallback({ output: event.output }, event.metadata); await Promise.all(artifactPromises); + await new Promise((resolve) => setImmediate(resolve)); expect(processCodeOutput).toHaveBeenCalledWith( expect.objectContaining({ @@ -699,7 +714,139 @@ describe('createToolEndCallback', () => { executionProfile: 'stateful', }), ); + expect(res.write).toHaveBeenCalledTimes(2); + expect(parseSseAttachment(res.write.mock.calls[0]).workspaceChange).toEqual({ + profile: 'stateful', + operation: 'created', + path: 'created.txt', + }); + expect(parseSseAttachment(res.write.mock.calls[1]).workspaceChange).toEqual({ + profile: 'stateful', + operation: 'created', + path: 'created.txt', + }); + await expect(artifactPromises[0]).resolves.toEqual( + expect.objectContaining({ + workspaceChange: { + profile: 'stateful', + operation: 'created', + path: 'created.txt', + }, + }), + ); + }); + + it('does not mark stateless file authoring outputs as stateful workspace changes', async () => { + res.headersSent = true; + processCodeOutput.mockResolvedValue({ + file: { + file_id: 'fid-default', + filename: 'default.txt', + filepath: '/uploads/default.txt', + type: 'text/plain', + conversationId: 'thread789', + messageId: 'run-default', + toolCallId: 'tool-default', + status: 'ready', + }, + }); + + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + const event = makeCodeExecutionEvent({ + runId: 'run-default', + threadId: 'thread789', + toolCallId: 'tool-default', + fileId: 'fid-default', + name: 'default.txt', + toolName: 'create_file', + hostFileAuthoring: true, + created: true, + codeExecutionContext: { + baseUrl: 'https://code-default.example.com', + executionProfile: 'default', + }, + }); + await toolEndCallback({ output: event.output }, event.metadata); + await Promise.all(artifactPromises); + expect(res.write).toHaveBeenCalledTimes(1); + expect(parseSseAttachment(res.write.mock.calls[0]).workspaceChange).toBeUndefined(); + }); + + it('preserves stateful workspace changes in Open Responses attachment events', async () => { + const { writeAttachmentEvent } = require('@librechat/api'); + const { createResponsesToolEndCallback } = require('../callbacks'); + res.headersSent = true; + res.writableEnded = false; + processCodeOutput.mockResolvedValue({ + file: { + file_id: 'fid-responses', + filename: 'summary.csv', + filepath: '/uploads/summary.csv', + type: 'text/csv', + conversationId: 'thread789', + messageId: 'run-responses', + toolCallId: 'tool-responses', + status: 'pending', + }, + finalize: jest.fn().mockResolvedValue({ + file_id: 'fid-responses', + filename: 'summary.csv', + filepath: '/uploads/summary.csv', + type: 'text/csv', + conversationId: 'thread789', + messageId: 'run-responses', + status: 'ready', + }), + }); + + const tracker = { nextSequence: jest.fn().mockReturnValueOnce(1).mockReturnValueOnce(2) }; + const toolEndCallback = createResponsesToolEndCallback({ + req, + res, + tracker, + artifactPromises, + }); + const event = makeCodeExecutionEvent({ + runId: 'run-responses', + threadId: 'thread789', + toolCallId: 'tool-responses', + fileId: 'fid-responses', + name: 'summary.csv', + toolName: 'edit_file', + hostFileAuthoring: true, + created: false, + codeExecutionContext: { + baseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', + }, + }); + event.output.artifact.path = 'reports/summary.csv'; + + await toolEndCallback({ output: event.output }, event.metadata); + await Promise.all(artifactPromises); + await new Promise((resolve) => setImmediate(resolve)); + + expect(writeAttachmentEvent).toHaveBeenCalledTimes(2); + expect(writeAttachmentEvent.mock.calls[0][2].workspaceChange).toEqual({ + profile: 'stateful', + operation: 'updated', + path: 'reports/summary.csv', + }); + expect(writeAttachmentEvent.mock.calls[1][2].workspaceChange).toEqual({ + profile: 'stateful', + operation: 'updated', + path: 'reports/summary.csv', + }); + await expect(artifactPromises[0]).resolves.toEqual( + expect.objectContaining({ + workspaceChange: { + profile: 'stateful', + operation: 'updated', + path: 'reports/summary.csv', + }, + }), + ); }); it('does not process arbitrary user tool artifacts named create_file as code outputs', async () => { diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index cb72f2ae37d..ae9683b3e7a 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -40,6 +40,25 @@ function isCodeArtifactToolOutput(output) { return isCodeSessionToolName(output.name) || isHostFileAuthoringArtifact(output.artifact); } +function addStatefulWorkspaceChange(attachment, artifact, executionProfile) { + if (!attachment || executionProfile !== 'stateful' || !isHostFileAuthoringArtifact(artifact)) { + return attachment; + } + const path = + typeof artifact.path === 'string' && artifact.path.length > 0 + ? artifact.path + : attachment.filename; + if (typeof path !== 'string' || path.length === 0) { + return attachment; + } + attachment.workspaceChange = { + profile: 'stateful', + operation: artifact.created === true ? 'created' : 'updated', + path, + }; + return attachment; +} + class ModelEndHandler { /** * @param {Array} collectedUsage @@ -979,7 +998,11 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl, executionProfile: metadata.codeExecutionContext?.executionProfile, }); - const fileMetadata = result?.file ?? null; + const fileMetadata = addStatefulWorkspaceChange( + result?.file ?? null, + output.artifact, + metadata.codeExecutionContext?.executionProfile, + ); const finalize = result?.finalize; if (!fileMetadata) { return null; @@ -1027,6 +1050,9 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo ...updated, messageId: metadata.run_id, toolCallId, + ...(fileMetadata.workspaceChange + ? { workspaceChange: fileMetadata.workspaceChange } + : {}), }, jobCreatedAt, ); @@ -1303,7 +1329,11 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises }) codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl, executionProfile: metadata.codeExecutionContext?.executionProfile, }); - const fileMetadata = result?.file ?? null; + const fileMetadata = addStatefulWorkspaceChange( + result?.file ?? null, + output.artifact, + metadata.codeExecutionContext?.executionProfile, + ); const finalize = result?.finalize; if (!fileMetadata) { return null; @@ -1336,7 +1366,12 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises }) writeResponsesAttachment( res, tracker, - buildResponsesAttachment(updated, toolCallId), + buildResponsesAttachment( + fileMetadata.workspaceChange + ? { ...updated, workspaceChange: fileMetadata.workspaceChange } + : updated, + toolCallId, + ), metadata, ); }, @@ -1371,6 +1406,7 @@ function buildResponsesAttachment(fileMetadata, toolCallId) { textFormat: fileMetadata.textFormat ?? null, status: fileMetadata.status, previewError: fileMetadata.previewError, + workspaceChange: fileMetadata.workspaceChange, }; } diff --git a/api/server/routes/files/files.js b/api/server/routes/files/files.js index 72295c4a9b8..4f0caee52a4 100644 --- a/api/server/routes/files/files.js +++ b/api/server/routes/files/files.js @@ -367,7 +367,10 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => { req, { baseUrl, executionProfile }, ); - res.set(response.headers); + res.setHeader('Content-Disposition', 'attachment'); + res.setHeader('Content-Type', 'application/octet-stream'); + res.setHeader('X-Content-Type-Options', 'nosniff'); + res.setHeader('Cache-Control', 'private, no-store'); response.data.pipe(res); } catch (error) { /* `logAxiosError` redacts buffer/stream response bodies — without diff --git a/api/server/routes/files/files.test.js b/api/server/routes/files/files.test.js index 4fce9cb30e7..4fcdd3a62a1 100644 --- a/api/server/routes/files/files.test.js +++ b/api/server/routes/files/files.test.js @@ -1097,7 +1097,10 @@ describe('File Routes - Delete with Agent Access', () => { describe('GET /files/code/download/:session_id/:fileId', () => { it('routes a persisted stateful fallback through the stateful Code API', async () => { const getDownloadStream = jest.fn().mockResolvedValue({ - headers: { 'content-type': 'text/plain' }, + headers: { + 'content-type': 'text/html', + 'set-cookie': 'internal-service-cookie=secret', + }, data: Readable.from(['stateful output']), }); getStrategyFunctions.mockReturnValue({ getDownloadStream }); @@ -1111,7 +1114,12 @@ describe('File Routes - Delete with Agent Access', () => { ); expect(response.status).toBe(200); - expect(response.text).toBe('stateful output'); + expect(response.body.toString()).toBe('stateful output'); + expect(response.headers['content-disposition']).toBe('attachment'); + expect(response.headers['content-type']).toBe('application/octet-stream'); + expect(response.headers['x-content-type-options']).toBe('nosniff'); + expect(response.headers['cache-control']).toBe('private, no-store'); + expect(response.headers['set-cookie']).toBeUndefined(); expect(getDownloadStream).toHaveBeenCalledWith( `${sessionId}/${codeFileId}`, { kind: 'user', id: otherUserId.toString() }, diff --git a/client/src/components/Chat/Messages/Content/ContentParts.tsx b/client/src/components/Chat/Messages/Content/ContentParts.tsx index 0afa4ab1be9..a09a9ff5bd7 100644 --- a/client/src/components/Chat/Messages/Content/ContentParts.tsx +++ b/client/src/components/Chat/Messages/Content/ContentParts.tsx @@ -9,6 +9,7 @@ import type { import type { ReactNode, ReactElement } from 'react'; import type { ToolCallGroupExpansionState } from './ToolCallGroup'; import { mapAttachments, filterAttachmentsForPart, groupSequentialToolCalls } from '~/utils'; +import WorkspaceChanges, { partitionWorkspaceChanges } from './Parts/WorkspaceChanges'; import { groupActivityPhases, lastVisibleContentIdx } from '~/utils/activityLabels'; import { ParallelContentRenderer, type PartWithIndex } from './ParallelContent'; import MemoryArtifacts, { hasMemoryArtifacts } from './MemoryArtifacts'; @@ -163,6 +164,8 @@ type ContentPartsProps = { | undefined; /** Internal recursion guard for nested phase segments. */ nestedActivityPhase?: boolean; + /** Internal signal that the parent already removed message-level workspace attachments. */ + workspaceAttachmentsPartitioned?: boolean; /** Absolute transcript index represented by `content[0]` in a phase slice. */ contentIndexOffset?: number; /** Absolute transcript index for each compacted sparse segment entry. */ @@ -197,12 +200,20 @@ const ContentParts = memo(function ContentParts({ isLatestMessage, createdAt, nestedActivityPhase = false, + workspaceAttachmentsPartitioned = false, contentIndexOffset = 0, contentIndices, resumeAuthors, toolGroupExpansionState, }: ContentPartsProps) { - const attachmentMap = useMemo(() => mapAttachments(attachments ?? []), [attachments]); + const { inlineAttachments, workspaceChanges } = useMemo( + () => + workspaceAttachmentsPartitioned + ? { inlineAttachments: attachments ?? [], workspaceChanges: [] } + : partitionWorkspaceChanges(attachments), + [attachments, workspaceAttachmentsPartitioned], + ); + const attachmentMap = useMemo(() => mapAttachments(inlineAttachments), [inlineAttachments]); const effectiveIsSubmitting = isLatestMessage ? isSubmitting : false; const localToolGroupExpansionRef = useRef(new Map()); const expansionState = toolGroupExpansionState ?? localToolGroupExpansionRef.current; @@ -459,7 +470,7 @@ const ContentParts = memo(function ContentParts({ ); // Early return: no content to render AND no pending skill cards - if (!content && !hasPendingSkills) { + if (!content && !hasPendingSkills && workspaceChanges.length === 0) { return null; } @@ -479,6 +490,7 @@ const ContentParts = memo(function ContentParts({ setSiblingIdx={setSiblingIdx} renderReadOnlyPart={(part, idx, isLastPart) => renderPart(part, idx, isLastPart)} /> + ); @@ -502,13 +514,14 @@ const ContentParts = memo(function ContentParts({ createdAt={createdAt} authorHeader={authorHeader} conversationId={conversationId} - attachments={attachments} + attachments={inlineAttachments} searchResults={searchResults} isCreatedByUser={isCreatedByUser} isLast={isLast && segmentIndices.includes(globalLastContentIdx)} isSubmitting={isSubmitting} isLatestMessage={isLatestMessage} nestedActivityPhase + workspaceAttachmentsPartitioned contentIndexOffset={segmentStartIndex} contentIndices={segmentIndices} resumeAuthors={postSteerAuthors} @@ -559,6 +572,7 @@ const ContentParts = memo(function ContentParts({ ) ), )} + ); @@ -598,6 +612,7 @@ const ContentParts = memo(function ContentParts({ contentIndexOffset={contentIndexOffset} contentIndices={contentIndices} /> + {!nestedActivityPhase && } ); return nestedActivityPhase ? ( @@ -660,6 +675,7 @@ const ContentParts = memo(function ContentParts({ ); return nodes; })} + {!nestedActivityPhase && } ); if (nestedActivityPhase) { diff --git a/client/src/components/Chat/Messages/Content/Parts/WorkspaceChanges.tsx b/client/src/components/Chat/Messages/Content/Parts/WorkspaceChanges.tsx new file mode 100644 index 00000000000..d49cc137b0f --- /dev/null +++ b/client/src/components/Chat/Messages/Content/Parts/WorkspaceChanges.tsx @@ -0,0 +1,142 @@ +import { memo, useId, useMemo, useState } from 'react'; +import { Button, IconButton } from '@librechat/client'; +import { ChevronDown, Download, Files } from 'lucide-react'; +import type { + TAttachment, + TFile, + WorkspaceChange as WorkspaceChangeMetadata, +} from 'librechat-data-provider'; +import { useExpandCollapse, useLocalize } from '~/hooks'; +import { useAttachmentLink } from './LogLink'; +import { cn } from '~/utils'; + +type StatefulWorkspaceAttachment = TAttachment & { + workspaceChange: WorkspaceChangeMetadata; +}; + +export function partitionWorkspaceChanges(attachments?: TAttachment[]): { + inlineAttachments: TAttachment[]; + workspaceChanges: StatefulWorkspaceAttachment[]; +} { + const inlineAttachments: TAttachment[] = []; + const changesByFile = new Map(); + + for (const attachment of attachments ?? []) { + const change = attachment.workspaceChange; + if (change?.profile !== 'stateful' || !attachment.filepath) { + inlineAttachments.push(attachment); + continue; + } + + const file = attachment as Partial & { agentId?: string }; + const key = file.file_id ?? `${file.agentId ?? ''}:${change.path}`; + changesByFile.delete(key); + changesByFile.set(key, attachment as StatefulWorkspaceAttachment); + } + + return { inlineAttachments, workspaceChanges: Array.from(changesByFile.values()) }; +} + +const WorkspaceChange = memo(({ attachment }: { attachment: StatefulWorkspaceAttachment }) => { + const localize = useLocalize(); + const file = attachment as TFile; + const path = attachment.workspaceChange.path; + const filename = path.split('/').pop() || path; + const { handleDownload } = useAttachmentLink({ + href: attachment.filepath ?? '', + filename, + file_id: file.file_id, + user: file.user, + source: file.source, + }); + + return ( +
+
+
+ {filename} +
+ {path !== filename && ( +
+ {path} +
+ )} +
+ void handleDownload(event)} + label={`${localize('com_ui_download')} ${filename}`} + title={localize('com_ui_download')} + variant="ghost" + size="sm" + shape="square" + className="text-text-secondary" + > + +
+ ); +}); + +WorkspaceChange.displayName = 'WorkspaceChange'; + +export default function WorkspaceChanges({ + attachments, +}: { + attachments: StatefulWorkspaceAttachment[]; +}) { + const localize = useLocalize(); + const panelId = useId(); + const [isExpanded, setIsExpanded] = useState(false); + const { style, ref } = useExpandCollapse(isExpanded); + const count = attachments.length; + const countLabel = localize(count === 1 ? 'com_ui_one_file_changed' : 'com_ui_n_files_changed', { + 0: String(count), + }); + const summary = useMemo( + () => attachments.map((attachment) => attachment.workspaceChange.path).join(', '), + [attachments], + ); + + if (count === 0) { + return null; + } + + return ( +
+ +
+
+
+ {attachments.map((attachment) => ( + + ))} +
+
+
+
+ ); +} diff --git a/client/src/components/Chat/Messages/Content/Parts/__tests__/WorkspaceChanges.test.tsx b/client/src/components/Chat/Messages/Content/Parts/__tests__/WorkspaceChanges.test.tsx new file mode 100644 index 00000000000..4566e2d0318 --- /dev/null +++ b/client/src/components/Chat/Messages/Content/Parts/__tests__/WorkspaceChanges.test.tsx @@ -0,0 +1,107 @@ +import React from 'react'; +import { FileSources } from 'librechat-data-provider'; +import { fireEvent, render, screen } from '@testing-library/react'; +import type { TAttachment } from 'librechat-data-provider'; +import WorkspaceChanges, { partitionWorkspaceChanges } from '../WorkspaceChanges'; + +const mockHandleDownload = jest.fn(); + +jest.mock('../LogLink', () => ({ + useAttachmentLink: () => ({ handleDownload: mockHandleDownload }), +})); + +jest.mock('~/hooks', () => ({ + useExpandCollapse: () => ({ style: {}, ref: { current: null } }), + useLocalize: () => (key: string, values?: Record) => { + const translations: Record = { + com_ui_download: 'Download', + com_ui_n_files_changed: `${values?.[0]} files changed`, + com_ui_one_file_changed: '1 file changed', + com_ui_workspace_changes: 'Workspace changes', + }; + return translations[key] ?? key; + }, +})); + +jest.mock('~/utils', () => ({ + cn: (...classes: Array) => classes.filter(Boolean).join(' '), +})); + +function makeAttachment({ + fileId, + path, + profile = 'stateful', + filepath, +}: { + fileId: string; + path: string; + profile?: 'stateful' | 'default'; + filepath?: string; +}): TAttachment { + return { + file_id: fileId, + filename: path, + filepath: filepath ?? `/uploads/${fileId}`, + source: FileSources.local, + user: 'user-1', + conversationId: 'conversation-1', + messageId: 'message-1', + toolCallId: `tool-${fileId}`, + workspaceChange: { + profile, + operation: 'updated', + path, + }, + } as TAttachment; +} + +describe('WorkspaceChanges', () => { + beforeEach(() => { + mockHandleDownload.mockReset(); + }); + + it('partitions only downloadable stateful changes and keeps the latest file entry', () => { + const first = makeAttachment({ fileId: 'shared', path: 'reports/result.csv' }); + const latest = { + ...makeAttachment({ fileId: 'shared', path: 'reports/result.csv' }), + filepath: '/uploads/latest', + } as TAttachment; + const stateless = makeAttachment({ + fileId: 'default', + path: 'default.txt', + profile: 'default', + }); + const unavailable = makeAttachment({ fileId: 'missing', path: 'missing.txt' }); + unavailable.filepath = ''; + + const result = partitionWorkspaceChanges([first, stateless, unavailable, latest]); + + expect(result.inlineAttachments).toEqual([stateless, unavailable]); + expect(result.workspaceChanges).toEqual([latest]); + }); + + it('renders one collapsed row and downloads through the existing attachment handler', () => { + const changes = partitionWorkspaceChanges([ + makeAttachment({ fileId: 'one', path: 'reports/summary.csv' }), + makeAttachment({ fileId: 'two', path: 'notes.txt' }), + ]).workspaceChanges; + + render(); + + const toggle = screen.getByRole('button', { + name: 'Workspace changes: 2 files changed', + }); + const panel = document.getElementById(toggle.getAttribute('aria-controls') ?? ''); + expect(toggle).toHaveAttribute('aria-expanded', 'false'); + expect(panel).toHaveAttribute('inert'); + + fireEvent.click(toggle); + expect(toggle).toHaveAttribute('aria-expanded', 'true'); + expect(panel).not.toHaveAttribute('inert'); + expect(screen.getByText('summary.csv')).toBeInTheDocument(); + expect(screen.getByText('reports/summary.csv')).toBeInTheDocument(); + + fireEvent.click(screen.getByRole('button', { name: 'Download summary.csv' })); + expect(mockHandleDownload).toHaveBeenCalledTimes(1); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/__tests__/ContentParts.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/ContentParts.test.tsx index b735bc04937..f4e1fbc3bc1 100644 --- a/client/src/components/Chat/Messages/Content/__tests__/ContentParts.test.tsx +++ b/client/src/components/Chat/Messages/Content/__tests__/ContentParts.test.tsx @@ -54,6 +54,18 @@ jest.mock('../Parts/PendingSkillCall', () => ({ ), })); +jest.mock('../Parts/WorkspaceChanges', () => ({ + __esModule: true, + default: ({ attachments }: { attachments: TAttachment[] }) => + attachments.length > 0 ? ( +
+ ) : null, + partitionWorkspaceChanges: (attachments?: TAttachment[]) => ({ + inlineAttachments: (attachments ?? []).filter((attachment) => !attachment.workspaceChange), + workspaceChanges: (attachments ?? []).filter((attachment) => attachment.workspaceChange), + }), +})); + jest.mock('../ToolCallGroup', () => ({ __esModule: true, default: ({ @@ -144,6 +156,48 @@ beforeEach(() => { }); describe('ContentParts — interim skill cards', () => { + it('renders stateful workspace changes once at message level', () => { + const content: TMessageContentParts[] = [ + { type: ContentTypes.TEXT, text: 'done' } as TMessageContentParts, + ]; + const attachment = { + filename: 'report.csv', + filepath: '/uploads/report.csv', + conversationId: 'conversation-1', + messageId: 'msg-1', + toolCallId: 'tool-1', + workspaceChange: { + profile: 'stateful', + operation: 'created', + path: 'report.csv', + }, + } as TAttachment; + + render(); + + expect(screen.getAllByTestId('workspace-changes')).toHaveLength(1); + expect(screen.getByTestId('workspace-changes')).toHaveAttribute('data-count', '1'); + }); + + it('renders stateful workspace changes when the assistant message has no content yet', () => { + const attachment = { + filename: 'report.csv', + filepath: '/uploads/report.csv', + conversationId: 'conversation-1', + messageId: 'msg-1', + toolCallId: 'tool-1', + workspaceChange: { + profile: 'stateful', + operation: 'created', + path: 'report.csv', + }, + } as TAttachment; + + render(); + + expect(screen.getByTestId('workspace-changes')).toHaveAttribute('data-count', '1'); + }); + it('renders a PendingSkillCall per manual skill on assistant messages', () => { render(); const cards = screen.getAllByTestId('pending-skill-call'); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index f4fa8311205..3f2405e8a53 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1595,6 +1595,7 @@ "com_ui_my_prompts": "My Prompts", "com_ui_my_skills": "My Skills", "com_ui_n_files": "{{0}} files", + "com_ui_n_files_changed": "{{0}} files changed", "com_ui_name": "Name", "com_ui_name_sort": "Sort by Name", "com_ui_navigate_results": "Navigate results", @@ -1658,6 +1659,7 @@ "com_ui_offline": "Offline", "com_ui_omitted": "Omitted", "com_ui_on": "On", + "com_ui_one_file_changed": "1 file changed", "com_ui_open_archived_chat_new_tab_title": "{{title}} (opens in new tab)", "com_ui_open_artifact": "Open artifact", "com_ui_open_as_artifact": "Open as artifact", @@ -2294,6 +2296,7 @@ "com_ui_web_searched": "Searched the web", "com_ui_web_searching": "Searching the web", "com_ui_web_searching_again": "Searching the web again", + "com_ui_workspace_changes": "Workspace changes", "com_ui_write": "Writing", "com_ui_writing_command": "Writing command", "com_ui_x_selected": "{{0}} selected", diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 97fbd700d2a..ba40501687a 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -895,10 +895,17 @@ export type UIResource = { [key: string]: unknown; }; +export type WorkspaceChange = { + profile: 'stateful'; + operation: 'created' | 'updated'; + path: string; +}; + export type TAttachmentMetadata = { type?: Tools; messageId: string; toolCallId: string; + workspaceChange?: WorkspaceChange; [Tools.memory]?: MemoryArtifact; [Tools.ui_resources]?: UIResource[]; [Tools.web_search]?: SearchResultData; From 5f956312835230c8265d6ab9f66657689cf816d5 Mon Sep 17 00:00:00 2001 From: Ravi Kumar L Date: Wed, 19 Aug 2026 04:06:33 +0200 Subject: [PATCH 05/15] =?UTF-8?q?=F0=9F=AA=A2=20fix:=20Harden=20Langfuse?= =?UTF-8?q?=20Media=20Upload=20Targets=20(#14974)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../cmd/langfuse-fanout/main.go | 26 +++++- .../cmd/langfuse-fanout/main_test.go | 81 +++++++++++++++++-- 2 files changed, 98 insertions(+), 9 deletions(-) diff --git a/otel/langfuse-fanout/cmd/langfuse-fanout/main.go b/otel/langfuse-fanout/cmd/langfuse-fanout/main.go index a1fcf5e6ec2..d85e696c83c 100644 --- a/otel/langfuse-fanout/cmd/langfuse-fanout/main.go +++ b/otel/langfuse-fanout/cmd/langfuse-fanout/main.go @@ -658,6 +658,9 @@ func (g *gateway) patchMedia(ctx context.Context, dest destination, path string, } func (g *gateway) putMedia(ctx context.Context, dest uploadDestination, body []byte, originalHeaders http.Header) (int, error) { + if err := validateMediaUploadURL(dest.UploadURL); err != nil { + return 0, err + } req, err := http.NewRequestWithContext(ctx, http.MethodPut, dest.UploadURL, bytes.NewReader(body)) if err != nil { return 0, err @@ -680,7 +683,11 @@ func (g *gateway) putMedia(ctx context.Context, dest uploadDestination, body []b req.Header.Set("x-amz-checksum-sha256", value) } } - resp, err := g.doUpstream(req, "media_upload", dest.Name) + uploadClient := *g.cfg.client + uploadClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + } + resp, err := g.doUpstreamWithClient(&uploadClient, req, "media_upload", dest.Name) if err != nil { return 0, err } @@ -706,8 +713,12 @@ func (g *gateway) doExpect2xx(operation string, destination string, req *http.Re } func (g *gateway) doUpstream(req *http.Request, operation string, destination string) (*http.Response, error) { + return g.doUpstreamWithClient(g.cfg.client, req, operation, destination) +} + +func (g *gateway) doUpstreamWithClient(client *http.Client, req *http.Request, operation string, destination string) (*http.Response, error) { startedAt := time.Now() - resp, err := g.cfg.client.Do(req) + resp, err := client.Do(req) if err != nil { duration := time.Since(startedAt) if g.metrics != nil { @@ -1141,6 +1152,17 @@ func isGCSUploadURL(value string) bool { return host == "storage.googleapis.com" || strings.HasSuffix(host, ".storage.googleapis.com") } +func validateMediaUploadURL(value string) error { + parsed, err := url.Parse(value) + if err != nil || parsed.Hostname() == "" { + return errors.New("media upload URL must be an absolute HTTPS URL") + } + if parsed.Scheme != "https" { + return errors.New("media upload URL must use HTTPS") + } + return nil +} + func isAzureUploadURL(value string) bool { parsed, err := url.Parse(value) if err != nil { diff --git a/otel/langfuse-fanout/cmd/langfuse-fanout/main_test.go b/otel/langfuse-fanout/cmd/langfuse-fanout/main_test.go index fac8edbb45c..f71b0da5c9a 100644 --- a/otel/langfuse-fanout/cmd/langfuse-fanout/main_test.go +++ b/otel/langfuse-fanout/cmd/langfuse-fanout/main_test.go @@ -99,6 +99,66 @@ func TestNormalizeBaseURLAllowsOnlyHTTPAndHTTPS(t *testing.T) { } } +func TestValidateMediaUploadURLRequiresHTTPS(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + url string + wantErr bool + }{ + {name: "public storage", url: "https://bucket.s3.amazonaws.com/object?X-Amz-Signature=value"}, + {name: "self-hosted storage", url: "https://minio.internal:9000/object"}, + {name: "private address", url: "https://10.0.0.8/object"}, + {name: "http", url: "http://minio.internal:9000/object", wantErr: true}, + {name: "unsupported scheme", url: "ftp://storage.example.com/object", wantErr: true}, + {name: "relative", url: "/object", wantErr: true}, + {name: "missing host", url: "https:///object", wantErr: true}, + {name: "malformed", url: "://storage.example.com/object", wantErr: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := validateMediaUploadURL(test.url) + if (err != nil) != test.wantErr { + t.Fatalf("validateMediaUploadURL(%q) error = %v, wantErr %t", test.url, err, test.wantErr) + } + }) + } +} + +func TestPutMediaDoesNotFollowRedirects(t *testing.T) { + t.Parallel() + + var targetRequests int + target := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + targetRequests++ + w.WriteHeader(http.StatusOK) + })) + defer target.Close() + + redirect := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, target.URL+"/upload", http.StatusTemporaryRedirect) + })) + defer redirect.Close() + + gw := newTestGateway(redirect.URL, nil) + gw.cfg.client = redirect.Client() + status, err := gw.putMedia(context.Background(), uploadDestination{ + Name: centralName, + UploadURL: redirect.URL + "/upload", + }, []byte("hello"), http.Header{"Content-Type": []string{"image/png"}}) + if status != http.StatusTemporaryRedirect { + t.Fatalf("status = %d, want %d", status, http.StatusTemporaryRedirect) + } + if err == nil { + t.Fatal("expected redirect response to fail the upload") + } + if targetRequests != 0 { + t.Fatalf("redirect target requests = %d, want 0", targetRequests) + } +} + func TestTraceProxyForwardsExistingRoutingAttributesToCollector(t *testing.T) { t.Parallel() @@ -290,10 +350,10 @@ func TestMediaUploadFansOutToCentralAndTenant(t *testing.T) { var mu sync.Mutex uploads := map[string]string{} upstream := func(name string) *httptest.Server { - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + return httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case r.Method == http.MethodPost && r.URL.Path == mediaPath: - uploadURL := "http://" + r.Host + "/upload/" + name + uploadURL := "https://" + r.Host + "/upload/" + name writeJSON(w, http.StatusCreated, mediaUploadResponse{ MediaID: "same-media-id", UploadURL: &uploadURL, @@ -319,6 +379,8 @@ func TestMediaUploadFansOutToCentralAndTenant(t *testing.T) { store := newFakeUploadPlanStore() createGateway := newTestGatewayWithStore(central.URL, map[string]string{"eu": tenant.URL}, store) uploadGateway := newTestGatewayWithStore(central.URL, map[string]string{"eu": tenant.URL}, store) + createGateway.cfg.client = central.Client() + uploadGateway.cfg.client = central.Client() createBody := `{"traceId":"trace","contentType":"image/png","contentLength":5,"sha256Hash":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","field":"input"}` req := httptest.NewRequest(http.MethodPost, tenantPrefix+"eu"+mediaPath, strings.NewReader(createBody)) req.Header.Set("Authorization", "Basic tenant") @@ -376,10 +438,10 @@ func TestMediaUploadSkipsCentralForCentralMediaDisabledTenantRoute(t *testing.T) var mu sync.Mutex uploads := map[string]string{} upstream := func(name string) *httptest.Server { - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + return httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case r.Method == http.MethodPost && r.URL.Path == mediaPath: - uploadURL := "http://" + r.Host + "/upload/" + name + uploadURL := "https://" + r.Host + "/upload/" + name writeJSON(w, http.StatusCreated, mediaUploadResponse{ MediaID: "same-media-id", UploadURL: &uploadURL, @@ -405,6 +467,8 @@ func TestMediaUploadSkipsCentralForCentralMediaDisabledTenantRoute(t *testing.T) store := newFakeUploadPlanStore() createGateway := newTestGatewayWithStore(central.URL, map[string]string{"eu": tenant.URL}, store) uploadGateway := newTestGatewayWithStore(central.URL, map[string]string{"eu": tenant.URL}, store) + createGateway.cfg.client = central.Client() + uploadGateway.cfg.client = central.Client() createBody := `{"traceId":"trace","contentType":"image/png","contentLength":5,"sha256Hash":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","field":"input"}` req := httptest.NewRequest(http.MethodPost, tenantPrefix+"eu/"+centralMediaDisabled+mediaPath, strings.NewReader(createBody)) req.Header.Set("Authorization", "Basic tenant") @@ -588,7 +652,7 @@ func TestMediaUploadIsOneTime(t *testing.T) { t.Parallel() var uploads int - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPut || r.URL.Path != "/upload" { http.NotFound(w, r) return @@ -609,6 +673,7 @@ func TestMediaUploadIsOneTime(t *testing.T) { }}, }) gw := newTestGatewayWithStore(upstream.URL, nil, store) + gw.cfg.client = upstream.Client() for index, expectedStatus := range []int{http.StatusOK, http.StatusNotFound} { req := httptest.NewRequest(http.MethodPut, mediaUploadProxyPath+uploadID, strings.NewReader("hello")) @@ -628,7 +693,7 @@ func TestMediaUploadOversizeRestoresPlanForRetry(t *testing.T) { t.Parallel() var uploads int - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPut || r.URL.Path != "/upload" { http.NotFound(w, r) return @@ -649,6 +714,7 @@ func TestMediaUploadOversizeRestoresPlanForRetry(t *testing.T) { }}, }) gw := newTestGatewayWithStore(upstream.URL, nil, store) + gw.cfg.client = upstream.Client() oversizeReq := httptest.NewRequest( http.MethodPut, @@ -678,7 +744,7 @@ func TestMediaUploadUnsupportedContentTypeRestoresPlanForRetry(t *testing.T) { t.Parallel() var uploads int - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPut || r.URL.Path != "/upload" { http.NotFound(w, r) return @@ -699,6 +765,7 @@ func TestMediaUploadUnsupportedContentTypeRestoresPlanForRetry(t *testing.T) { }}, }) gw := newTestGatewayWithStore(upstream.URL, nil, store) + gw.cfg.client = upstream.Client() badReq := httptest.NewRequest(http.MethodPut, mediaUploadProxyPath+uploadID, strings.NewReader("hello")) badReq.Header.Set("Content-Type", "text/html") From 259f1e0c3232fe6f2c3486e6eb2761b615ee4c95 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 02:16:35 -0400 Subject: [PATCH 06/15] =?UTF-8?q?=F0=9F=9B=B0=EF=B8=8F=20feat:=20Route=20L?= =?UTF-8?q?ive=20Subagent=20Controls=20Across=20Replicas=20(#14971)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: route live subagent controls across replicas * fix: initialize task routing in cluster workers * fix: harden cross-replica task routing * fix: expire routed task owners independently * fix: close cross-replica routing edge cases * fix: bound owner refresh and close routed cancellation gaps Refresh owned task registrations in bounded parallel batches so a full heartbeat pass stays well inside the 30-second directory lease instead of serializing one Redis EVAL per registration. Route conversation-deletion cancellation through a dedicated owner-side scope operation. The owner applies the deletion predicate to its complete local task set, so a scope holding more children than the model-facing list cap no longer leaves live executors running after their parent is removed. Key a consumed claim's retained response by its operation rather than by one caller's correlation id, so a later poll recovers a terminal result whose responses were all lost. Live claim statuses stay uncached so a poll always observes the task's current state. Type the model-facing `maxLength` bounds with a narrow local string schema; the SDK's JsonSchemaType does not declare the keyword, and the runtime checks continue to enforce the same limits. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: retain claimed results apart from control replays A consumed claim is the only routed response whose loss destroys data, so it no longer shares one bounded cache with control replays that unrelated command traffic can evict. Claims are retained under their own budget, and the requester acknowledges a result it received so the owner releases the copy immediately instead of holding it for the full replay window. Resolve the post-delete cancellation pass from durable leases. The deleted conversations cannot be read back, so re-reading each one only scaled the cascade while probing the owner directory once per removed id; one lease read now resolves every live child address instead. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: never consume a result the owner cannot replay Retention for consumed claims is bounded, so a burst of undelivered results could evict an earlier one and lose it for good. The owner now admits a claim only while it can retain a worst-case result, and refuses the routed claim otherwise instead of consuming it, leaving the result on the task for a later poll. Retained claims are never displaced; control replays keep evicting. Key a control replay by the command itself rather than by one caller's correlation id. The transport's own retry reuses a single envelope, but a caller that saw the owner as unavailable reissues the command under a new id, which steered, queued, or interrupted the child a second time. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: own a claimed result until it is acknowledged A consumed terminal result is task-owned state, not a cache entry. It now carries no expiry at all: the owner holds it until a caller acknowledges receipt, and only then is it released. Retention stays bounded by the existing admission gate, which refuses a claim the owner could not keep rather than consuming a result it might drop. Identify a control by the caller's invocation instead of by its content. The tool mints one id per invocation and routing carries it, so a routed retransmission of that invocation replays the owner's result while two deliberate identical commands arrive under distinct ids and both apply. Content-derived identity could not tell those apart and would have answered the second from a stale snapshot. Wait for the dpkg frontend lock in the best-effort Playwright font step. Its timeout kills npx while the apt-get it spawned keeps the lock, which then failed the fatal Redis install and ended the MCP replica jobs before any test ran (#14983). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: treat acknowledgement as part of delivering a result Publishing an acknowledgement once and ignoring the outcome meant a result could be reported as delivered while the owner never learned it could let go, and since that retention neither expires nor evicts, enough lost acknowledgements would fill it and refuse every later remote claim. An acknowledgement is now confirmed: publishing to zero subscribers is not success, it retries inside the ordinary request window, and a claim whose acknowledgement cannot be confirmed reports the retryable unavailable path instead of handing back a result the owner still holds. A later poll recovers that result and acknowledges it, and releasing is idempotent. Owner registration also outlives the task while a result is unacknowledged, so the retained result cannot become unreachable. Take the control invocation identity from the provider's tool-call id rather than minting one per execution, so replaying the same tool call stays idempotent while two distinct calls with identical payloads both apply. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * style: sort the widened node:crypto import Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: own control invocations and cancellation plans at the task seam Applies one logical control exactly once for its owning task rather than in the transport, so a local caller and a routed caller of the same invocation agree, and reusing an invocation id for different content is refused instead of silently applied. Invocation identity now comes from the run, agent, and provider tool-call id hashed to a bounded 32 characters, so a repeated `call_0` never bleeds across tasks and no id can overrun the routed bound. Cancellation for conversation deletion is now resolved into a plan while those rows are still readable, then replayed against the owner directory after the cascade is deleted. Owner registration is awaited before any provider work, so a child that cannot be addressed never starts. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * style: separate the control invocation map from the next member Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: close subagent deletion, claim, and control invocation gaps Bulk conversation deletion now runs behind a durable owner admission fence. Draining alone could not close the race: a child admitted on another replica after the drain read its leases would start provider work against a parent about to disappear. The fence is written before any lease is read and each child revalidates it after its own lease is written, so one of the two always observes the other. It expires on its own, so a process lost mid-deletion cannot leave an account unable to run subagents. A terminal child result is no longer kept alive in the owning replica's memory until someone acknowledges it. Collection is recorded durably on the child's own message against the polling invocation, so the poll whose response was lost recovers its own result while a different invocation is told the result was already collected. Owner-side retention returns to an ordinary bounded cache that expires, which is what abandoned polls needed: they can no longer occupy claim capacity until the process restarts. The deletion drain now cancels each task under one invocation held for the whole drain, stops re-sending once the owner answers, and retries only deliveries it could not confirm. A routed control replay also validates the command fingerprint, so one invocation id carrying different content reaches the owner to be refused instead of collecting the earlier command's success. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: assert the drain's calls before restoring its spies Restoring a spy also clears its recorded calls, so the drain assertions ran against an emptied mock. Formats the durable claim method tests. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: close the follow-on gaps in the deletion fence and result claim The admission fence now carries an ownership token, so an overlapping deletion's fence is never lifted by the one that finishes first, and both fence writes invalidate the cached auth user document. It also covers the other bulk-delete path: `DELETE /` with no conversation filter removes every conversation, so it runs behind the same fence rather than a bare drain. The durable record now decides who holds a one-shot result. An owner replaying a retained response could hand the same terminal claim to a second invocation; that invocation is told the result was already collected, while the one that consumed it still recovers its own. A task with no durable record to arbitrate keeps whatever the owner answered. Drain cancellation treats `not_found` as unconfirmed: a missing registration while the durable lease is still live means the child may be running, so the command is retried under its invocation once the owner republishes itself. Control fingerprints are hashed, so retaining one per invocation costs a fixed few bytes instead of a bounded message. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: hold every deletion fence and keep live idempotency records An owner now holds one admission fence per concurrent bulk deletion instead of one at a time, so admission reopens only when the last deletion finishes regardless of completion order. Expired fences are pruned as new ones arrive and the set is bounded, so an abandoned fence cannot accumulate or lock an account out. A failed durable claim write is no longer read as an absent record. Handing a terminal result over without recording its claimant would let another invocation collect the same one-shot output once the database recovered, so the collection reports the retryable unavailable path and leaves the result for a later poll. Control invocation records now evict tasks the store no longer holds before live ones, over a bounded scan. Dropping a live task's record would let a caller retry apply its queue, steer, or interrupt a second time once the transport replay had also expired. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: keep the deletion fence portable and never drop a live record The admission fence is written with plain update operators again. DocumentDB rejects pipeline-form updates, and this runs before any deletion, so the pipeline form would have failed both bulk-delete endpoints outright on a supported database target. An excess deletion is now refused rather than silently displacing the oldest active fence, which would have reopened admission for a deletion still running. Expired fences are pruned before the cap is tested, so only genuinely concurrent deletions count against it. Control invocation records now sweep every settled task's entry when the window fills, and a window of entirely live records refuses the new control before touching the child instead of evicting one. Applying a command with no room to record it would let the caller's own retry apply it twice. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: hold the fence, bound recovered results, and expire stale commands The admission fence is renewed for as long as its deletion runs, so a very large account or a stalled database cannot let it lapse while conversations are still being removed. Only the deletion's own fence is renewed, and the renewal stops with the operation. Cancellation now covers every conversation the cascade removed, not only the ones a plan named: a grandchild lives in its own parent's scope, which a plan naming the deleted root never reaches. A routed request carries the deadline its caller waits for, and an owner drops one that arrives past it. A publisher disconnected mid-request queues the envelope offline and delivers it after the caller was told the owner was unavailable, which would otherwise steer a child the caller believes untouched. A result recovered from its durable child message is bounded like a routed one. The message keeps the child's untruncated output, so recovery could otherwise return far more than the routed result limit allows. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: size the fence window so a renewal can be observed The renewal test set a 30ms drain timeout but the five-minute grace window dominates it, so the interval was 100 seconds and no renewal could fire inside the test's deletion. The grace window is an option now, matching the store's other timings, and the test sizes the window to 90ms. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: wire the durable claim method and close the fence follow-ons The production store never received `claimSubagentTaskResult`, so every terminal result would have surfaced as unavailable once a task settled. The host wires that object from JavaScript, where the factory's parameter type checks nothing, so the factory now refuses a store missing any method it calls rather than failing at the first claim. The routing transport takes a dedicated publisher with the offline queue disabled. The shared client held commands issued during a disconnect and delivered them after the caller had given up, which the request deadline narrowed but could not close inside the clock-skew allowance. Fence renewal invalidates the cached auth document like the fence and release paths, and a renewal reporting its entry gone re-takes the fence instead of letting the deletion run on unfenced. The post-delete cancellation retries a transiently unreachable owner: the conversations are already gone, so it is the only pass that can still stop a late-admitted child. A replaced replay entry no longer leaves its bytes counted, which would have inflated the cache's total until unrelated responses were evicted. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: wait on observed lease renewal instead of a fixed delay The shared-lease renewal test held a 60ms lease and slept 100ms before asserting an overlapping worker was refused, so a loaded runner that starved the 10ms heartbeat past the TTL let the lease lapse and the second worker run. Spy on acquisition and renewal, then wait until a renewal succeeds past the acquired lease's own deadline — direct evidence the heartbeat carried it past expiry, with no timing assumption — and give the lease enough headroom that a stalled timer no longer decides the outcome. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): close the routing, fence, and cache gaps found in review Five separate seams, each with its own failure: `Cluster.duplicate` reads its first argument as a startup-node list and its second as the overrides, unlike `Redis.duplicate`, so the publisher's `enableOfflineQueue: false` was silently dropped under `USE_REDIS_CLUSTER` and a command issued mid-disconnect could still reach a child after its caller was told `unavailable`. Route both through `duplicateIoRedisClient`. The control window's capacity refusal ran before the store knew whether it owned the task, so unrelated local load could veto a cancellation bound for another replica. Establish that the task is local first and leave a remote one to its owner's window. `clearInterval` stops only future fence renewals. One already waiting on the database could resolve after the release, read its own lifted fence as expiry, and write a replacement that nothing remained to lift — closing subagent admission for the account until it aged out. Track the in-flight renewal, refuse overlapping passes, and await it before releasing. Every owner bounds its own task list, but the aggregation appended each batch whole, so the model-facing list grew with the number of replicas holding the scope. Cap the merged list while still reading every reply for the stale-registration sweep. The admission-fence prune commits independently of the fence that follows it, so a refused or failed push left the cached auth document describing entries the collection no longer held. Invalidate whichever way the second write goes. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): cap the merged task list the poll tool actually reads Each owner bounds its own reply and the remote aggregation bounds their sum, but `listTasks` merged that bounded remote list with however many children this replica owns and returned it whole. `check_background_task` could therefore still receive roughly twice the advertised cap. Bound the deduplicated, sorted result and export the cap so both seams share one number. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: admit every task the merged-list cap test starts The base store admits ten concurrent runs per scope by default, so starting 150 at once left most refused for capacity and the assertion never reached the merge it was written to check. Raise the cap for this store only; admission is a different invariant with its own tests. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): let a deletion notice its admission fence lapsing Renewal failures were logged and swallowed, so a run of rejected writes let the last confirmed `fencedUntil` pass while the deletion carried on believing admission was still closed — long enough for another replica to admit a child against conversations about to be removed. Track the deadline only a confirmed write advances, and check it after the drain, before anything is deleted: nothing has been removed at that point, so the operation fails closed and the caller retries once the fence can be held. A lapse detected after the rows are gone is logged instead, since reporting failure there would invite a retry against conversations that no longer exist. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: raise both concurrency caps the merged-list test trips Raising the per-scope limit left the store-wide `maxRunningTotal` at its default hundred, so fifty of the hundred and fifty starts were still refused. Verified against the base store directly this time: with only the per-scope cap raised it admits a hundred, and with both raised it admits all hundred and fifty. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): close the fence renewal gap and keep running tasks listed A renewal that started before its deadline but landed after it was still credited with extending the fence from its own start time, so a window in which admission stood open was papered over: a child could take a lease the drain had already read past and the deletion would proceed without cancelling it. The deadline now only advances when the write lands while the previous one still holds; anything later records a lapse the fence cannot be restored backwards over. The model-facing cap sorted oldest-first and sliced, which dropped the newest tasks — including children that had only just started running, and which the poll tool offers no other way to discover. Bound by status instead: running children first, then the most recent settled results. Both caps share one helper, and the routed aggregation now bounds after its loop so the choice is made across every owner's reply rather than by whichever answered first. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): finish the cap and the fence at the seams they still missed The status-aware cap only reached the requester: an owner's own reply still sliced positionally, so a replica holding more than the cap dropped its running children before the requester could bound anything. Both sides now share `boundedTaskList`. A fence that lapsed during the deletion itself was only logged. The rows are gone by then, so failing is still wrong, but the child another replica admitted while the fence was down is not: the fence is retaken and the drain repeated to cancel it. A child's lease renewal had the same retroactive hole the admission fence had — Mongo filters on the `now` captured before the call, so a write landing after the lease expired still moves the row forward, while an owner drain reading active leases in that gap saw the thread as free. The lease now carries its own deadline and a late renewal stops the executor. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: cover the lease lapse and the post-deletion re-drain The owner-side cap shipped with a regression test; these two did not. One drives a lease renewal that succeeds only after the lease it was extending had expired and asserts the executor stops; the other lets the fence lapse during the deletion itself and asserts a second drain runs while the request still reports success. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix(agents): close live-task lifecycle gaps * test(redis): exercise cluster node discovery * fix(test): type cluster discovery seam * fix(ci): wait for orphaned apt processes * fix(ci): reserve time for apt drain * fix(ci): skip optional fonts in MCP jobs * fix(agents): recover tasks after owner loss * fix(agents): preserve local task discovery * fix(agents): initialize fail-fast cluster publisher * style(agents): sort routing test imports --------- Co-authored-by: Claude --- .github/workflows/playwright-mock.yml | 6 - CONTEXT.md | 1 + api/server/experimental.js | 5 + api/server/experimental.spec.js | 10 + api/server/index.js | 2 + api/server/index.spec.js | 8 + .../__test-utils__/convos-route-mocks.js | 15 +- api/server/routes/__tests__/convos.spec.js | 58 +- api/server/routes/convos.js | 58 +- .../services/Endpoints/agents/initialize.js | 19 +- .../Endpoints/agents/initialize.spec.js | 16 +- .../Endpoints/agents/subagentThreadStore.js | 58 +- packages/api/src/agents/background.spec.ts | 210 ++- packages/api/src/agents/background.ts | 186 ++- packages/api/src/agents/guard.spec.ts | 2 + packages/api/src/agents/handlers.ts | 5 +- packages/api/src/agents/index.ts | 1 + .../src/agents/subagentTaskRouting.spec.ts | 1301 ++++++++++++++++ .../api/src/agents/subagentTaskRouting.ts | 1379 +++++++++++++++++ .../api/src/agents/subagentThreads.spec.ts | 1192 +++++++++++++- packages/api/src/agents/subagentThreads.ts | 789 +++++++++- packages/api/src/cache/redisUtils.spec.ts | 81 + packages/api/src/cache/redisUtils.ts | 44 +- .../src/methods/conversation.spec.ts | 9 + .../data-schemas/src/methods/conversation.ts | 38 + packages/data-schemas/src/methods/index.ts | 8 +- .../data-schemas/src/methods/message.spec.ts | 79 + packages/data-schemas/src/methods/message.ts | 68 + .../src/methods/user.methods.spec.ts | 153 ++ packages/data-schemas/src/methods/user.ts | 105 ++ packages/data-schemas/src/schema/message.ts | 8 + packages/data-schemas/src/schema/user.ts | 11 + packages/data-schemas/src/types/convo.ts | 6 + packages/data-schemas/src/types/message.ts | 5 + packages/data-schemas/src/types/user.ts | 5 + 35 files changed, 5839 insertions(+), 102 deletions(-) create mode 100644 packages/api/src/agents/subagentTaskRouting.spec.ts create mode 100644 packages/api/src/agents/subagentTaskRouting.ts create mode 100644 packages/api/src/cache/redisUtils.spec.ts diff --git a/.github/workflows/playwright-mock.yml b/.github/workflows/playwright-mock.yml index db3a275b255..6c57b4fe102 100644 --- a/.github/workflows/playwright-mock.yml +++ b/.github/workflows/playwright-mock.yml @@ -251,12 +251,6 @@ jobs: continue-on-error: true run: timeout -k 10 90 npx playwright install ffmpeg - # Optional fonts only — see the note in the e2e_shards job. - - name: Install optional Playwright font dependencies (best effort) - timeout-minutes: 4 - continue-on-error: true - run: .github/scripts/install-playwright-fonts.sh - # Redis is a hard requirement for this job, so this step stays fatal. - name: Install Redis runtime dependencies timeout-minutes: 5 diff --git a/CONTEXT.md b/CONTEXT.md index 80f2c15602b..6e411689c07 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -2,4 +2,5 @@ - **Agent run envelope**: the versioned, JSON-safe request contract created after ingress authentication and protocol validation but before agent, provider, tool, or MCP initialization. It carries only the validated protocol payload and the minimum trusted principal identifiers. The execution host rehydrates all runtime state from those identifiers. - **Subagent thread**: a durable, view-only child conversation owned by one parent conversation and subagent identity. A parent agent may continue it by stable `threadId`; each continuation uses a fresh execution lease restored from the canonical child transcript. It is not an ordinary human-writable chat. +- **Live subagent task owner**: the one API process holding a detached child execution, its abort controller, and its bounded control queue. Redis may route trusted poll/control envelopes to that owner, but it does not migrate or persist the executor; Mongo persists only the logical child thread and its continuation fence. - **Theme definition**: a versioned, data-only description of LibreChat semantic colors and shared appearance roles, optionally specialized by light or dark mode. The theme module validates and resolves partial definitions against bundled defaults before adapters apply them. A theme definition does not contain arbitrary CSS, application behavior, or alternate feature layouts. diff --git a/api/server/experimental.js b/api/server/experimental.js index 2d2ff5ec6fc..698cb52ff26 100644 --- a/api/server/experimental.js +++ b/api/server/experimental.js @@ -28,6 +28,7 @@ const { setupGracefulShutdown, configureMessageFilterRegexValidator, configureFileConfigRegexEngine, + waitForKeyvRedisClient, } = require('@librechat/api'); const { connectDb, indexSync } = require('~/db'); const initializeOAuthReconnectManager = require('./services/initializeOAuthReconnectManager'); @@ -36,6 +37,7 @@ const createValidateImageRequest = require('./middleware/validateImageRequest'); const { startExpiredFileSweep } = require('./services/Files/process'); const { initializeGitHubSkillSync } = require('./services/Skills/sync'); const { initializeAgentTriggerService } = require('./services/Agents/triggers'); +const { configureSubagentTaskRouting } = require('./services/Endpoints/agents/subagentThreadStore'); const { jwtLogin, ldapLogin, passportLogin } = require('~/strategies'); const { updateInterfacePermissions: updateInterfacePerms } = require('@librechat/api'); const { @@ -304,6 +306,9 @@ if (cluster.isMaster) { const startServer = async () => { logger.info(`Worker ${process.pid} initializing...`); + await waitForKeyvRedisClient(); + await configureSubagentTaskRouting(); + if (typeof Bun !== 'undefined') { axios.defaults.headers.common['Accept-Encoding'] = 'gzip'; } diff --git a/api/server/experimental.spec.js b/api/server/experimental.spec.js index c4364fb6b8e..4ff9d6b924a 100644 --- a/api/server/experimental.spec.js +++ b/api/server/experimental.spec.js @@ -30,6 +30,16 @@ describe('Experimental server configuration', () => { ); }); + it('configures routed subagent controls before a worker accepts requests', () => { + const redisReadyIndex = source.indexOf('await waitForKeyvRedisClient();'); + const routingIndex = source.indexOf('await configureSubagentTaskRouting();'); + const listenIndex = source.indexOf('const server = app.listen'); + + expect(redisReadyIndex).toBeGreaterThan(-1); + expect(routingIndex).toBeGreaterThan(redisReadyIndex); + expect(listenIndex).toBeGreaterThan(routingIndex); + }); + it('matches the standard server pre-authentication tenant routes', () => { expect(source).toContain("app.use('/oauth', preAuthTenantMiddleware, routes.oauth);"); expect(source).toContain("app.use('/api/auth', preAuthTenantMiddleware, routes.auth);"); diff --git a/api/server/index.js b/api/server/index.js index f3cf7c67ec8..0768a27c50d 100644 --- a/api/server/index.js +++ b/api/server/index.js @@ -62,6 +62,7 @@ const { startExpiredFileSweep } = require('./services/Files/process'); const { checkMigrations } = require('./services/start/migration'); const optionalJwtAuth = require('./middleware/optionalJwtAuth'); const initializeMCPs = require('./services/initializeMCPs'); +const { configureSubagentTaskRouting } = require('./services/Endpoints/agents/subagentThreadStore'); const configureSocialLogins = require('./socialLogins'); const createSpaFallback = require('./utils/fallback'); const { getAppConfig } = require('./services/Config'); @@ -124,6 +125,7 @@ const configureGenerationStreams = () => { const startServer = async () => { await waitForKeyvRedisClient(); + await configureSubagentTaskRouting(); const { metricsMiddleware, metricsRouter } = createMetrics(); if (!process.env.METRICS_SECRET) { logger.warn('[metrics] METRICS_SECRET is not set - /metrics will return 401 for all requests'); diff --git a/api/server/index.spec.js b/api/server/index.spec.js index 6461a148e38..013ece204ef 100644 --- a/api/server/index.spec.js +++ b/api/server/index.spec.js @@ -136,6 +136,14 @@ describe('Startup readiness wiring', () => { expect(streamConfigIndex).toBeLessThan(postListenMcpIndex); }); + it('configures subagent task routing before the server accepts requests', () => { + const routingIndex = source.indexOf('await configureSubagentTaskRouting();'); + const listenIndex = source.indexOf('const server = app.listen'); + + expect(routingIndex).toBeGreaterThan(-1); + expect(listenIndex).toBeGreaterThan(routingIndex); + }); + it('registers generation stream cleanup with the graceful shutdown coordinator', () => { const shutdownRegistrationIndex = source.indexOf( "registerShutdownTask('generation job manager'", diff --git a/api/server/routes/__test-utils__/convos-route-mocks.js b/api/server/routes/__test-utils__/convos-route-mocks.js index d6ac3ca0dab..5233d27ef0d 100644 --- a/api/server/routes/__test-utils__/convos-route-mocks.js +++ b/api/server/routes/__test-utils__/convos-route-mocks.js @@ -123,7 +123,20 @@ module.exports = { assistantEndpoint: () => ({ initializeClient: jest.fn() }), subagentThreadStore: () => ({ - cancelForConversations: jest.fn(), + cancelAndDrainForOwner: jest.fn().mockResolvedValue(undefined), + withOwnerDeletionFence: jest.fn().mockImplementation(async (_userId, _tenantId, deletion) => { + return deletion(); + }), + planCancellationForConversations: jest + .fn() + .mockImplementation(async (userId, conversationIds, tenantId) => ({ + userId, + tenantId, + conversationIds: [...conversationIds], + scopes: [], + leases: [], + })), + cancelPlan: jest.fn().mockResolvedValue(0), cancelForOwner: jest.fn(), }), }; diff --git a/api/server/routes/__tests__/convos.spec.js b/api/server/routes/__tests__/convos.spec.js index ec5b23f263b..c09a8ac0753 100644 --- a/api/server/routes/__tests__/convos.spec.js +++ b/api/server/routes/__tests__/convos.spec.js @@ -65,7 +65,13 @@ describe('Convos Routes', () => { expect(response.status).toBe(201); expect(deleteAgentCheckpoints).toHaveBeenCalledTimes(1); expect(deleteAgentCheckpoints.mock.calls[0][0]).toEqual(conversationIds); - expect(subagentThreadStore.cancelForOwner).toHaveBeenCalledWith('test-user-123', undefined); + /** The deletion runs inside the owner admission fence, not around it. */ + expect(subagentThreadStore.withOwnerDeletionFence).toHaveBeenCalledTimes(1); + const [fencedUserId, fencedTenantId] = + subagentThreadStore.withOwnerDeletionFence.mock.calls[0]; + expect(fencedUserId).toBe('test-user-123'); + expect(fencedTenantId).toBeUndefined(); + expect(subagentThreadStore.cancelAndDrainForOwner).not.toHaveBeenCalled(); }); it('should delete all conversations, tool calls, and shared links for a user', async () => { @@ -132,6 +138,21 @@ describe('Convos Routes', () => { expect(logger.error).toHaveBeenCalledWith('Error clearing conversations', expect.any(Error)); }); + it('does not delete conversations when cross-replica task draining fails', async () => { + /** Draining happens inside the admission fence, so its failure fails the fence. */ + subagentThreadStore.withOwnerDeletionFence.mockRejectedValueOnce( + new Error('task owner unavailable'), + ); + + const response = await request(app).delete('/api/convos/all'); + + expect(response.status).toBe(500); + expect(deleteConvos).not.toHaveBeenCalled(); + expect(deleteAgentCheckpoints).not.toHaveBeenCalled(); + expect(deleteToolCalls).not.toHaveBeenCalled(); + expect(deleteAllSharedLinksWithCleanup).not.toHaveBeenCalled(); + }); + it('should return 500 if deleteToolCalls fails', async () => { deleteConvos.mockResolvedValue({ deletedCount: 5 }); deleteToolCalls.mockRejectedValue(new Error('Tool calls deletion failed')); @@ -239,6 +260,21 @@ describe('Convos Routes', () => { }); describe('DELETE /', () => { + it('fences the owner when DELETE / is called without a conversation filter', async () => { + deleteConvos.mockResolvedValue({ deletedCount: 3, conversationIds: ['a', 'b', 'c'] }); + + const response = await request(app) + .delete('/api/convos') + .send({ arg: { thread_id: 'thread-abc' } }); + + expect(response.status).toBe(201); + /** An empty filter deletes everything, so it takes the same admission fence. */ + expect(subagentThreadStore.withOwnerDeletionFence).toHaveBeenCalledTimes(1); + expect(subagentThreadStore.withOwnerDeletionFence.mock.calls[0][0]).toBe('test-user-123'); + expect(subagentThreadStore.cancelAndDrainForOwner).not.toHaveBeenCalled(); + expect(deleteConvos).toHaveBeenCalledWith('test-user-123', {}); + }); + it('cancels root and descendant leases and cleans every cascaded conversation', async () => { deleteConvos.mockResolvedValue({ deletedCount: 2, @@ -254,18 +290,22 @@ describe('Convos Routes', () => { }); expect(response.status).toBe(201); - expect(subagentThreadStore.cancelForConversations).toHaveBeenNthCalledWith( - 1, + /** The plan is resolved before deletion, while those rows can still be read. */ + expect(subagentThreadStore.planCancellationForConversations).toHaveBeenCalledWith( 'test-user-123', ['parent-conversation'], undefined, ); - expect(subagentThreadStore.cancelForConversations).toHaveBeenNthCalledWith( - 2, - 'test-user-123', - ['parent-conversation', 'child-conversation'], - undefined, - ); + expect( + subagentThreadStore.planCancellationForConversations.mock.invocationCallOrder[0], + ).toBeLessThan(deleteConvos.mock.invocationCallOrder[0]); + /** It is applied once before deletion and replayed after with the cascade. */ + expect(subagentThreadStore.cancelPlan).toHaveBeenCalledTimes(2); + expect(subagentThreadStore.cancelPlan.mock.calls[0][1]).toBeUndefined(); + expect(subagentThreadStore.cancelPlan.mock.calls[1][1]).toEqual([ + 'parent-conversation', + 'child-conversation', + ]); expect(deleteToolCalls.mock.calls.map((call) => call[1])).toEqual([ 'parent-conversation', 'child-conversation', diff --git a/api/server/routes/convos.js b/api/server/routes/convos.js index 87a9d0a719d..fcba4f9a83c 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -118,6 +118,26 @@ router.get('/gen_title/:conversationId', async (req, res) => { } }); +const POST_DELETE_CANCEL_ATTEMPTS = 3; +const POST_DELETE_CANCEL_BACKOFF_MS = 250; + +/** Replays a cancellation plan after deletion, retrying a transiently unreachable + * owner rather than losing the only pass that can stop a late-admitted child. */ +async function retryPostDeleteCancellation(cancellationPlan, deletedConversationIds) { + for (let attempt = 1; attempt <= POST_DELETE_CANCEL_ATTEMPTS; attempt += 1) { + try { + await subagentThreadTaskStore.cancelPlan(cancellationPlan, deletedConversationIds); + return; + } catch (error) { + if (attempt === POST_DELETE_CANCEL_ATTEMPTS) { + logger.warn('Post-delete subagent cancellation failed', error); + return; + } + await new Promise((resolve) => setTimeout(resolve, POST_DELETE_CANCEL_BACKOFF_MS * attempt)); + } + } +} + router.delete('/', configMiddleware, async (req, res) => { let filter = {}; const { conversationId, source, thread_id, endpoint } = req.body?.arg ?? {}; @@ -154,17 +174,36 @@ router.delete('/', configMiddleware, async (req, res) => { typeof req.user.tenantId === 'string' && req.user.tenantId !== '' ? req.user.tenantId : undefined; + let cancellationPlan; + let dbResponse; if (filter.conversationId) { - subagentThreadTaskStore.cancelForConversations( + /** Resolve the targets while the conversations still exist: the second pass + * runs after their rows are gone and can only reach registered owners. */ + cancellationPlan = await subagentThreadTaskStore.planCancellationForConversations( req.user.id, [filter.conversationId], tenantId, ); + await subagentThreadTaskStore.cancelPlan(cancellationPlan); + dbResponse = await db.deleteConvos(req.user.id, filter); + } else { + /** An empty filter deletes every conversation this owner has, so it runs behind + * the same admission fence as `DELETE /all` rather than a bare drain. */ + dbResponse = await subagentThreadTaskStore.withOwnerDeletionFence(req.user.id, tenantId, () => + db.deleteConvos(req.user.id, filter), + ); } - const dbResponse = await db.deleteConvos(req.user.id, filter); const deletedConversationIds = dbResponse.conversationIds ?? (filter.conversationId ? [filter.conversationId] : []); - subagentThreadTaskStore.cancelForConversations(req.user.id, deletedConversationIds, tenantId); + /** Root deletion closes new child admission. Replay the plan to catch a task + * admitted after the first pass but before that fence, extended with the cascade + * this deletion reported. */ + if (cancellationPlan != null && deletedConversationIds.length > 0) { + /** The conversations are gone, so this pass is the only thing that can still + * stop a child admitted after the first one. It cannot fail the request — the + * deletion already committed — so it retries briefly before giving up. */ + await retryPostDeleteCancellation(cancellationPlan, deletedConversationIds); + } // HITL: prune the deleted conversations' durable checkpoints — a paused run's // checkpoint would otherwise persist until the Mongo TTL. Never throws. await deleteAgentCheckpoints( @@ -186,13 +225,18 @@ router.delete('/', configMiddleware, async (req, res) => { router.delete('/all', configMiddleware, async (req, res) => { try { - subagentThreadTaskStore.cancelForOwner( - req.user.id, + const tenantId = typeof req.user.tenantId === 'string' && req.user.tenantId !== '' ? req.user.tenantId - : undefined, + : undefined; + /** Fences new child admission for this owner, drains the live ones, and deletes + * inside that fence: a child admitted on another replica mid-deletion would + * otherwise keep running against conversations that no longer exist. */ + const dbResponse = await subagentThreadTaskStore.withOwnerDeletionFence( + req.user.id, + tenantId, + () => db.deleteConvos(req.user.id, {}), ); - const dbResponse = await db.deleteConvos(req.user.id, {}); // HITL: prune ALL the deleted conversations' durable checkpoints in one bulk pass. await deleteAgentCheckpoints( dbResponse.conversationIds, diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 46be766f7cb..a3256f54bd2 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -1210,7 +1210,7 @@ const initializeClient = async ({ } /** Build detached execution only for an attributable owner/thread. New - * tasks still require a spawnable child, while an existing process-local + * tasks still require a spawnable child, while an existing registered live * task keeps its poll/control seam after agent configuration changes. The * SDK receives only this trusted host scope; models can select a child * `threadId`, never the owner or parent-thread namespace. */ @@ -1236,9 +1236,20 @@ const initializeClient = async ({ : {}), }) : undefined; - const hasExistingSubagentTask = - trustedSubagentTasks != null && - trustedSubagentTasks.store.list(trustedSubagentTasks.scopeId).length > 0; + let hasExistingSubagentTask = false; + if (trustedSubagentTasks != null && !(subagentsAvailableForRun && hasSpawnableSubagent)) { + try { + hasExistingSubagentTask = await subagentThreadTaskStore.hasTasks( + trustedSubagentTasks.scopeId, + ); + } catch (error) { + /** Keep the poll/control tool visible when the owner directory is briefly + * unavailable. The tool then returns an honest `unavailable` status + * instead of making a live task look nonexistent. */ + logger.warn('[initializeClient] Failed to inspect routed subagent tasks', error); + hasExistingSubagentTask = true; + } + } const subagentTasks = trustedSubagentTasks != null && ((subagentsAvailableForRun && hasSpawnableSubagent) || hasExistingSubagentTask) diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index a5be3f98528..59910aa2a0c 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -756,19 +756,7 @@ describe('initializeClient — subagent loading', () => { endpointOption: makeEndpointOption(), }); const existingConfig = agentClientArgs.subagentTasks; - const listSpy = jest.spyOn(existingConfig.store, 'list').mockReturnValueOnce([ - { - taskId: 'existing-task', - threadId: 'existing-thread', - subagentType: 'researcher', - status: 'running', - createdAt: Date.now(), - updatedAt: Date.now(), - resultAvailable: false, - resultClaimed: false, - pendingControls: 0, - }, - ]); + const hasTasksSpy = jest.spyOn(existingConfig.store, 'hasTasks').mockResolvedValueOnce(true); mockInitializeAgent.mockResolvedValue(makePrimaryConfig({})); const changedReq = makeSubagentReq(); changedReq.config.endpoints.agents.capabilities.push('run_in_background'); @@ -783,7 +771,7 @@ describe('initializeClient — subagent loading', () => { expect(agentClientArgs.subagentTasks).toEqual(existingConfig); expect(capturedToolExecuteOptions.subagentTasks).toEqual(existingConfig); expect(agentClientArgs.agent.subagents).toBeUndefined(); - listSpy.mockRestore(); + hasTasksSpy.mockRestore(); }); it('disables every nested subagent path at the durable child-thread depth limit', async () => { diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.js b/api/server/services/Endpoints/agents/subagentThreadStore.js index ff1dd6c60bc..f11c6a22670 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -1,16 +1,25 @@ -const { createSubagentThreadTaskStore } = require('@librechat/api'); +const { + cacheConfig, + ioredisClient, + registerShutdownTask, + duplicateIoRedisClient, + createSubagentThreadTaskStore, + RedisSubagentTaskControlTransport, +} = require('@librechat/api'); const db = require('~/models'); -/** Durable logical threads use normal LibreChat conversations/messages. Live - * controls stay process-local; Mongo fences continuation across API replicas. */ +/** Durable logical threads use normal LibreChat conversations/messages. Mongo + * fences continuation; optional Redis routing reaches the live owning process. */ const subagentThreadTaskStore = createSubagentThreadTaskStore( { acquireSubagentThreadLease: db.acquireSubagentThreadLease, + claimSubagentTaskResult: db.claimSubagentTaskResult, countActiveSubagentThreadLeases: db.countActiveSubagentThreadLeases, deleteConvos: db.deleteConvos, deleteMessages: db.deleteMessages, getConvo: db.getConvo, getMessages: db.getMessages, + listActiveSubagentThreadLeases: db.listActiveSubagentThreadLeases, releaseSubagentThreadLease: db.releaseSubagentThreadLease, reserveSubagentThread: db.reserveSubagentThread, renewSubagentThreadLease: db.renewSubagentThreadLease, @@ -18,8 +27,49 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( saveMessage: db.saveMessage, }, { - isOwnerActive: db.isAgentTriggerPrincipalActive, + isOwnerActive: db.isSubagentOwnerAdmissible, + fenceOwnerAdmission: db.fenceSubagentAdmission, + renewOwnerAdmission: db.renewSubagentAdmission, + releaseOwnerAdmission: db.releaseSubagentAdmission, }, ); +let taskRoutingConfigured = false; + +/** Starts the optional Redis owner directory before HTTP admission opens. */ +async function configureSubagentTaskRouting() { + if (taskRoutingConfigured || !cacheConfig.USE_REDIS) { + return; + } + if (ioredisClient == null || typeof ioredisClient.duplicate !== 'function') { + throw new Error('Redis subagent task routing requires a dedicated subscriber connection.'); + } + const subscriber = ioredisClient.duplicate(); + /** A dedicated publisher without the offline queue: the shared client would hold a + * command issued during a disconnect and deliver it after the caller gave up, so a + * steer the caller was told had failed could still reach the child. Failing fast + * turns that into the honest `unavailable` the caller already handles. */ + const publisher = duplicateIoRedisClient(ioredisClient, { enableOfflineQueue: false }); + const transport = new RedisSubagentTaskControlTransport(publisher, subscriber, { + namespace: cacheConfig.REDIS_KEY_PREFIX, + }); + try { + await subagentThreadTaskStore.configureTaskControlTransport(transport); + } catch (error) { + subscriber.disconnect(); + publisher.disconnect(); + throw error; + } + taskRoutingConfigured = true; + registerShutdownTask( + 'subagent task control transport', + async () => { + await subagentThreadTaskStore.destroyTaskControlTransport(); + publisher.disconnect(); + }, + { priority: 90 }, + ); +} + module.exports = subagentThreadTaskStore; +module.exports.configureSubagentTaskRouting = configureSubagentTaskRouting; diff --git a/packages/api/src/agents/background.spec.ts b/packages/api/src/agents/background.spec.ts index 4c604c8b989..97fbc2eed69 100644 --- a/packages/api/src/agents/background.spec.ts +++ b/packages/api/src/agents/background.spec.ts @@ -19,6 +19,7 @@ import { CHECK_BACKGROUND_TASK_NAME, RUN_IN_BACKGROUND_ARG, } from './background'; +import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting'; import { TOOL_SELECTION_WILDCARD } from './selection'; import { toolOptionsSchema } from './validation'; @@ -1056,8 +1057,8 @@ describe('getBackgroundCodeDelivery (singleton)', () => { }); describe('runCheckBackgroundTask (singleton)', () => { - it('returns not_found for an unknown id', () => { - const content = runCheckBackgroundTask({ + it('returns not_found for an unknown id', async () => { + const content = await runCheckBackgroundTask({ userId: 'poll_user', conversationId: 'poll_convo', args: { background_task_id: 'nope' }, @@ -1067,7 +1068,29 @@ describe('runCheckBackgroundTask (singleton)', () => { ); }); - it('returns a single task by id and lists all when omitted', () => { + it('rejects an oversized task id before local or cross-replica lookup', async () => { + const store = Object.assign(new InMemorySubagentTaskStore(), { + claimTask: jest.fn(), + controlTask: jest.fn(), + listTasks: jest.fn(), + }); + const content = await runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { background_task_id: 'x'.repeat(257) }, + subagentTasks: { store, scopeId: 'owner:parent-thread' }, + }); + + expect(JSON.parse(content)).toEqual({ + status: 'invalid', + message: 'A background_task_id cannot exceed 256 characters.', + }); + expect(store.claimTask).not.toHaveBeenCalled(); + expect(store.controlTask).not.toHaveBeenCalled(); + expect(store.listTasks).not.toHaveBeenCalled(); + }); + + it('returns a single task by id and lists all when omitted', async () => { const created = backgroundTaskRegistry.create({ userId: 'poll_user', conversationId: 'poll_convo2', @@ -1082,7 +1105,7 @@ describe('runCheckBackgroundTask (singleton)', () => { }); const single = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'poll_user', conversationId: 'poll_convo2', args: { background_task_id: created.task.id }, @@ -1097,7 +1120,11 @@ describe('runCheckBackgroundTask (singleton)', () => { ); const listed = JSON.parse( - runCheckBackgroundTask({ userId: 'poll_user', conversationId: 'poll_convo2', args: {} }), + await runCheckBackgroundTask({ + userId: 'poll_user', + conversationId: 'poll_convo2', + args: {}, + }), ); expect(listed.tasks).toHaveLength(1); expect(listed.tasks[0].background_task_id).toBe(created.task.id); @@ -1109,7 +1136,7 @@ describe('runCheckBackgroundTask (singleton)', () => { // stringified args must still resolve the specific task (with its full result) const singleFromString = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'poll_user', conversationId: 'poll_convo2', args: `{"background_task_id":"${created.task.id}"}`, @@ -1120,7 +1147,68 @@ describe('runCheckBackgroundTask (singleton)', () => { ); }); - it('retrieves a task across turns: the poll is keyed only by id, not the dispatch run/turn', () => { + it('preserves local task lists when cross-replica subagent discovery is unavailable', async () => { + const ordinary = backgroundTaskRegistry.create({ + userId: 'partial-list-owner', + conversationId: 'partial-list-parent', + toolCallId: 'ordinary-call', + toolName: 'search_mcp_docs', + }); + if ('atCapacity' in ordinary) { + throw new Error('unexpected capacity'); + } + + const store = new InMemorySubagentTaskStore(); + const started = store.start({ + scopeId: 'partial-list-owner:partial-list-parent', + idempotencyKey: 'partial-list-run:parent-agent:subagent-call', + parentRunId: 'partial-list-run', + parentAgentId: 'parent-agent', + parentToolCallId: 'subagent-call', + input: 'Keep working locally.', + subagentKind: 'agent', + subagentType: 'researcher', + run: async () => ({ content: 'local result' }), + }); + if (!started.accepted) { + throw new Error('Expected subagent task to start.'); + } + await waitForSubagentTaskToSettle( + store, + 'partial-list-owner:partial-list-parent', + started.task.taskId, + ); + + const routedStore = Object.assign(store, { + claimTask: jest.fn(), + controlTask: jest.fn(), + listTasks: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()), + }); + const listed = JSON.parse( + await runCheckBackgroundTask({ + userId: 'partial-list-owner', + conversationId: 'partial-list-parent', + args: {}, + subagentTasks: { + store: routedStore, + scopeId: 'partial-list-owner:partial-list-parent', + }, + }), + ); + + expect(listed).toEqual( + expect.objectContaining({ + partial: true, + warning: + 'Cross-replica subagent tasks could not be listed: The process running this subagent task is temporarily unavailable.', + }), + ); + expect( + listed.tasks.map((task: { background_task_id: string }) => task.background_task_id), + ).toEqual(expect.arrayContaining([ordinary.task.id, started.task.taskId])); + }); + + it('retrieves a task across turns: the poll is keyed only by id, not the dispatch run/turn', async () => { // Turn 1 dispatches under run-turn-1 and the result lands after the turn. const dispatched = backgroundTaskRegistry.create({ userId: 'poll_user', @@ -1139,7 +1227,7 @@ describe('runCheckBackgroundTask (singleton)', () => { // Turn 2 (a later run) polls with just the id; get/list carry no run/turn scope. const polled = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'poll_user', conversationId: 'poll_xturn', args: { background_task_id: dispatched.task.id }, @@ -1174,7 +1262,7 @@ describe('runCheckBackgroundTask (singleton)', () => { await waitForSubagentTaskToSettle(store, subagentTasks.scopeId, started.task.taskId); const first = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'owner', conversationId: 'parent-thread', args: { background_task_id: started.task.taskId }, @@ -1192,7 +1280,7 @@ describe('runCheckBackgroundTask (singleton)', () => { ); const second = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'owner', conversationId: 'parent-thread', args: { background_task_id: started.task.taskId }, @@ -1227,7 +1315,7 @@ describe('runCheckBackgroundTask (singleton)', () => { await Promise.resolve(); const queued = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'owner', conversationId: 'parent-thread', args: { @@ -1243,7 +1331,7 @@ describe('runCheckBackgroundTask (singleton)', () => { ); const cancelledMessage = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'owner', conversationId: 'parent-thread', args: { @@ -1257,7 +1345,7 @@ describe('runCheckBackgroundTask (singleton)', () => { expect(cancelledMessage.status).toBe('accepted'); const cancelledTask = JSON.parse( - runCheckBackgroundTask({ + await runCheckBackgroundTask({ userId: 'owner', conversationId: 'parent-thread', args: { background_task_id: started.task.taskId, action: 'cancel' }, @@ -1267,6 +1355,102 @@ describe('runCheckBackgroundTask (singleton)', () => { expect(cancelledTask.status).toBe('cancelled'); finish({ content: 'late result' }); }); + + it('derives a bounded control invocation identity from the tool call', async () => { + const controlTask = jest.fn().mockResolvedValue({ + status: 'not_running', + task: { + taskId: 'remote-task', + subagentType: 'researcher', + status: 'completed', + createdAt: 1, + updatedAt: 2, + resultAvailable: false, + resultClaimed: true, + pendingControls: 0, + }, + }); + const store = Object.assign(new InMemorySubagentTaskStore(), { + claimTask: jest.fn(), + controlTask, + listTasks: jest.fn(), + }); + const control = (toolCallId: string | undefined) => + runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { + background_task_id: 'remote-task', + action: 'queue', + message: 'Check one more source.', + }, + toolCallId, + subagentTasks: { store, scopeId: 'owner:parent-thread' }, + }); + + /** Replaying one tool call keeps its identity, so routing can replay the result. */ + await control('call_abc'); + await control('call_abc'); + const [firstInvocation, replayedInvocation] = controlTask.mock.calls.map((call) => call[3]); + expect(firstInvocation).toBe(replayedInvocation); + expect(firstInvocation).toHaveLength(32); + + /** A separate tool call is a separate command even with an identical payload. */ + await control('call_def'); + expect(controlTask.mock.calls[2][3]).not.toBe(firstInvocation); + + /** The same provider id in another run or agent is a different command. */ + await runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { + background_task_id: 'remote-task', + action: 'queue', + message: 'Check one more source.', + }, + toolCallId: 'call_abc', + runId: 'run-2:0', + subagentTasks: { store, scopeId: 'owner:parent-thread' }, + }); + expect(controlTask.mock.calls[3][3]).not.toBe(firstInvocation); + + /** A provider id far past the protocol bound still routes as a bounded identity. */ + const longToolCallId = `call_${'x'.repeat(200)}`; + await control(longToolCallId); + await control(longToolCallId); + const [longInvocation, replayedLongInvocation] = controlTask.mock.calls + .slice(4) + .map((call) => call[3]); + expect(longInvocation).toHaveLength(32); + expect(replayedLongInvocation).toBe(longInvocation); + + /** Without a tool-call id each invocation stays distinct rather than colliding. */ + await control(undefined); + await control(undefined); + const [fallback, otherFallback] = controlTask.mock.calls.slice(6).map((call) => call[3]); + expect(fallback).not.toBe(otherFallback); + expect(fallback.length).toBeLessThanOrEqual(128); + }); + + it('reports an unreachable remote subagent owner without pretending the task is missing', async () => { + const store = Object.assign(new InMemorySubagentTaskStore(), { + claimTask: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()), + controlTask: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()), + listTasks: jest.fn().mockRejectedValue(new SubagentTaskOwnerUnavailableError()), + }); + const content = await runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { background_task_id: 'remote-task' }, + subagentTasks: { store, scopeId: 'owner:parent-thread' }, + }); + + expect(JSON.parse(content)).toEqual({ + status: 'unavailable', + background_task_id: 'remote-task', + message: 'The process running this subagent task is temporarily unavailable.', + }); + }); }); describe('stripBackgroundFromToolRegistry', () => { diff --git a/packages/api/src/agents/background.ts b/packages/api/src/agents/background.ts index 0f06c2ad6d9..113bd2c7fbe 100644 --- a/packages/api/src/agents/background.ts +++ b/packages/api/src/agents/background.ts @@ -17,7 +17,10 @@ * are lost on restart and are not shared across replicas (durable follow-up), * and ephemeral request-scoped MCP tools (runtime `{{LIBRECHAT_BODY_*}}` * placeholders) are never backgrounded — their connection is torn down at - * request end, so the executor runs them in the foreground instead. + * request end, so the executor runs them in the foreground instead. Detached + * subagents use the separate host task store; Redis-backed hosts may route + * their poll/control operations to the owning process without moving the live + * executor or making ordinary background tool results durable. * * Opt-in mirrors `deferred_tools`: an admin capability * (`AgentCapabilities.run_in_background`) gates the feature, and a per-tool @@ -30,8 +33,8 @@ * @module packages/api/src/agents/background */ -import { randomUUID } from 'node:crypto'; import { logger } from '@librechat/data-schemas'; +import { createHash, randomUUID } from 'node:crypto'; import { Constants as AgentConstants } from '@librechat/agents'; import { Tools, Constants, imageGenTools } from 'librechat-data-provider'; import type { @@ -43,6 +46,7 @@ import type { SubagentTaskSnapshot, SubagentTaskControlCommand, SubagentTaskControlResult, + SubagentTaskStore, } from '@librechat/agents'; import type { AgentToolOptions } from 'librechat-data-provider'; import type { CapabilityToolNames } from './selection'; @@ -52,6 +56,7 @@ import { warnUnmatchedSelectionNames, synthesizeSelectionToolOptions, } from './selection'; +import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting'; import { SET_MEMORY_TOOL_NAME, DELETE_MEMORY_TOOL_NAME } from './memory'; import { ASK_USER_QUESTION_TOOL_NAME } from './hitl/askUserQuestionTool'; import { CREATE_FILE_TOOL_NAME, EDIT_FILE_TOOL_NAME } from './tools'; @@ -62,6 +67,9 @@ export const RUN_IN_BACKGROUND_ARG = 'run_in_background'; /** Log prefix for selection diagnostics, phrased in the spec's own field name. */ const BACKGROUND_SELECTION_LABEL = '[background] runInBackground'; +const MAX_BACKGROUND_TASK_ID_CHARS = 256; +const MAX_BACKGROUND_CONTROL_ID_CHARS = 256; +const MAX_BACKGROUND_CONTROL_MESSAGE_CHARS = 64 * 1024; /** * `type` of the synthetic attachment emitted on a poll turn when a harvested @@ -297,13 +305,36 @@ export function stripBackgroundFromToolRegistry( const CHECK_BACKGROUND_TASK_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 poll one task; omit it to list every background task in this thread. A task is only finished when its status is "completed", "error", or "cancelled" — never assume completion without polling. Results are not pushed to you; you must call this tool to collect them. Subagent tasks additionally accept steer, queue, interrupt, cancel, and cancel_message actions while running. Execution leases remain available only while requests reach the owning server process; they do not survive a restart or cross-worker routing. A completed subagent thread may be continued later through the subagent tool's durable thread id.`; +Provide a background_task_id to poll one task; omit it to list every background task in this thread. A task is only finished when its status is "completed", "error", or "cancelled" — never assume completion without polling. Results are not pushed to you; you must call this tool to collect them. Subagent tasks additionally accept steer, queue, interrupt, cancel, and cancel_message actions while running. 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.`; -const CHECK_BACKGROUND_TASK_PARAMETERS: JsonSchemaType = Object.freeze({ +/** + * `maxLength` is valid JSON Schema and is honored by providers, but the SDK's + * `JsonSchemaType` does not declare it, so the model-facing bounds are typed here. + * Runtime argument validation enforces the same limits as defense in depth. + */ +interface BoundedStringSchema { + type: 'string'; + maxLength: number; + description: string; +} + +interface CheckBackgroundTaskParameters { + type: 'object'; + properties: { + background_task_id: BoundedStringSchema; + action: { type: 'string'; enum: string[]; description: string }; + message: BoundedStringSchema; + control_id: BoundedStringSchema; + }; + required: string[]; +} + +const CHECK_BACKGROUND_TASK_PARAMETERS = Object.freeze({ type: 'object', properties: { background_task_id: { type: 'string', + maxLength: MAX_BACKGROUND_TASK_ID_CHARS, description: 'The id returned when the tool or subagent was dispatched. Omit to list all background tasks in this thread.', }, @@ -314,10 +345,12 @@ const CHECK_BACKGROUND_TASK_PARAMETERS: JsonSchemaType = Object.freeze; + controlTask( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise; + listTasks(scopeId: string): Promise; +} + +function routedSubagentStore(store: SubagentTaskStore): RoutedSubagentTaskStore | undefined { + const candidate = store as SubagentTaskStore & Partial; + return typeof candidate.claimTask === 'function' && + typeof candidate.controlTask === 'function' && + typeof candidate.listTasks === 'function' + ? (candidate as RoutedSubagentTaskStore) + : undefined; +} + +export async function runCheckBackgroundTask(params: { userId: string; conversationId: string; args: unknown; + /** The provider's tool-call id: one control invocation, stable across replays. */ + toolCallId?: string; + /** Scopes that tool-call id, whose provider ids repeat across runs and agents. */ + agentId?: string; + runId?: string; subagentTasks?: SubagentTaskConfig; -}): string { +}): Promise { const { userId, conversationId } = params; const args = coerceArgsObject(params.args) ?? {}; const rawId = args.background_task_id; + if (typeof rawId === 'string' && rawId.trim().length > MAX_BACKGROUND_TASK_ID_CHARS) { + return JSON.stringify({ + status: 'invalid', + message: `A background_task_id cannot exceed ${MAX_BACKGROUND_TASK_ID_CHARS} characters.`, + }); + } const taskId = typeof rawId === 'string' && rawId.trim() !== '' ? rawId.trim() : undefined; const action = typeof args.action === 'string' && args.action !== '' ? args.action : 'poll'; + const invocationId = controlInvocationId(params); if (taskId) { const task = backgroundTaskRegistry.get(userId, conversationId, taskId); @@ -1118,28 +1209,44 @@ export function runCheckBackgroundTask(params: { const subagentTasks = params.subagentTasks; if (subagentTasks != null) { - if (action === 'poll') { - const claimed = serializeSubagentClaim( - subagentTasks.store.claim(subagentTasks.scopeId, taskId), - ); - if (claimed != null) { - return JSON.stringify(claimed); + try { + const routedStore = routedSubagentStore(subagentTasks.store); + if (action === 'poll') { + const claim = + routedStore == null + ? subagentTasks.store.claim(subagentTasks.scopeId, taskId) + : await routedStore.claimTask(subagentTasks.scopeId, taskId, invocationId); + const claimed = serializeSubagentClaim(claim); + if (claimed != null) { + return JSON.stringify(claimed); + } + } else { + const command = buildSubagentControlCommand(args, action); + if (command == null) { + return JSON.stringify({ + status: 'invalid', + background_task_id: taskId, + message: 'This subagent control action is unknown or missing its required argument.', + }); + } + const result = + routedStore == null + ? subagentTasks.store.control(subagentTasks.scopeId, taskId, command) + : await routedStore.controlTask(subagentTasks.scopeId, taskId, command, invocationId); + const controlled = serializeSubagentControl(result); + if (controlled != null) { + return JSON.stringify(controlled); + } } - } else { - const command = buildSubagentControlCommand(args, action); - if (command == null) { + } catch (error) { + if (error instanceof SubagentTaskOwnerUnavailableError) { return JSON.stringify({ - status: 'invalid', + status: 'unavailable', background_task_id: taskId, - message: 'This subagent control action is unknown or missing its required argument.', + message: error.message, }); } - const controlled = serializeSubagentControl( - subagentTasks.store.control(subagentTasks.scopeId, taskId, command), - ); - if (controlled != null) { - return JSON.stringify(controlled); - } + throw error; } } @@ -1158,10 +1265,30 @@ export function runCheckBackgroundTask(params: { } const tasks = backgroundTaskRegistry.list(userId, conversationId); - const subagentTasks = - params.subagentTasks?.store - .list(params.subagentTasks.scopeId) - .map((task) => serializeSubagentSnapshot(task)) ?? []; + let subagentTasks: SerializedSubagentTask[] = []; + let listWarning: string | undefined; + if (params.subagentTasks != null) { + try { + const routedStore = routedSubagentStore(params.subagentTasks.store); + const snapshots = + routedStore == null + ? params.subagentTasks.store.list(params.subagentTasks.scopeId) + : await routedStore.listTasks(params.subagentTasks.scopeId); + subagentTasks = snapshots.map((task) => serializeSubagentSnapshot(task)); + } catch (error) { + if (error instanceof SubagentTaskOwnerUnavailableError) { + /** Cross-replica discovery is an additive source. A Redis outage must not + * hide ordinary tasks or subagents owned by this process; surface the + * incomplete view explicitly so the caller can retry for remote tasks. */ + subagentTasks = params.subagentTasks.store + .list(params.subagentTasks.scopeId) + .map((task) => serializeSubagentSnapshot(task)); + listWarning = `Cross-replica subagent tasks could not be listed: ${error.message}`; + } else { + throw error; + } + } + } logger.debug( `[background] check_background_task listed ${tasks.length + subagentTasks.length} task(s)`, ); @@ -1170,6 +1297,7 @@ export function runCheckBackgroundTask(params: { ...tasks.map((task) => serializeTask(task, { includeResult: false })), ...subagentTasks, ], + ...(listWarning != null && { partial: true, warning: listWarning }), }); } diff --git a/packages/api/src/agents/guard.spec.ts b/packages/api/src/agents/guard.spec.ts index 0673ac1071b..ee2a69c9b27 100644 --- a/packages/api/src/agents/guard.spec.ts +++ b/packages/api/src/agents/guard.spec.ts @@ -31,11 +31,13 @@ function makeStore(): SubagentThreadTaskStore { const unused = jest.fn(); return new SubagentThreadTaskStore({ acquireSubagentThreadLease: unused as AllMethods['acquireSubagentThreadLease'], + claimSubagentTaskResult: unused as AllMethods['claimSubagentTaskResult'], countActiveSubagentThreadLeases: unused as AllMethods['countActiveSubagentThreadLeases'], deleteConvos: unused as AllMethods['deleteConvos'], deleteMessages: unused as AllMethods['deleteMessages'], getConvo: unused as AllMethods['getConvo'], getMessages: unused as AllMethods['getMessages'], + listActiveSubagentThreadLeases: unused as AllMethods['listActiveSubagentThreadLeases'], releaseSubagentThreadLease: unused as AllMethods['releaseSubagentThreadLease'], reserveSubagentThread: unused as AllMethods['reserveSubagentThread'], renewSubagentThreadLease: unused as AllMethods['renewSubagentThreadLease'], diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index c99e4aea247..4e0442b4ef6 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -4089,10 +4089,13 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand const results: ToolExecuteResult[] = await Promise.all( toolCalls.map(async (tc: ToolCallRequest) => { if (backgroundControlEnabled && tc.name === CHECK_BACKGROUND_TASK_NAME) { - const pollContent = runCheckBackgroundTask({ + const pollContent = await runCheckBackgroundTask({ userId: backgroundUserId, conversationId: backgroundConversationId, args: tc.args, + toolCallId: tc.id, + agentId, + runId: `${backgroundRunId ?? ''}:${tc.turn ?? ''}`, subagentTasks, }); /** Deliver a completed task's artifact through THIS live poll diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index 95d1b34b53c..f79b560f11a 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -32,6 +32,7 @@ export * from './skills'; export * from './phases'; export * from './startup'; export * from './subagentThreads'; +export * from './subagentTaskRouting'; export * from './skillConfigurable'; export * from './skillFiles'; export * from './codeFilesSession'; diff --git a/packages/api/src/agents/subagentTaskRouting.spec.ts b/packages/api/src/agents/subagentTaskRouting.spec.ts new file mode 100644 index 00000000000..044dec0e233 --- /dev/null +++ b/packages/api/src/agents/subagentTaskRouting.spec.ts @@ -0,0 +1,1301 @@ +import { EventEmitter } from 'node:events'; +import type { + SubagentTaskControlCommand, + SubagentTaskControlResult, + SubagentTaskSnapshot, +} from '@librechat/agents'; +import type { Cluster, Redis } from 'ioredis'; +import type { SubagentTaskControlHandler } from './subagentTaskRouting'; +import { + controlFingerprint, + RedisSubagentTaskControlTransport, + SubagentTaskOwnerUnavailableError, +} from './subagentTaskRouting'; + +type MessageListener = (channel: string, message: string) => void; + +class FakeRedisBus { + readonly hashes = new Map>(); + readonly clients = new Set(); + dropResponses = 0; + /** Acknowledgements that reach nobody, as Redis reports during a resubscribe. */ + ackFailures = 0; + registrationFailures = 0; + registrationHook?: (taskId: string) => Promise; + + createClient(): FakeRedisClient { + const client = new FakeRedisClient(this); + this.clients.add(client); + return client; + } + + publish(channel: string, message: string): number { + if (this.dropResponses > 0 && channel.endsWith(':requester')) { + this.dropResponses -= 1; + return 1; + } + if (this.ackFailures > 0 && message.includes('"kind":"ack"')) { + this.ackFailures -= 1; + return 0; + } + let delivered = 0; + for (const client of this.clients) { + if (!client.disconnected && client.channels.has(channel)) { + delivered += 1; + for (const listener of client.listeners) { + queueMicrotask(() => listener(channel, message)); + } + } + } + return delivered; + } +} + +class FakeRedisClient { + readonly channels = new Set(); + readonly listeners = new Set(); + disconnected = false; + + constructor(private readonly bus: FakeRedisBus) {} + + on(event: string, listener: MessageListener): this { + if (event === 'message') { + this.listeners.add(listener); + } + return this; + } + + off(event: string, listener: MessageListener): this { + if (event === 'message') { + this.listeners.delete(listener); + } + return this; + } + + async subscribe(channel: string): Promise { + this.channels.add(channel); + return this.channels.size; + } + + async unsubscribe(channel: string): Promise { + this.channels.delete(channel); + return this.channels.size; + } + + disconnect(): void { + this.disconnected = true; + this.channels.clear(); + } + + async publish(channel: string, message: string): Promise { + return this.bus.publish(channel, message); + } + + async eval( + _script: string, + _keyCount: number, + key: string, + ...args: string[] + ): Promise { + const hash = this.bus.hashes.get(key) ?? new Map(); + if (args.length === 3) { + if (this.bus.registrationFailures > 0) { + this.bus.registrationFailures -= 1; + throw new Error('temporary registration failure'); + } + const [taskId, ownerId, ttlMs] = args; + if (this.bus.registrationHook != null) { + await this.bus.registrationHook(taskId); + } + hash.set(taskId, `${Date.now() + Number(ttlMs)}|${ownerId}`); + this.bus.hashes.set(key, hash); + return 1; + } + const readOwner = (taskId: string): string | null => { + const value = hash.get(taskId); + const separator = value?.indexOf('|') ?? -1; + const expiresAt = separator < 0 ? Number.NaN : Number(value?.slice(0, separator)); + if (value == null || !Number.isFinite(expiresAt) || expiresAt <= Date.now()) { + hash.delete(taskId); + return null; + } + return value.slice(separator + 1); + }; + if (args.length === 1) { + return readOwner(args[0]); + } + return [...hash.keys()].flatMap((taskId) => { + const ownerId = readOwner(taskId); + return ownerId == null ? [] : [taskId, ownerId]; + }); + } + + async hget(key: string, field: string): Promise { + return this.bus.hashes.get(key)?.get(field) ?? null; + } + + async hgetall(key: string): Promise> { + return Object.fromEntries(this.bus.hashes.get(key) ?? []); + } + + async hlen(key: string): Promise { + return this.bus.hashes.get(key)?.size ?? 0; + } + + async hdel(key: string, ...fields: string[]): Promise { + let deleted = 0; + for (const field of fields) { + deleted += this.bus.hashes.get(key)?.delete(field) ? 1 : 0; + } + return deleted; + } +} + +function asRedis(client: FakeRedisClient): Redis | Cluster { + return client as unknown as Redis; +} + +function snapshot(overrides: Partial = {}): SubagentTaskSnapshot { + return { + taskId: 'task-1', + threadId: 'thread-1', + subagentType: 'researcher', + status: 'running', + createdAt: 1, + updatedAt: 1, + resultAvailable: false, + resultClaimed: false, + pendingControls: 0, + ...overrides, + }; +} + +function taskHandler( + overrides: Partial = {}, +): SubagentTaskControlHandler { + return { + claim: () => ({ status: 'not_found' }), + control: () => ({ status: 'not_found' }), + list: () => [], + cancelScope: () => 0, + ...overrides, + }; +} + +describe('RedisSubagentTaskControlTransport', () => { + it('waits for the fail-fast publisher before reporting itself bound', async () => { + const bus = new FakeRedisBus(); + const publisher = new EventEmitter() as EventEmitter & { status: string }; + publisher.status = 'connecting'; + const transport = new RedisSubagentTaskControlTransport( + publisher as unknown as Redis, + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'waiting-owner' }, + ); + + let bound = false; + const binding = transport.bind(taskHandler()).then(() => { + bound = true; + }); + await Promise.resolve(); + expect(bound).toBe(false); + + publisher.status = 'ready'; + publisher.emit('ready'); + await binding; + expect(bound).toBe(true); + await transport.destroy(); + }); + + it('routes list, claim, and controls to the owner and deduplicates a retried command', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const control = jest.fn( + (_scopeId: string, _taskId: string, _command: SubagentTaskControlCommand) => + ({ + status: 'accepted', + task: snapshot(), + controlId: 'control-1', + }) satisfies SubagentTaskControlResult, + ); + const claim = jest.fn(() => ({ status: 'running', task: snapshot() }) as const); + await owner.bind(taskHandler({ claim, control, list: () => [snapshot()] })); + await requester.bind(taskHandler({ claim, control })); + await owner.registerTask('scope-1', 'task-1', 60_000); + + await expect(requester.hasTasks('scope-1')).resolves.toBe(true); + await expect(requester.list('scope-1')).resolves.toEqual([snapshot()]); + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'running', + }); + + bus.dropResponses = 1; + await expect( + requester.control( + 'scope-1', + 'task-1', + { action: 'queue', message: 'Check one more source.' }, + 'invocation-1', + ), + ).resolves.toMatchObject({ status: 'accepted', controlId: 'control-1' }); + expect(control).toHaveBeenCalledTimes(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('recomputes an idempotent list when its first response is lost', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const list = jest.fn(() => [snapshot()]); + const handler = taskHandler({ + claim: () => ({ status: 'running', task: snapshot() }) as const, + list, + }); + await owner.bind(handler); + await requester.bind({ ...handler, list: () => [] }); + await owner.registerTask('scope-1', 'task-1', 60_000); + bus.dropResponses = 1; + + await expect(requester.list('scope-1')).resolves.toEqual([snapshot()]); + expect(list).toHaveBeenCalledTimes(2); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('reports a registered but unreachable task owner as unavailable', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const handler = taskHandler(); + await owner.bind(handler); + await requester.bind(handler); + await owner.registerTask('scope-1', 'task-1', 60_000); + await owner.destroy(); + + await expect( + requester.control('scope-1', 'task-1', { action: 'cancel' }, 'invocation-dead-owner'), + ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); + await requester.destroy(); + }); + + it('delivers the largest default task result without consuming it on the owner', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const result = '\u0000'.repeat(100_000); + const claim = jest.fn(() => ({ + status: 'completed' as const, + task: snapshot({ status: 'completed', resultAvailable: true }), + result, + })); + const handler = taskHandler({ + claim, + list: () => [snapshot({ status: 'completed', resultAvailable: true })], + }); + await owner.bind(handler); + await requester.bind({ ...handler, list: () => [] }); + await owner.registerTask('scope-1', 'task-1', 60_000); + + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + result, + }); + expect(claim).toHaveBeenCalledTimes(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('keeps consuming claims when earlier results were never acknowledged', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const result = '\u0000'.repeat(100_000); + const claim = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result, + })); + await owner.bind(taskHandler({ claim })); + await requester.bind(taskHandler()); + const taskIds = Array.from({ length: 40 }, (_, index) => `task-${index + 1}`); + await Promise.all(taskIds.map((taskId) => owner.registerTask('scope-1', taskId, 60_000))); + + /** Every response is lost, so nothing is ever acknowledged or released. */ + bus.dropResponses = taskIds.length * 2; + for (const taskId of taskIds) { + await expect(requester.claim('scope-1', taskId)).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + } + + /** Retention is a fast path over a durable result, so abandoned copies bound + * themselves instead of refusing later callers until the process restarts. */ + const { claimReplays } = owner as unknown as { + claimReplays: { entries: Map; bytes: number }; + }; + expect(claim).toHaveBeenCalledTimes(taskIds.length); + expect(claimReplays.entries.size).toBeLessThanOrEqual(2_000); + expect(claimReplays.bytes).toBeLessThanOrEqual(16 * 1024 * 1024); + + bus.dropResponses = 0; + await owner.registerTask('scope-1', 'task-late', 60_000); + await expect(requester.claim('scope-1', 'task-late')).resolves.toMatchObject({ + status: 'completed', + }); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('replays one invocation and applies two identical invocations separately', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const control = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'accepted' as const, + task: snapshot({ taskId }), + controlId: 'control-1', + })); + await owner.bind(taskHandler({ control })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + const steer = { action: 'queue' as const, message: 'Check one more source.' }; + + /** The command is applied, but both responses for that invocation are lost. */ + bus.dropResponses = 2; + await expect( + requester.control('scope-1', 'task-1', steer, 'invocation-a'), + ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); + expect(control).toHaveBeenCalledTimes(1); + + /** Retransmitting that invocation replays the owner's result. */ + await expect( + requester.control('scope-1', 'task-1', steer, 'invocation-a'), + ).resolves.toMatchObject({ status: 'accepted', controlId: 'control-1' }); + expect(control).toHaveBeenCalledTimes(1); + + /** A separate invocation of the identical command is a second command. */ + await expect( + requester.control('scope-1', 'task-1', steer, 'invocation-b'), + ).resolves.toMatchObject({ status: 'accepted' }); + expect(control).toHaveBeenCalledTimes(2); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('retries and refreshes owner registration while the local task is retained', async () => { + const bus = new FakeRedisBus(); + bus.registrationFailures = 1; + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', registrationHeartbeatMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester' }, + ); + const handler = taskHandler({ + claim: () => ({ status: 'running', task: snapshot() }) as const, + list: () => [snapshot()], + }); + await owner.bind(handler); + await requester.bind({ ...handler, list: () => [] }); + await expect(owner.registerTask('scope-1', 'task-1', 60_000)).rejects.toThrow( + 'temporary registration failure', + ); + + for (let attempt = 0; attempt < 20; attempt += 1) { + if (await requester.hasTasks('scope-1')) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await expect(requester.hasTasks('scope-1')).resolves.toBe(true); + + bus.hashes.clear(); + for (let attempt = 0; attempt < 20; attempt += 1) { + if (await requester.hasTasks('scope-1')) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await expect(requester.hasTasks('scope-1')).resolves.toBe(true); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('expires a dead owner independently while another owner keeps the scope active', async () => { + const bus = new FakeRedisBus(); + const deadOwner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'dead-owner', registrationHeartbeatMs: 5 }, + ); + const liveOwner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'live-owner', registrationHeartbeatMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const deadTask = snapshot({ taskId: 'dead-task' }); + const liveTask = snapshot({ taskId: 'live-task' }); + await deadOwner.bind( + taskHandler({ claim: () => ({ status: 'running', task: deadTask }), list: () => [deadTask] }), + ); + await liveOwner.bind( + taskHandler({ claim: () => ({ status: 'running', task: liveTask }), list: () => [liveTask] }), + ); + await requester.bind(taskHandler()); + await deadOwner.registerTask('scope-1', deadTask.taskId, 20); + await liveOwner.registerTask('scope-1', liveTask.taskId, 20); + await deadOwner.destroy(); + + await new Promise((resolve) => setTimeout(resolve, 35)); + + await expect(requester.list('scope-1')).resolves.toEqual([liveTask]); + await Promise.all([liveOwner.destroy(), requester.destroy()]); + }); + + it('does not prune registered tasks omitted from a capped owner response', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const tasks = Array.from({ length: 201 }, (_, index) => + snapshot({ taskId: `task-${index + 1}` }), + ); + const handler = taskHandler({ + claim: (_scopeId: string, taskId: string) => ({ + status: 'running' as const, + task: snapshot({ taskId }), + }), + list: () => tasks, + }); + await owner.bind(handler); + await requester.bind({ ...handler, list: () => [] }); + await Promise.all(tasks.map((task) => owner.registerTask('scope-1', task.taskId, 60_000))); + + await expect(requester.list('scope-1')).resolves.toHaveLength(200); + await expect(requester.claim('scope-1', 'task-201')).resolves.toMatchObject({ + status: 'running', + task: { taskId: 'task-201' }, + }); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('refreshes owner registrations in bounded parallel batches', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', registrationHeartbeatMs: 5 }, + ); + const taskIds = Array.from({ length: 80 }, (_, index) => `task-${index + 1}`); + const [staleTaskId, ...retainedTaskIds] = taskIds; + await owner.bind( + taskHandler({ list: () => retainedTaskIds.map((taskId) => snapshot({ taskId })) }), + ); + await Promise.all(taskIds.map((taskId) => owner.registerTask('scope-1', taskId, 60_000))); + + const started: string[] = []; + let release = (): void => {}; + const gate = new Promise((resolve) => { + release = resolve; + }); + bus.registrationHook = async (taskId) => { + started.push(taskId); + await gate; + }; + + for (let attempt = 0; attempt < 100 && started.length < 32; attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + /** A serialized pass would hold exactly one refresh open; the batch bound, not + * the pass, is what limits concurrency. */ + expect(started).toHaveLength(32); + await new Promise((resolve) => setTimeout(resolve, 25)); + expect(started).toHaveLength(32); + + release(); + for ( + let attempt = 0; + attempt < 200 && new Set(started).size < retainedTaskIds.length; + attempt += 1 + ) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(new Set(started)).toEqual(new Set(retainedTaskIds)); + const [registry] = [...bus.hashes.values()]; + expect(registry.has(staleTaskId)).toBe(false); + expect(registry.size).toBe(retainedTaskIds.length); + + bus.registrationHook = undefined; + await owner.destroy(); + }); + + it('keeps refreshing other registrations when one registration fails', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', registrationHeartbeatMs: 5 }, + ); + const taskIds = ['task-1', 'task-2', 'task-3', 'task-4', 'task-5']; + await owner.bind(taskHandler({ list: () => taskIds.map((taskId) => snapshot({ taskId })) })); + await Promise.all(taskIds.map((taskId) => owner.registerTask('scope-1', taskId, 60_000))); + + bus.hashes.clear(); + bus.registrationHook = async (taskId) => { + if (taskId === 'task-1') { + throw new Error('registration failed'); + } + }; + + const healthyTaskIds = taskIds.slice(1); + for (let attempt = 0; attempt < 100; attempt += 1) { + const [registry] = [...bus.hashes.values()]; + if (registry != null && registry.size >= healthyTaskIds.length) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 5)); + } + const [registry] = [...bus.hashes.values()]; + expect([...registry.keys()].sort()).toEqual(healthyTaskIds); + + bus.registrationHook = undefined; + await owner.destroy(); + }); + + it('cancels every task in a scope beyond the model-facing list cap', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const tasks = Array.from({ length: 201 }, (_, index) => + snapshot({ taskId: `task-${index + 1}`, threadId: `thread-${index + 1}` }), + ); + const requests: Array = []; + await owner.bind( + taskHandler({ + list: () => tasks, + cancelScope: (_scopeId, threadIds) => { + requests.push(threadIds); + return threadIds == null ? tasks.length : threadIds.length; + }, + }), + ); + await requester.bind(taskHandler()); + await Promise.all(tasks.map((task) => owner.registerTask('scope-1', task.taskId, 60_000))); + + /** The model-facing list stays capped, but cancellation still reaches every task. */ + await expect(requester.list('scope-1')).resolves.toHaveLength(200); + await expect(requester.cancelScope('scope-1', null)).resolves.toBe(201); + expect(requests).toEqual([null]); + + const threadIds = tasks.map((_task, index) => `thread-${index + 1}`); + await expect(requester.cancelScope('scope-1', threadIds)).resolves.toBe(201); + expect(requests.slice(1).map((batch) => batch?.length)).toEqual([200, 1]); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('caps the aggregated list across owners rather than per owner', async () => { + const bus = new FakeRedisBus(); + const owners = ['owner-a', 'owner-b'].map( + (instanceId) => + new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId, requestTimeoutMs: 200, retryDelayMs: 10 }, + ), + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + await Promise.all( + owners.map(async (owner, ownerIndex) => { + const tasks = Array.from({ length: 150 }, (_unused, index) => + snapshot({ + taskId: `owner-${ownerIndex}-task-${index + 1}`, + threadId: `owner-${ownerIndex}-thread-${index + 1}`, + }), + ); + await owner.bind(taskHandler({ list: () => tasks })); + await Promise.all(tasks.map((task) => owner.registerTask('scope-1', task.taskId, 60_000))); + }), + ); + await requester.bind(taskHandler()); + + /** Each owner bounds its own reply, so an unbounded merge would hand the model + * every replica's batch and grow the poll response with the deployment. */ + await expect(requester.list('scope-1')).resolves.toHaveLength(200); + + await Promise.all([...owners.map((owner) => owner.destroy()), requester.destroy()]); + }); + + it('keeps running tasks when one owner caps its own reply', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + /** One owner holding more than the cap, oldest settled first: a positional slice in + * the reply drops the running children before the requester can bound anything. */ + const tasks = [ + ...Array.from({ length: 190 }, (_unused, index) => + snapshot({ + taskId: `settled-${index + 1}`, + threadId: `settled-thread-${index + 1}`, + status: 'completed', + createdAt: index + 1, + resultAvailable: true, + }), + ), + ...Array.from({ length: 30 }, (_unused, index) => + snapshot({ + taskId: `running-${index + 1}`, + threadId: `running-thread-${index + 1}`, + status: 'running', + createdAt: 1_000 + index, + }), + ), + ]; + await owner.bind(taskHandler({ list: () => tasks })); + await Promise.all(tasks.map((task) => owner.registerTask('scope-1', task.taskId, 60_000))); + await requester.bind(taskHandler()); + + const listed = await requester.list('scope-1'); + expect(listed).toHaveLength(200); + expect(listed.filter((task) => task.status === 'running')).toHaveLength(30); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('keeps running tasks when the aggregate cap drops the rest', async () => { + const bus = new FakeRedisBus(); + const owners = ['owner-old', 'owner-new'].map( + (instanceId) => + new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId, requestTimeoutMs: 200, retryDelayMs: 10 }, + ), + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + /** The settled tasks are the oldest, and the running ones the newest, so an + * oldest-first slice would drop exactly the children still worth polling. */ + const settled = Array.from({ length: 150 }, (_unused, index) => + snapshot({ + taskId: `settled-${index + 1}`, + threadId: `settled-thread-${index + 1}`, + status: 'completed', + createdAt: index + 1, + resultAvailable: true, + }), + ); + const running = Array.from({ length: 150 }, (_unused, index) => + snapshot({ + taskId: `running-${index + 1}`, + threadId: `running-thread-${index + 1}`, + status: 'running', + createdAt: 1_000 + index, + }), + ); + await Promise.all( + [settled, running].map(async (tasks, ownerIndex) => { + const owner = owners[ownerIndex]; + await owner.bind(taskHandler({ list: () => tasks })); + await Promise.all(tasks.map((task) => owner.registerTask('scope-1', task.taskId, 60_000))); + }), + ); + await requester.bind(taskHandler()); + + const listed = await requester.list('scope-1'); + expect(listed).toHaveLength(200); + expect(listed.filter((task) => task.status === 'running')).toHaveLength(150); + + await Promise.all([...owners.map((owner) => owner.destroy()), requester.destroy()]); + }); + + it('releases a claim replay once the requester acknowledges it, and keeps it otherwise', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + await owner.bind( + taskHandler({ + claim: (_scopeId: string, taskId: string) => ({ + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + }), + control: (_scopeId: string, taskId: string) => ({ + status: 'accepted' as const, + task: snapshot({ taskId }), + controlId: 'control-1', + }), + }), + ); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + await owner.registerTask('scope-1', 'task-2', 60_000); + const { claimReplays, controlReplays } = owner as unknown as { + claimReplays: { entries: Map }; + controlReplays: { entries: Map }; + }; + + /** A delivered result needs no replay copy. */ + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + }); + for (let attempt = 0; attempt < 50 && claimReplays.entries.size > 0; attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(claimReplays.entries.size).toBe(0); + + /** An undelivered one is retained, and control traffic cannot displace it. */ + bus.dropResponses = 2; + await expect(requester.claim('scope-1', 'task-2')).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + expect(claimReplays.entries.size).toBe(1); + for (let index = 0; index < 50; index += 1) { + await requester.control( + 'scope-1', + 'task-1', + { action: 'queue', message: `m-${index}` }, + `invocation-churn-${index}`, + ); + } + expect(controlReplays.entries.size).toBe(50); + expect(claimReplays.entries.size).toBe(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('lets an abandoned result expire out of retention', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 30, retryDelayMs: 5 }, + ); + let consumed = false; + const claim = jest.fn((_scopeId: string, taskId: string) => { + if (consumed) { + return { status: 'claimed' as const, task: snapshot({ taskId, status: 'completed' }) }; + } + consumed = true; + return { + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + }; + }); + await owner.bind(taskHandler({ claim })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 4 * 60 * 60_000); + + bus.dropResponses = 2; + await expect(requester.claim('scope-1', 'task-1')).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + + /** A requester that never comes back cannot hold owner memory forever: the copy + * carries an expiry, and the result stays recoverable from its durable thread. */ + const { claimReplays } = owner as unknown as { + claimReplays: { entries: Map }; + }; + expect([...claimReplays.entries.values()][0]?.expiresAt).toBeGreaterThan(Date.now()); + + const realNow = Date.now(); + const clock = jest.spyOn(Date, 'now').mockReturnValue(realNow + 6 * 60_000); + try { + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'claimed', + }); + } finally { + clock.mockRestore(); + } + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('delivers a result whose acknowledgement could not be confirmed', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const claim = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + })); + await owner.bind(taskHandler({ claim })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + const { claimReplays } = owner as unknown as { + claimReplays: { entries: Map }; + }; + + /** Every acknowledgement reaches nobody, so the owner is never told it landed. */ + bus.ackFailures = 1_000; + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + result: 'child result', + }); + /** The caller keeps the result it is holding; only the owner's copy lingers. */ + expect(claimReplays.entries.size).toBe(1); + expect(claim).toHaveBeenCalledTimes(1); + + bus.ackFailures = 0; + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + result: 'child result', + }); + for (let attempt = 0; attempt < 50 && claimReplays.entries.size > 0; attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(claimReplays.entries.size).toBe(0); + expect(claim).toHaveBeenCalledTimes(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('retries an acknowledgement that briefly reaches no subscriber', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 5 }, + ); + const claim = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + })); + await owner.bind(taskHandler({ claim })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + + /** The first two acknowledgements land during a resubscribe; the third succeeds. */ + bus.ackFailures = 2; + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + }); + expect(bus.ackFailures).toBe(0); + const { claimReplays } = owner as unknown as { + claimReplays: { entries: Map }; + }; + for (let attempt = 0; attempt < 50 && claimReplays.entries.size > 0; attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(claimReplays.entries.size).toBe(0); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('keeps a retained result addressable after its task leaves the store', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { + namespace: 'test', + instanceId: 'owner', + requestTimeoutMs: 40, + retryDelayMs: 5, + registrationHeartbeatMs: 5, + }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + let retained = true; + const claim = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + })); + await owner.bind( + taskHandler({ + claim, + /** The task ages out of the store while its result is still retained. */ + list: () => (retained ? [snapshot({ taskId: 'task-1' })] : []), + }), + ); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + + bus.dropResponses = 2; + await expect(requester.claim('scope-1', 'task-1')).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + + retained = false; + await new Promise((resolve) => setTimeout(resolve, 30)); + + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + result: 'child result', + }); + expect(claim).toHaveBeenCalledTimes(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('keeps a control fingerprint small no matter how large its message is', () => { + const large = controlFingerprint({ action: 'queue', message: 'x'.repeat(64 * 1024) }); + const other = controlFingerprint({ action: 'queue', message: 'y'.repeat(64 * 1024) }); + + /** Fingerprints are retained per invocation, so they must not carry the message. */ + expect(large).toHaveLength(43); + expect(other).toHaveLength(43); + expect(large).not.toBe(other); + expect(controlFingerprint({ action: 'queue', message: 'same' })).toBe( + controlFingerprint({ action: 'queue', message: 'same' }), + ); + }); + + it('drops a routed command whose caller already stopped waiting', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const control = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'accepted' as const, + task: snapshot({ taskId }), + controlId: 'control-1', + })); + await owner.bind(taskHandler({ control })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + const ownerChannel = [...bus.clients] + .flatMap((client) => [...client.channels]) + .find((channel) => channel.endsWith(':owner')); + expect(ownerChannel).toBeDefined(); + + /** A disconnected publisher queues an envelope offline and delivers it after the + * caller has already been told the owner was unavailable. */ + bus.publish( + ownerChannel as string, + JSON.stringify({ + version: 1, + kind: 'request', + requestId: 'stale-request', + requesterId: 'requester', + expiresAt: Date.now() - 10 * 60_000, + operation: 'control', + scopeId: 'scope-1', + taskId: 'task-1', + command: { action: 'queue', message: 'a steer the caller gave up on' }, + invocationId: 'invocation-stale', + }), + ); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(control).not.toHaveBeenCalled(); + + /** A command still inside its deadline applies normally. */ + await expect( + requester.control( + 'scope-1', + 'task-1', + { action: 'queue', message: 'Check one more source.' }, + 'invocation-fresh', + ), + ).resolves.toMatchObject({ status: 'accepted' }); + expect(control).toHaveBeenCalledTimes(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('never answers one invocation id from a different command it already ran', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const control = jest.fn( + (_scopeId: string, taskId: string, command: SubagentTaskControlCommand) => ({ + status: 'accepted' as const, + task: snapshot({ taskId }), + controlId: `control-${'message' in command ? command.message : command.action}`, + }), + ); + await owner.bind(taskHandler({ control })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + + await expect( + requester.control('scope-1', 'task-1', { action: 'queue', message: 'first' }, 'invocation-1'), + ).resolves.toMatchObject({ controlId: 'control-first' }); + + /** A retransmission of that invocation replays without applying again. */ + await expect( + requester.control('scope-1', 'task-1', { action: 'queue', message: 'first' }, 'invocation-1'), + ).resolves.toMatchObject({ controlId: 'control-first' }); + expect(control).toHaveBeenCalledTimes(1); + + /** Reusing the id for different content is a caller error, so it reaches the + * owner to be refused rather than collecting the earlier command's success. */ + await expect( + requester.control( + 'scope-1', + 'task-1', + { action: 'queue', message: 'second' }, + 'invocation-1', + ), + ).resolves.toMatchObject({ controlId: 'control-second' }); + expect(control).toHaveBeenCalledTimes(2); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('keeps one repeated provider invocation id from bleeding across tasks', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 40, retryDelayMs: 5 }, + ); + const control = jest.fn((_scopeId: string, taskId: string) => ({ + status: 'accepted' as const, + task: snapshot({ taskId }), + controlId: `control-${taskId}`, + })); + await owner.bind(taskHandler({ control })); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + await owner.registerTask('scope-2', 'task-2', 60_000); + const steer = { action: 'queue' as const, message: 'Check one more source.' }; + + /** `call_0` repeats across runs and agents, so it must not answer one task from + * another task's retained response. */ + await expect(requester.control('scope-1', 'task-1', steer, 'call_0')).resolves.toMatchObject({ + controlId: 'control-task-1', + }); + await expect(requester.control('scope-2', 'task-2', steer, 'call_0')).resolves.toMatchObject({ + controlId: 'control-task-2', + }); + expect(control).toHaveBeenCalledTimes(2); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('never retains a live claim status behind a later poll', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 200, retryDelayMs: 10 }, + ); + let settled = false; + await owner.bind( + taskHandler({ + claim: (_scopeId: string, taskId: string) => { + if (!settled) { + return { status: 'running' as const, task: snapshot({ taskId }) }; + } + return { + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + }; + }, + }), + ); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'running', + }); + settled = true; + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + result: 'child result', + }); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); + + it('returns a consumed result to a later claim after both responses are lost', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', requestTimeoutMs: 60, retryDelayMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 60, retryDelayMs: 5 }, + ); + let claims = 0; + await owner.bind( + taskHandler({ + claim: (_scopeId: string, taskId: string) => { + claims += 1; + return claims === 1 + ? { + status: 'completed' as const, + task: snapshot({ taskId, status: 'completed', resultAvailable: true }), + result: 'child result', + } + : { + status: 'claimed' as const, + task: snapshot({ taskId, status: 'completed', resultClaimed: true }), + }; + }, + }), + ); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + + /** Both the first response and its retry are lost after the owner consumed the result. */ + bus.dropResponses = 2; + await expect(requester.claim('scope-1', 'task-1')).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + expect(claims).toBe(1); + + await expect(requester.claim('scope-1', 'task-1')).resolves.toMatchObject({ + status: 'completed', + result: 'child result', + }); + expect(claims).toBe(1); + + await Promise.all([owner.destroy(), requester.destroy()]); + }); +}); diff --git a/packages/api/src/agents/subagentTaskRouting.ts b/packages/api/src/agents/subagentTaskRouting.ts new file mode 100644 index 00000000000..7496bbdafd9 --- /dev/null +++ b/packages/api/src/agents/subagentTaskRouting.ts @@ -0,0 +1,1379 @@ +import { logger } from '@librechat/data-schemas'; +import { createHash, randomUUID } from 'node:crypto'; +import type { + SubagentTaskClaim, + SubagentTaskControlCommand, + SubagentTaskControlResult, + SubagentTaskSnapshot, +} from '@librechat/agents'; +import type { Cluster, Redis } from 'ioredis'; +import { createConcurrencyLimiter } from '~/utils/promise'; + +const PROTOCOL_VERSION = 1; +const DEFAULT_REQUEST_TIMEOUT_MS = 2_000; +const DEFAULT_RETRY_DELAY_MS = 500; +const DEFAULT_READY_TIMEOUT_MS = 10_000; +const DEFAULT_REGISTRATION_HEARTBEAT_MS = 10_000; +const MAX_PENDING_REQUESTS = 1_000; +/** A consumed claim is retained apart from control replays so unrelated command + * traffic cannot displace it while its caller retries. Retention is a fast path, not + * the guarantee: the terminal result is recoverable from its durable child message. */ +const MAX_CLAIM_REPLAY_ENTRIES = 2_000; +const MAX_CLAIM_REPLAY_BYTES = 16 * 1024 * 1024; +const MAX_CONTROL_REPLAY_ENTRIES = 2_000; +const MAX_CONTROL_REPLAY_BYTES = 4 * 1024 * 1024; +const RESPONSE_CACHE_TTL_MS = 5 * 60_000; +/** Absorbs ordinary clock drift between replicas when honouring a request deadline. */ +const REQUEST_CLOCK_SKEW_MS = 30_000; +const MAX_SCOPE_ID_CHARS = 4_096; +const MAX_TASK_ID_CHARS = 256; +const MAX_CONTROL_MESSAGE_CHARS = 64 * 1_024; +const MAX_RESULT_CHARS = 100_000; +const MAX_ERROR_CHARS = 4 * 1_024; +const MAX_THREAD_ID_CHARS = 256; +const MAX_SUBAGENT_TYPE_CHARS = 256; +const MAX_PROGRESS_LABEL_CHARS = 1_024; +/** Bounds the model-facing task list, per owner reply and across the merged result. */ +export const MAX_TASK_SNAPSHOTS = 200; +const MAX_CANCEL_THREAD_IDS = 200; +/** Matches the deletion drain so bounded fan-out stays well inside the lease TTL. */ +const ROUTING_FANOUT_CONCURRENCY = 32; +/** Contains every bounded response even when JSON escapes each retained character. */ +const MAX_ROUTED_MESSAGE_CHARS = 8 * 1_024 * 1_024; + +const REGISTER_TASK_SCRIPT = + "local now = redis.call('TIME'); " + + 'local ttl = tonumber(ARGV[3]); ' + + 'local expiresAt = (tonumber(now[1]) * 1000) + math.floor(tonumber(now[2]) / 1000) + ttl; ' + + "redis.call('HSET', KEYS[1], ARGV[1], tostring(expiresAt) .. '|' .. ARGV[2]); " + + "local directoryTtl = redis.call('PTTL', KEYS[1]); " + + "if directoryTtl < ttl then redis.call('PEXPIRE', KEYS[1], ttl); end; " + + 'return 1'; + +const READ_ACTIVE_REGISTRATIONS_SCRIPT = + "local now = redis.call('TIME'); " + + 'local nowMs = (tonumber(now[1]) * 1000) + math.floor(tonumber(now[2]) / 1000); ' + + "local entries = redis.call('HGETALL', KEYS[1]); " + + 'local active = {}; ' + + 'for i = 1, #entries, 2 do ' + + 'local value = entries[i + 1]; ' + + "local separator = string.find(value, '|', 1, true); " + + 'local expiresAt = separator and tonumber(string.sub(value, 1, separator - 1)); ' + + 'if expiresAt and expiresAt > nowMs then ' + + 'table.insert(active, entries[i]); ' + + 'table.insert(active, string.sub(value, separator + 1)); ' + + "else redis.call('HDEL', KEYS[1], entries[i]); end; " + + 'end; ' + + 'return active'; + +const READ_TASK_OWNER_SCRIPT = + "local value = redis.call('HGET', KEYS[1], ARGV[1]); " + + 'if not value then return nil; end; ' + + "local separator = string.find(value, '|', 1, true); " + + 'local expiresAt = separator and tonumber(string.sub(value, 1, separator - 1)); ' + + "local now = redis.call('TIME'); " + + 'local nowMs = (tonumber(now[1]) * 1000) + math.floor(tonumber(now[2]) / 1000); ' + + "if not expiresAt or expiresAt <= nowMs then redis.call('HDEL', KEYS[1], ARGV[1]); return nil; end; " + + 'return string.sub(value, separator + 1)'; + +type RedisClient = Redis | Cluster; +interface RoutedRequestBase { + version: typeof PROTOCOL_VERSION; + kind: 'request'; + requestId: string; + requesterId: string; + scopeId: string; + /** Epoch milliseconds after which the requester has stopped waiting. */ + expiresAt: number; +} + +type RoutedRequest = RoutedRequestBase & + ( + | { operation: 'claim'; taskId: string } + | { + operation: 'control'; + taskId: string; + command: SubagentTaskControlCommand; + invocationId: string; + } + | { operation: 'list' } + | { operation: 'cancel'; threadIds: string[] | null } + ); + +type RoutedRequestPayload = + | { operation: 'claim'; scopeId: string; taskId: string } + | { + operation: 'control'; + scopeId: string; + taskId: string; + command: SubagentTaskControlCommand; + invocationId: string; + } + | { operation: 'list'; scopeId: string } + | { operation: 'cancel'; scopeId: string; threadIds: string[] | null }; + +interface RoutedResponse { + version: typeof PROTOCOL_VERSION; + kind: 'response'; + requestId: string; + ok: boolean; + result?: unknown; +} + +/** Tells the owner a consumed result reached a caller and no longer needs retaining. */ +interface RoutedAck { + version: typeof PROTOCOL_VERSION; + kind: 'ack'; + scopeId: string; + taskId: string; +} + +interface PendingRequest { + resolve: (value: unknown) => void; + reject: (error: Error) => void; + retry: ReturnType; + timeout: ReturnType; +} + +interface CachedResponse { + value: string; + bytes: number; + expiresAt: number; + /** Content this response answered, so one id cannot replay a different command. */ + fingerprint?: string; +} + +interface ReplayCache { + entries: Map; + bytes: number; + maxEntries: number; + maxBytes: number; +} + +interface RoutedTaskList { + snapshots: SubagentTaskSnapshot[]; + truncated: boolean; +} + +interface RoutedCancelResult { + cancelled: number; +} + +interface OwnedTaskRegistration { + scopeId: string; + taskId: string; + ttlMs: number; +} + +export interface SubagentTaskControlHandler { + claim(scopeId: string, taskId: string): SubagentTaskClaim; + control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): SubagentTaskControlResult; + list(scopeId: string): SubagentTaskSnapshot[]; + cancelScope(scopeId: string, threadIds: string[] | null): number; +} + +/** Optional host transport for reaching the process that owns a live child task. */ +export interface SubagentTaskControlTransport { + bind(handler: SubagentTaskControlHandler): Promise; + registerTask(scopeId: string, taskId: string, ttlMs: number): Promise; + hasTasks(scopeId: string): Promise; + claim(scopeId: string, taskId: string): Promise; + control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise; + list(scopeId: string): Promise; + cancelScope(scopeId: string, threadIds: string[] | null): Promise; + destroy(): Promise; +} + +export class SubagentTaskOwnerUnavailableError extends Error { + constructor() { + super('The process running this subagent task is temporarily unavailable.'); + } +} + +export interface RedisSubagentTaskControlTransportOptions { + /** Separates pub/sub channels for deployments sharing one Redis service. */ + namespace?: string; + instanceId?: string; + requestTimeoutMs?: number; + retryDelayMs?: number; + registrationHeartbeatMs?: number; +} + +function positiveInteger(value: number | undefined, fallback: number): number { + return Number.isSafeInteger(value) && value != null && value > 0 ? value : fallback; +} + +function shortHash(value: string): string { + return createHash('sha256').update(value).digest('base64url').slice(0, 24); +} + +function isBoundedString(value: unknown, maxChars: number): value is string { + return typeof value === 'string' && value.length > 0 && value.length <= maxChars; +} + +function isStringWithin(value: unknown, maxChars: number): value is string { + return typeof value === 'string' && value.length <= maxChars; +} + +function truncateMiddle(value: string, maxChars: number): string { + if (value.length <= maxChars) { + return value; + } + const marker = '\n…[truncated]…\n'; + const available = Math.max(0, maxChars - marker.length); + const head = Math.ceil(available / 2); + return `${value.slice(0, head)}${marker}${value.slice(value.length - (available - head))}`; +} + +function boundedSnapshot(snapshot: SubagentTaskSnapshot): SubagentTaskSnapshot { + return { + taskId: truncateMiddle(snapshot.taskId, MAX_TASK_ID_CHARS), + ...(snapshot.threadId == null + ? {} + : { threadId: truncateMiddle(snapshot.threadId, MAX_THREAD_ID_CHARS) }), + subagentType: truncateMiddle(snapshot.subagentType, MAX_SUBAGENT_TYPE_CHARS), + status: snapshot.status, + createdAt: snapshot.createdAt, + updatedAt: snapshot.updatedAt, + resultAvailable: snapshot.resultAvailable, + resultClaimed: snapshot.resultClaimed, + pendingControls: snapshot.pendingControls, + ...(snapshot.progress == null + ? {} + : { + progress: { + ...snapshot.progress, + ...(snapshot.progress.label == null + ? {} + : { label: truncateMiddle(snapshot.progress.label, MAX_PROGRESS_LABEL_CHARS) }), + }, + }), + ...(snapshot.error == null ? {} : { error: truncateMiddle(snapshot.error, MAX_ERROR_CHARS) }), + }; +} + +/** + * Bounds a model-facing task list, keeping what a caller can still act on: running + * children first, then the most recent settled results. A plain oldest-first slice + * would drop the newest tasks, hiding a child that just started from the only tool + * able to poll it. + */ +export function boundedTaskList(tasks: SubagentTaskSnapshot[]): SubagentTaskSnapshot[] { + const byCreatedAt = (left: SubagentTaskSnapshot, right: SubagentTaskSnapshot): number => + left.createdAt - right.createdAt; + if (tasks.length <= MAX_TASK_SNAPSHOTS) { + return tasks.sort(byCreatedAt); + } + const running: SubagentTaskSnapshot[] = []; + const settled: SubagentTaskSnapshot[] = []; + for (const task of tasks) { + (task.status === 'running' ? running : settled).push(task); + } + running.sort(byCreatedAt); + const keptRunning = running.slice(-MAX_TASK_SNAPSHOTS); + const remaining = MAX_TASK_SNAPSHOTS - keptRunning.length; + if (remaining <= 0) { + return keptRunning; + } + settled.sort(byCreatedAt); + return [...keptRunning, ...settled.slice(-remaining)].sort(byCreatedAt); +} + +/** Applies the routed result and snapshot bounds to a claim from any source. */ +export function boundedClaim(claim: SubagentTaskClaim): SubagentTaskClaim { + if (claim.status === 'not_found') { + return claim; + } + const task = boundedSnapshot(claim.task); + if (claim.status === 'completed') { + return { status: 'completed', task, result: truncateMiddle(claim.result, MAX_RESULT_CHARS) }; + } + if (claim.status === 'error' || claim.status === 'cancelled') { + return { status: claim.status, task, error: truncateMiddle(claim.error, MAX_ERROR_CHARS) }; + } + return { status: claim.status, task }; +} + +function boundedControlResult(result: SubagentTaskControlResult): SubagentTaskControlResult { + if (result.status === 'not_found') { + return result; + } + if (result.status === 'invalid') { + return { status: 'invalid', message: truncateMiddle(result.message, MAX_ERROR_CHARS) }; + } + return { + status: result.status, + task: boundedSnapshot(result.task), + ...(result.status === 'accepted' && result.controlId != null + ? { controlId: truncateMiddle(result.controlId, MAX_TASK_ID_CHARS) } + : {}), + }; +} + +async function waitForRedisConnectionReady(client: RedisClient): Promise { + if (client.status == null || client.status === 'ready') { + return; + } + if (client.status === 'end') { + throw new SubagentTaskOwnerUnavailableError(); + } + await new Promise((resolve, reject) => { + const cleanup = () => { + clearTimeout(timeout); + client.off('ready', onReady); + client.off('end', onEnd); + }; + const onReady = () => { + cleanup(); + resolve(); + }; + const onEnd = () => { + cleanup(); + reject(new SubagentTaskOwnerUnavailableError()); + }; + const timeout = setTimeout(() => { + cleanup(); + reject(new SubagentTaskOwnerUnavailableError()); + }, DEFAULT_READY_TIMEOUT_MS); + timeout.unref?.(); + client.once('ready', onReady); + client.once('end', onEnd); + /** Close the status-check/listener-registration race: ioredis may become ready + * synchronously between the check above and installing these listeners. */ + if (client.status === 'ready') { + onReady(); + } else if (client.status === 'end') { + onEnd(); + } else if (client.status === 'wait') { + client.connect().catch(onEnd); + } + }); +} + +async function waitForRedisReady( + client: RedisClient, + options: { eagerClusterMasters?: boolean } = {}, +): Promise { + await waitForRedisConnectionReady(client); + if (options.eagerClusterMasters !== true || !client.isCluster) { + return; + } + /** A ready Cluster has a slot map but its per-master connections are lazy. The + * fail-fast publisher cannot admit requests until every possible write target is + * connected; otherwise the first command to a cold shard would be rejected. */ + await Promise.all( + (client as Cluster).nodes('master').map((node) => waitForRedisConnectionReady(node)), + ); +} + +function isSnapshot(value: unknown): value is SubagentTaskSnapshot { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + return ( + isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS) && + typeof candidate.subagentType === 'string' && + ['running', 'completed', 'error', 'cancelled'].includes(candidate.status ?? '') && + typeof candidate.createdAt === 'number' && + typeof candidate.updatedAt === 'number' && + typeof candidate.resultAvailable === 'boolean' && + typeof candidate.resultClaimed === 'boolean' && + typeof candidate.pendingControls === 'number' + ); +} + +function isClaim(value: unknown): value is SubagentTaskClaim { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + if (candidate.status === 'not_found') { + return true; + } + if (!('task' in candidate) || !isSnapshot(candidate.task)) { + return false; + } + if (candidate.status === 'completed') { + return 'result' in candidate && typeof candidate.result === 'string'; + } + if (candidate.status === 'error' || candidate.status === 'cancelled') { + return 'error' in candidate && typeof candidate.error === 'string'; + } + return candidate.status === 'running' || candidate.status === 'claimed'; +} + +function isControlResult(value: unknown): value is SubagentTaskControlResult { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + if (candidate.status === 'not_found') { + return true; + } + if (candidate.status === 'invalid') { + return typeof candidate.message === 'string'; + } + return ( + ['accepted', 'cancelled', 'not_running', 'control_not_found'].includes( + candidate.status ?? '', + ) && + 'task' in candidate && + isSnapshot(candidate.task) + ); +} + +function isCancelThreadIds(value: unknown): value is string[] { + return ( + Array.isArray(value) && + value.length > 0 && + value.length <= MAX_CANCEL_THREAD_IDS && + value.every((threadId) => isBoundedString(threadId, MAX_THREAD_ID_CHARS)) + ); +} + +function isCancelResult(value: unknown): value is RoutedCancelResult { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const { cancelled } = value as Partial; + return Number.isSafeInteger(cancelled) && (cancelled as number) >= 0; +} + +function controlContent(command: SubagentTaskControlCommand): string { + if (command.action === 'cancel') { + return 'cancel'; + } + if (command.action === 'cancel_message') { + return `cancel_message\u0000${command.controlId}`; + } + return `${command.action}\u0000${command.message}`; +} + +/** + * Canonical identity of one control's content. Property order cannot vary it, so the + * transport and the owning task store agree on when two commands are the same, and it + * is hashed so retaining one costs a fixed few bytes rather than a whole message. + */ +export function controlFingerprint(command: SubagentTaskControlCommand): string { + return createHash('sha256').update(controlContent(command)).digest('base64url'); +} + +/** True once a claim has consumed the task's one-shot terminal result. */ +function consumesResult(result: SubagentTaskClaim): boolean { + return ( + result.status === 'completed' || result.status === 'error' || result.status === 'cancelled' + ); +} + +function isRoutedTaskList(value: unknown): value is RoutedTaskList { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + return ( + Array.isArray(candidate.snapshots) && + candidate.snapshots.every(isSnapshot) && + typeof candidate.truncated === 'boolean' + ); +} + +function parseControlCommand(value: unknown): SubagentTaskControlCommand | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as { action?: unknown; message?: unknown; controlId?: unknown }; + if (candidate.action === 'cancel') { + return { action: 'cancel' }; + } + if (candidate.action === 'cancel_message') { + return isStringWithin(candidate.controlId, MAX_TASK_ID_CHARS) + ? { action: 'cancel_message', controlId: candidate.controlId } + : undefined; + } + if ( + (candidate.action === 'steer' || + candidate.action === 'queue' || + candidate.action === 'interrupt') && + isStringWithin(candidate.message, MAX_CONTROL_MESSAGE_CHARS) + ) { + return { action: candidate.action, message: candidate.message }; + } + return undefined; +} + +function failureResponse(requestId: string): string { + const response: RoutedResponse = { + version: PROTOCOL_VERSION, + kind: 'response', + requestId, + ok: false, + }; + return JSON.stringify(response); +} + +function successResponse(requestId: string, result: string): string { + return `{"version":${PROTOCOL_VERSION},"kind":"response","requestId":${JSON.stringify( + requestId, + )},"ok":true,"result":${result}}`; +} + +function parseRequest(value: unknown): RoutedRequest | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as { + version?: unknown; + kind?: unknown; + requestId?: unknown; + requesterId?: unknown; + operation?: unknown; + scopeId?: unknown; + taskId?: unknown; + command?: unknown; + threadIds?: unknown; + invocationId?: unknown; + expiresAt?: unknown; + }; + if ( + candidate.version !== PROTOCOL_VERSION || + candidate.kind !== 'request' || + !isBoundedString(candidate.requestId, 128) || + !isBoundedString(candidate.requesterId, 128) || + !['claim', 'control', 'list', 'cancel'].includes( + typeof candidate.operation === 'string' ? candidate.operation : '', + ) || + !isBoundedString(candidate.scopeId, MAX_SCOPE_ID_CHARS) || + !Number.isSafeInteger(candidate.expiresAt) + ) { + return undefined; + } + const expiresAt = candidate.expiresAt as number; + if (candidate.operation === 'list') { + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'list', + scopeId: candidate.scopeId, + }; + } + if (candidate.operation === 'cancel') { + if (candidate.threadIds !== null && !isCancelThreadIds(candidate.threadIds)) { + return undefined; + } + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'cancel', + scopeId: candidate.scopeId, + threadIds: candidate.threadIds, + }; + } + if (!isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS)) { + return undefined; + } + if (candidate.operation === 'claim') { + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'claim', + scopeId: candidate.scopeId, + taskId: candidate.taskId, + }; + } + const command = parseControlCommand(candidate.command); + if (command == null || !isBoundedString(candidate.invocationId, 128)) { + return undefined; + } + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'control', + scopeId: candidate.scopeId, + taskId: candidate.taskId, + command, + invocationId: candidate.invocationId, + }; +} + +function createReplayCache(maxEntries: number, maxBytes: number): ReplayCache { + return { entries: new Map(), bytes: 0, maxEntries, maxBytes }; +} + +function parseAck(value: unknown): RoutedAck | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as Partial; + if ( + candidate.version !== PROTOCOL_VERSION || + candidate.kind !== 'ack' || + !isBoundedString(candidate.scopeId, MAX_SCOPE_ID_CHARS) || + !isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS) + ) { + return undefined; + } + return { + version: PROTOCOL_VERSION, + kind: 'ack', + scopeId: candidate.scopeId, + taskId: candidate.taskId, + }; +} + +function parseResponse(value: unknown): RoutedResponse | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as Partial; + if ( + candidate.version !== PROTOCOL_VERSION || + candidate.kind !== 'response' || + !isBoundedString(candidate.requestId, 128) || + typeof candidate.ok !== 'boolean' + ) { + return undefined; + } + return candidate as RoutedResponse; +} + +/** + * Routes bounded live-task operations to their owning API replica. Redis keeps + * only an expiring owner directory and request/reply envelopes; the executor, + * transcript, and checkpoint never move between processes. + */ +export class RedisSubagentTaskControlTransport implements SubagentTaskControlTransport { + private readonly instanceId: string; + private readonly namespaceHash: string; + private readonly requestTimeoutMs: number; + private readonly retryDelayMs: number; + private readonly registrationHeartbeatMs: number; + private readonly pending = new Map(); + private readonly claimReplays = createReplayCache( + MAX_CLAIM_REPLAY_ENTRIES, + MAX_CLAIM_REPLAY_BYTES, + ); + + private readonly controlReplays = createReplayCache( + MAX_CONTROL_REPLAY_ENTRIES, + MAX_CONTROL_REPLAY_BYTES, + ); + + private readonly ownedTasks = new Map(); + private handler?: SubagentTaskControlHandler; + private ready?: Promise; + private registrationHeartbeat?: ReturnType; + private registrationRefresh?: Promise; + private destroyed = false; + + constructor( + private readonly publisher: RedisClient, + private readonly subscriber: RedisClient, + options: RedisSubagentTaskControlTransportOptions = {}, + ) { + this.instanceId = options.instanceId?.trim() || randomUUID(); + this.namespaceHash = shortHash(options.namespace?.trim() || 'default'); + this.requestTimeoutMs = positiveInteger(options.requestTimeoutMs, DEFAULT_REQUEST_TIMEOUT_MS); + this.retryDelayMs = Math.min( + positiveInteger(options.retryDelayMs, DEFAULT_RETRY_DELAY_MS), + Math.max(1, Math.floor(this.requestTimeoutMs / 2)), + ); + this.registrationHeartbeatMs = positiveInteger( + options.registrationHeartbeatMs, + DEFAULT_REGISTRATION_HEARTBEAT_MS, + ); + } + + async bind(handler: SubagentTaskControlHandler): Promise { + if (this.destroyed) { + throw new Error('Subagent task control transport is closed.'); + } + if (this.handler != null) { + throw new Error('Subagent task control transport is already bound.'); + } + this.handler = handler; + /** The publisher fails fast instead of queueing commands, so opening HTTP + * admission before it is ready would turn healthy startup lag into false + * `unavailable` results. Both dedicated connections are part of readiness. */ + await Promise.all([ + waitForRedisReady(this.publisher, { eagerClusterMasters: true }), + waitForRedisReady(this.subscriber), + ]); + this.subscriber.on('message', this.onMessage); + this.ready = this.subscriber.subscribe(this.channel(this.instanceId)).then(() => undefined); + await this.ready; + } + + async registerTask(scopeId: string, taskId: string, ttlMs: number): Promise { + this.assertTaskAddress(scopeId, taskId); + const registration = { + scopeId, + taskId, + ttlMs: positiveInteger(ttlMs, 1), + }; + this.ownedTasks.set(this.registrationKey(scopeId, taskId), registration); + this.ensureRegistrationHeartbeat(); + await this.publishRegistration(registration); + } + + async hasTasks(scopeId: string): Promise { + this.assertScope(scopeId); + await this.requireReady(); + try { + return Object.keys(await this.readActiveRegistrations(scopeId)).length > 0; + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to inspect the task owner directory', error); + throw new SubagentTaskOwnerUnavailableError(); + } + } + + async claim(scopeId: string, taskId: string): Promise { + const routed = await this.requestTaskOwner(scopeId, taskId, 'claim'); + if (routed == null) { + return undefined; + } + const { ownerId, result } = routed; + if (!isClaim(result)) { + throw new SubagentTaskOwnerUnavailableError(); + } + if (result.status === 'not_found') { + await this.removeRegistrations(scopeId, [taskId]); + return undefined; + } + if (consumesResult(result)) { + /** Frees the owner's retained copy immediately. An acknowledgement that cannot + * be confirmed only leaves that copy to expire, so the caller still keeps the + * result it is holding rather than trading it for a retry. */ + await this.acknowledgeClaim(ownerId, scopeId, taskId); + } + return result; + } + + async control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise { + const routed = await this.requestTaskOwner(scopeId, taskId, 'control', command, invocationId); + if (routed == null) { + return undefined; + } + const { result } = routed; + if (!isControlResult(result)) { + throw new SubagentTaskOwnerUnavailableError(); + } + if (result.status === 'not_found') { + await this.removeRegistrations(scopeId, [taskId]); + return undefined; + } + return result; + } + + async list(scopeId: string): Promise { + this.assertScope(scopeId); + await this.requireReady(); + let ownersByTask: Record; + try { + ownersByTask = await this.readActiveRegistrations(scopeId); + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to read the task owner directory', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const owners = new Set(Object.values(ownersByTask)); + owners.delete(this.instanceId); + if (owners.size === 0) { + return []; + } + const ownerIds = [...owners]; + const results = await Promise.all( + ownerIds.map((ownerId) => this.sendRequest(ownerId, { operation: 'list', scopeId })), + ); + const snapshots: SubagentTaskSnapshot[] = []; + const staleTaskIds: string[] = []; + for (const [index, value] of results.entries()) { + if (!isRoutedTaskList(value)) { + throw new SubagentTaskOwnerUnavailableError(); + } + const ownerId = ownerIds[index]; + const reportedTaskIds = new Set(value.snapshots.map((snapshot) => snapshot.taskId)); + for (const snapshot of value.snapshots) { + if (ownersByTask[snapshot.taskId] === ownerId) { + snapshots.push(snapshot); + } + } + if (!value.truncated) { + for (const [taskId, registeredOwnerId] of Object.entries(ownersByTask)) { + if (registeredOwnerId === ownerId && !reportedTaskIds.has(taskId)) { + staleTaskIds.push(taskId); + } + } + } + } + if (staleTaskIds.length > 0) { + await this.removeRegistrations(scopeId, staleTaskIds); + } + /** Each owner bounds its own reply, so without an aggregate cap this grows with the + * number of replicas holding the scope. Bounding after the loop rather than during + * it keeps the sweep above reading every owner's reply, and lets the cap choose by + * status instead of by whichever owner answered first. */ + return boundedTaskList(snapshots); + } + + /** + * Cancels live children on every other owner of this scope. The owner applies the + * predicate to its complete local task set, so deletion never depends on the + * bounded model-facing list and cannot miss a task beyond that cap. + */ + async cancelScope(scopeId: string, threadIds: string[] | null): Promise { + this.assertScope(scopeId); + if (threadIds != null && threadIds.length === 0) { + return 0; + } + await this.requireReady(); + let ownersByTask: Record; + try { + ownersByTask = await this.readActiveRegistrations(scopeId); + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to read the task owner directory', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const owners = new Set(Object.values(ownersByTask)); + owners.delete(this.instanceId); + if (owners.size === 0) { + return 0; + } + const batches: Array = []; + if (threadIds == null) { + batches.push(null); + } else { + for (let index = 0; index < threadIds.length; index += MAX_CANCEL_THREAD_IDS) { + batches.push(threadIds.slice(index, index + MAX_CANCEL_THREAD_IDS)); + } + } + const cancelSlot = createConcurrencyLimiter(ROUTING_FANOUT_CONCURRENCY); + const requests: Array> = []; + for (const ownerId of owners) { + for (const batch of batches) { + requests.push( + cancelSlot(() => + this.sendRequest(ownerId, { operation: 'cancel', scopeId, threadIds: batch }), + ), + ); + } + } + let cancelled = 0; + for (const value of await Promise.all(requests)) { + if (!isCancelResult(value)) { + throw new SubagentTaskOwnerUnavailableError(); + } + cancelled += value.cancelled; + } + return cancelled; + } + + async destroy(): Promise { + if (this.destroyed) { + return; + } + this.destroyed = true; + for (const pending of this.pending.values()) { + clearTimeout(pending.retry); + clearTimeout(pending.timeout); + pending.reject(new SubagentTaskOwnerUnavailableError()); + } + this.pending.clear(); + for (const cache of [this.claimReplays, this.controlReplays]) { + cache.entries.clear(); + cache.bytes = 0; + } + this.ownedTasks.clear(); + if (this.registrationHeartbeat != null) { + clearInterval(this.registrationHeartbeat); + this.registrationHeartbeat = undefined; + } + this.subscriber.off('message', this.onMessage); + await this.subscriber.unsubscribe(this.channel(this.instanceId)).catch(() => undefined); + this.subscriber.disconnect(); + } + + private readonly onMessage = (channel: string, message: string): void => { + if (channel !== this.channel(this.instanceId) || message.length > MAX_ROUTED_MESSAGE_CHARS) { + return; + } + let parsed: unknown; + try { + parsed = JSON.parse(message) as unknown; + } catch { + return; + } + const response = parseResponse(parsed); + if (response != null) { + this.handleResponse(response); + return; + } + const ack = parseAck(parsed); + if (ack != null) { + this.releaseClaimReplay(ack.scopeId, ack.taskId); + return; + } + const request = parseRequest(parsed); + if (request != null) { + void this.handleRequest(request).catch((error) => { + logger.warn('[subagentTaskRouting] Failed to answer a routed command', error); + }); + } + }; + + private handleResponse(response: RoutedResponse): void { + const pending = this.pending.get(response.requestId); + if (pending == null) { + return; + } + this.pending.delete(response.requestId); + clearTimeout(pending.retry); + clearTimeout(pending.timeout); + if (!response.ok) { + pending.reject(new SubagentTaskOwnerUnavailableError()); + return; + } + pending.resolve(response.result); + } + + private async handleRequest(request: RoutedRequest): Promise { + if (Date.now() > request.expiresAt + REQUEST_CLOCK_SKEW_MS) { + /** The caller stopped waiting for this long ago and has been told it was + * unavailable, so applying it now would steer a child it believes untouched. */ + logger.warn('[subagentTaskRouting] Dropped a routed command past its deadline'); + return; + } + const replay = this.replayFor(request); + const cached = replay?.cache.entries.get(replay.key); + /** A retransmission replays; the same id carrying different content is a caller + * error, so it reaches the owner, which refuses it, rather than being answered + * from the earlier command's response. */ + if ( + cached != null && + cached.expiresAt > Date.now() && + cached.fingerprint === replay?.fingerprint + ) { + await this.publish( + this.channel(request.requesterId), + successResponse(request.requestId, cached.value), + ); + return; + } + const handler = this.handler; + if (handler == null) { + return; + } + let serialized: string; + try { + let result: + | SubagentTaskClaim + | SubagentTaskControlResult + | RoutedTaskList + | RoutedCancelResult; + /** A claim that consumed nothing stays uncached so a later poll still observes + * the task's live status. */ + let replayable = replay != null; + if (request.operation === 'list') { + const tasks = handler.list(request.scopeId); + /** Bounded the same way the requester bounds the merge: a positional slice here + * would drop this owner's running children before they ever reached it. */ + const bounded = boundedTaskList(tasks); + result = { + snapshots: bounded.map(boundedSnapshot), + truncated: tasks.length > bounded.length, + }; + } else if (request.operation === 'cancel') { + result = { cancelled: handler.cancelScope(request.scopeId, request.threadIds) }; + } else if (request.operation === 'claim') { + const claim = boundedClaim(handler.claim(request.scopeId, request.taskId)); + replayable = consumesResult(claim); + result = claim; + } else { + result = boundedControlResult( + handler.control(request.scopeId, request.taskId, request.command, request.invocationId), + ); + } + const serializedResult = JSON.stringify(result); + /** Retaining the result rather than the envelope lets a later caller retry, + * which carries its own correlation id, recover a response it never received. */ + if (replay != null && replayable) { + this.retainReplay(replay.cache, replay.key, serializedResult, replay.fingerprint); + } + serialized = successResponse(request.requestId, serializedResult); + } catch (error) { + logger.error('[subagentTaskRouting] Owner failed to process a routed command', error); + serialized = failureResponse(request.requestId); + } + await this.publish(this.channel(request.requesterId), serialized); + } + + /** + * Locates a destructive operation's replay slot. A claim consumes the one-shot + * result, so it is keyed by operation—stable across callers and replicas, so a + * later poll still resolves to the response the owner produced. Lists are + * idempotent and recomputable, so their large bodies are never retained. + */ + private replayFor( + request: RoutedRequest, + ): { cache: ReplayCache; key: string; fingerprint?: string } | undefined { + if (request.operation === 'list') { + return undefined; + } + if (request.operation === 'claim') { + return { + cache: this.claimReplays, + key: this.claimReplayKey(request.scopeId, request.taskId), + }; + } + if (request.operation === 'cancel') { + return { cache: this.controlReplays, key: `cancel\u0000${request.requestId}` }; + } + /** Task-scoped: a provider tool-call id such as `call_0` repeats across runs and + * agents, so keying on it alone would answer one task from another's snapshot. */ + return { + cache: this.controlReplays, + key: `control\u0000${shortHash(request.scopeId)}\u0000${request.taskId}\u0000${request.invocationId}`, + fingerprint: controlFingerprint(request.command), + }; + } + + private claimReplayKey(scopeId: string, taskId: string): string { + return `claim\u0000${shortHash(scopeId)}\u0000${taskId}`; + } + + /** + * Releases a retained result once a caller confirms holding it, so a delivered + * result frees its slot immediately instead of waiting out the replay window. + */ + private releaseClaimReplay(scopeId: string, taskId: string): void { + const key = this.claimReplayKey(scopeId, taskId); + const cached = this.claimReplays.entries.get(key); + if (cached == null) { + return; + } + this.claimReplays.entries.delete(key); + this.claimReplays.bytes -= cached.bytes; + } + + /** + * Tells the owner it may release a delivered result. Delivery to zero subscribers is + * not an acknowledgement, so this retries inside the ordinary request window. + */ + private async acknowledgeClaim(ownerId: string, scopeId: string, taskId: string): Promise { + const ack: RoutedAck = { version: PROTOCOL_VERSION, kind: 'ack', scopeId, taskId }; + const serialized = JSON.stringify(ack); + const destination = this.channel(ownerId); + const deadline = Date.now() + this.requestTimeoutMs; + for (;;) { + try { + if ((await this.publish(destination, serialized)) > 0) { + return; + } + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to acknowledge a claimed result', error); + } + if (Date.now() + this.retryDelayMs >= deadline) { + return; + } + await new Promise((resolve) => { + const timer = setTimeout(resolve, this.retryDelayMs); + timer.unref?.(); + }); + } + } + + private async requestTaskOwner( + scopeId: string, + taskId: string, + operation: 'claim' | 'control', + command?: SubagentTaskControlCommand, + invocationId?: string, + ): Promise<{ ownerId: string; result: unknown } | undefined> { + this.assertTaskAddress(scopeId, taskId); + await this.requireReady(); + let ownerId: string | null; + try { + ownerId = (await this.publisher.eval( + READ_TASK_OWNER_SCRIPT, + 1, + this.registryKey(scopeId), + taskId, + )) as string | null; + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to resolve the task owner', error); + throw new SubagentTaskOwnerUnavailableError(); + } + if (!isBoundedString(ownerId, 128)) { + return undefined; + } + if (operation === 'claim') { + return { ownerId, result: await this.sendRequest(ownerId, { operation, scopeId, taskId }) }; + } + if (command == null || invocationId == null) { + throw new Error('A routed subagent control command and invocation id are required.'); + } + return { + ownerId, + result: await this.sendRequest(ownerId, { + operation, + scopeId, + taskId, + command, + invocationId, + }), + }; + } + + private async sendRequest(ownerId: string, request: RoutedRequestPayload): Promise { + await this.requireReady(); + if (this.pending.size >= MAX_PENDING_REQUESTS) { + throw new SubagentTaskOwnerUnavailableError(); + } + const requestId = randomUUID(); + /** Carried so a request the caller has stopped waiting for cannot be applied + * later: a disconnected publisher queues the envelope offline and delivers it + * after this deadline, by which time the caller has been told `unavailable`. */ + const envelope: RoutedRequest = { + version: PROTOCOL_VERSION, + kind: 'request', + requestId, + requesterId: this.instanceId, + expiresAt: Date.now() + this.requestTimeoutMs, + ...request, + }; + const serialized = JSON.stringify(envelope); + const destination = this.channel(ownerId); + return new Promise((resolve, reject) => { + const retry = setTimeout(() => { + void this.publish(destination, serialized).catch((error) => { + logger.warn('[subagentTaskRouting] Routed command retry failed', error); + }); + }, this.retryDelayMs); + retry.unref?.(); + const timeout = setTimeout(() => { + this.pending.delete(requestId); + clearTimeout(retry); + reject(new SubagentTaskOwnerUnavailableError()); + }, this.requestTimeoutMs); + timeout.unref?.(); + this.pending.set(requestId, { resolve, reject, retry, timeout }); + void this.publish(destination, serialized).catch((error) => { + logger.warn('[subagentTaskRouting] Routed command publish failed', error); + }); + }); + } + + private pruneExpiredReplays(cache: ReplayCache): void { + const now = Date.now(); + for (const [id, cached] of cache.entries) { + if (cached.expiresAt != null && cached.expiresAt <= now) { + cache.entries.delete(id); + cache.bytes -= cached.bytes; + } + } + } + + private retainReplay(cache: ReplayCache, key: string, value: string, fingerprint?: string): void { + const bytes = Buffer.byteLength(value, 'utf8'); + if (bytes > cache.maxBytes) { + return; + } + this.pruneExpiredReplays(cache); + /** Replacing a key is not an additional entry: leaving the old one counted would + * inflate the cache's byte total permanently and evict unrelated responses. */ + const replaced = cache.entries.get(key); + if (replaced != null) { + cache.entries.delete(key); + cache.bytes -= replaced.bytes; + } + while (cache.entries.size >= cache.maxEntries || cache.bytes + bytes > cache.maxBytes) { + const oldest = cache.entries.keys().next().value as string | undefined; + if (oldest == null) { + break; + } + const evicted = cache.entries.get(oldest); + cache.entries.delete(oldest); + cache.bytes -= evicted?.bytes ?? 0; + } + cache.entries.set(key, { + value, + bytes, + expiresAt: Date.now() + RESPONSE_CACHE_TTL_MS, + ...(fingerprint == null ? {} : { fingerprint }), + }); + cache.bytes += bytes; + } + + private async publish(channel: string, value: string): Promise { + const delivered = await this.publisher.publish(channel, value); + return typeof delivered === 'number' ? delivered : 0; + } + + private ensureRegistrationHeartbeat(): void { + if (this.registrationHeartbeat != null || this.destroyed) { + return; + } + this.registrationHeartbeat = setInterval(() => { + if (this.registrationRefresh != null) { + return; + } + const refresh = this.refreshRegistrations() + .catch((error) => { + logger.warn('[subagentTaskRouting] Failed to refresh child-task owners', error); + }) + .finally(() => { + if (this.registrationRefresh === refresh) { + this.registrationRefresh = undefined; + } + }); + this.registrationRefresh = refresh; + }, this.registrationHeartbeatMs); + this.registrationHeartbeat.unref?.(); + } + + private async refreshRegistrations(): Promise { + const handler = this.handler; + if (handler == null || this.destroyed || this.ownedTasks.size === 0) { + return; + } + const localTaskIdsByScope = new Map>(); + const staleTaskIdsByScope = new Map(); + const retained: OwnedTaskRegistration[] = []; + for (const registration of this.ownedTasks.values()) { + const { scopeId, taskId } = registration; + let localTaskIds = localTaskIdsByScope.get(scopeId); + if (localTaskIds == null) { + localTaskIds = new Set(handler.list(scopeId).map((task) => task.taskId)); + localTaskIdsByScope.set(scopeId, localTaskIds); + } + /** A retained result is only reachable while its owner stays registered, so the + * address outlives the task itself until the result is acknowledged. */ + if ( + localTaskIds.has(taskId) || + this.claimReplays.entries.has(this.claimReplayKey(scopeId, taskId)) + ) { + retained.push(registration); + continue; + } + this.ownedTasks.delete(this.registrationKey(scopeId, taskId)); + const staleTaskIds = staleTaskIdsByScope.get(scopeId) ?? []; + staleTaskIds.push(taskId); + staleTaskIdsByScope.set(scopeId, staleTaskIds); + } + /** Serializing one EVAL per registration can outlast the lease TTL, so a pass + * refreshes in bounded parallel batches and one failure cannot cancel the rest. */ + const refreshSlot = createConcurrencyLimiter(ROUTING_FANOUT_CONCURRENCY); + await Promise.all([ + ...[...staleTaskIdsByScope].map(([scopeId, taskIds]) => + refreshSlot(() => this.removeRegistrations(scopeId, taskIds)), + ), + ...retained.map((registration) => + refreshSlot(() => + this.publishRegistration(registration).catch((error) => { + logger.warn('[subagentTaskRouting] Failed to refresh a child-task owner', error); + }), + ), + ), + ]); + } + + private async publishRegistration(registration: OwnedTaskRegistration): Promise { + await this.requireReady(); + await this.publisher.eval( + REGISTER_TASK_SCRIPT, + 1, + this.registryKey(registration.scopeId), + registration.taskId, + this.instanceId, + registration.ttlMs.toString(), + ); + } + + private async readActiveRegistrations(scopeId: string): Promise> { + const value = (await this.publisher.eval( + READ_ACTIVE_REGISTRATIONS_SCRIPT, + 1, + this.registryKey(scopeId), + )) as unknown; + if (!Array.isArray(value) || value.length % 2 !== 0) { + throw new SubagentTaskOwnerUnavailableError(); + } + const registrations: Record = {}; + for (let index = 0; index < value.length; index += 2) { + const taskId = value[index]; + const ownerId = value[index + 1]; + if (!isBoundedString(taskId, MAX_TASK_ID_CHARS) || !isBoundedString(ownerId, 128)) { + throw new SubagentTaskOwnerUnavailableError(); + } + registrations[taskId] = ownerId; + } + return registrations; + } + + private async removeRegistrations(scopeId: string, taskIds: string[]): Promise { + if (taskIds.length === 0) { + return; + } + await this.publisher.hdel(this.registryKey(scopeId), ...taskIds).catch((error) => { + logger.warn('[subagentTaskRouting] Failed to prune stale task owners', error); + }); + } + + private registryKey(scopeId: string): string { + return `subagent-task:{${shortHash(scopeId)}}:owners`; + } + + private registrationKey(scopeId: string, taskId: string): string { + return `${scopeId}\u0000${taskId}`; + } + + private channel(instanceId: string): string { + return `subagent-task-control:${this.namespaceHash}:${instanceId}`; + } + + private assertScope(scopeId: string): void { + if (!isBoundedString(scopeId, MAX_SCOPE_ID_CHARS)) { + throw new Error('Invalid subagent task routing scope.'); + } + } + + private assertTaskAddress(scopeId: string, taskId: string): void { + this.assertScope(scopeId); + if (!isBoundedString(taskId, MAX_TASK_ID_CHARS)) { + throw new Error('Invalid subagent task routing identity.'); + } + } + + private async requireReady(): Promise { + if (this.destroyed || this.ready == null) { + throw new SubagentTaskOwnerUnavailableError(); + } + await this.ready; + } +} diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index b75ec3b7ee4..a68b3bb0545 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -5,14 +5,27 @@ import { Constants, EModelEndpoint } from 'librechat-data-provider'; import { createMethods, createModels, logger } from '@librechat/data-schemas'; import { AIMessage, HumanMessage } from '@librechat/agents/langchain/messages'; import type { + SubagentTaskClaim, + SubagentTaskControlCommand, + SubagentTaskControlResult, SubagentTaskRuntime, + SubagentTaskSnapshot, SubagentTaskStartRequest, SubagentTaskStartResult, } from '@librechat/agents'; import type { AllMethods, IConversation, IMessage } from '@librechat/data-schemas'; import type { BaseMessage } from '@librechat/agents/langchain/messages'; +import type { + SubagentTaskControlHandler, + SubagentTaskControlTransport, +} from './subagentTaskRouting'; import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; -import { buildSubagentThreadTaskConfig, SubagentThreadTaskStore } from './subagentThreads'; +import { + buildSubagentThreadTaskConfig, + createSubagentThreadTaskStore, + SubagentThreadTaskStore, +} from './subagentThreads'; +import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting'; import { createSubagentAttemptKey } from './subagentThreadIds'; import { createSubagentUsageSink } from './usage'; @@ -20,6 +33,75 @@ let mongod: MongoMemoryServer; let methods: AllMethods; let loggerErrorSpy: jest.SpyInstance; +class TestTaskRoutingHub { + readonly owners = new Map(); + + key(scopeId: string, taskId: string): string { + return `${scopeId}\u0000${taskId}`; + } +} + +class TestTaskControlTransport implements SubagentTaskControlTransport { + private handler?: SubagentTaskControlHandler; + readonly registrations: Array<{ scopeId: string; taskId: string; ttlMs: number }> = []; + + constructor(private readonly hub: TestTaskRoutingHub) {} + + async bind(handler: SubagentTaskControlHandler): Promise { + this.handler = handler; + } + + async registerTask(scopeId: string, taskId: string, ttlMs: number): Promise { + this.registrations.push({ scopeId, taskId, ttlMs }); + this.hub.owners.set(this.hub.key(scopeId, taskId), this); + } + + async hasTasks(scopeId: string): Promise { + const prefix = `${scopeId}\u0000`; + return [...this.hub.owners.keys()].some((key) => key.startsWith(prefix)); + } + + async claim(scopeId: string, taskId: string): Promise { + return this.hub.owners.get(this.hub.key(scopeId, taskId))?.handler?.claim(scopeId, taskId); + } + + async control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + _invocationId: string, + ): Promise { + return this.hub.owners + .get(this.hub.key(scopeId, taskId)) + ?.handler?.control(scopeId, taskId, command, _invocationId); + } + + async list(scopeId: string): Promise { + return [...this.remoteOwners(scopeId)].flatMap((owner) => owner.handler?.list(scopeId) ?? []); + } + + async cancelScope(scopeId: string, threadIds: string[] | null): Promise { + let cancelled = 0; + for (const owner of this.remoteOwners(scopeId)) { + cancelled += owner.handler?.cancelScope(scopeId, threadIds) ?? 0; + } + return cancelled; + } + + async destroy(): Promise {} + + private remoteOwners(scopeId: string): Set { + const owners = new Set(); + const prefix = `${scopeId}\u0000`; + for (const [key, owner] of this.hub.owners) { + if (key.startsWith(prefix) && owner !== this) { + owners.add(owner); + } + } + return owners; + } +} + function taskRequest( scopeId: string, overrides: Partial = {}, @@ -51,6 +133,32 @@ function taskRequest( }; } +function replayTransport(claim: SubagentTaskClaim): SubagentTaskControlTransport { + return { + bind: async () => undefined, + registerTask: async () => undefined, + hasTasks: async () => true, + claim: async () => claim, + control: async () => undefined, + list: async () => [], + cancelScope: async () => 0, + destroy: async () => undefined, + }; +} + +function threadSnapshot(taskId: string): SubagentTaskSnapshot { + return { + taskId, + subagentType: 'researcher', + status: 'cancelled', + createdAt: 1, + updatedAt: 2, + resultAvailable: false, + resultClaimed: true, + pendingControls: 0, + }; +} + async function waitForSettled( store: SubagentThreadTaskStore, scopeId: string, @@ -67,6 +175,16 @@ async function waitForSettled( throw new Error('Timed out waiting for the subagent task.'); } +async function waitUntil(condition: () => boolean, description: string): Promise { + for (let attempt = 0; attempt < 400; attempt += 1) { + if (condition()) { + return; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error(`Timed out waiting for ${description}.`); +} + function requireAccepted( started: SubagentTaskStartResult, ): Extract { @@ -583,8 +701,26 @@ describe('SubagentThreadTaskStore', () => { const preparationRelease = new Promise((resolve) => { releasePreparation = resolve; }); + let leaseDeadline = new Date(0); + let renewedPastDeadline = false; const slowMethods = { ...methods, + acquireSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + const acquired = await methods.acquireSubagentThreadLease(...args); + if (acquired) { + leaseDeadline = args[0].expiresAt; + } + return acquired; + }, + ), + renewSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + const renewed = await methods.renewSubagentThreadLease(...args); + renewedPastDeadline ||= renewed && args[0].now > leaseDeadline; + return renewed; + }, + ), getMessages: jest.fn(async (...args: Parameters) => { if (blockNextRead && args[0].conversationId === slowThreadId) { blockNextRead = false; @@ -594,7 +730,7 @@ describe('SubagentThreadTaskStore', () => { return methods.getMessages(...args); }), }; - const options = { leaseTtlMs: 60, leaseHeartbeatMs: 10 }; + const options = { leaseTtlMs: 500, leaseHeartbeatMs: 50 }; const firstWorker = new SubagentThreadTaskStore(slowMethods, options); const secondWorker = new SubagentThreadTaskStore(methods, options); const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); @@ -612,7 +748,9 @@ describe('SubagentThreadTaskStore', () => { }), ); await preparing; - await new Promise((resolve) => setTimeout(resolve, 100)); + /** Wait for evidence rather than a fixed delay: a renewal that succeeds after the + * acquired lease's own deadline proves the heartbeat carried it past expiry. */ + await waitUntil(() => renewedPastDeadline, 'the shared lease to outlive its original deadline'); const overlappingRun = jest.fn(taskRequest(config.scopeId).run); const overlapping = secondWorker.start( @@ -630,6 +768,61 @@ describe('SubagentThreadTaskStore', () => { expect(firstRun).toHaveBeenCalledTimes(1); }); + it('cancels a child when its lease renewal only commits after expiry', async () => { + const userId = 'late-lease-renewal-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let providerEntered = false; + let previousExpiry = 0; + let markLateRenewal = (): void => undefined; + const lateRenewal = new Promise((resolve) => { + markLateRenewal = resolve; + }); + const slowMethods = { + ...methods, + renewSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + if (providerEntered) { + markLateRenewal(); + /** Wait on the last confirmed lease deadline rather than a tiny fixed TTL: + * the renewal definitely commits after the gap, without assuming how fast + * a loaded runner completes preparation and its first Mongo write. */ + await new Promise((resolve) => + setTimeout(resolve, Math.max(0, previousExpiry - Date.now() + 10)), + ); + } + const renewed = await methods.renewSubagentThreadLease(...args); + if (renewed) { + previousExpiry = args[0].expiresAt.getTime(); + } + return renewed; + }, + ), + }; + const store = new SubagentThreadTaskStore(slowMethods, { + leaseTtlMs: 1_000, + leaseHeartbeatMs: 20, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(async (runtime: SubagentTaskRuntime) => { + providerEntered = true; + return new Promise<{ content: string }>((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }); + }); + const started = store.start(taskRequest(config.scopeId, { run })); + + await lateRenewal; + await waitForSettled(store, config.scopeId, started); + + expect(run).toHaveBeenCalledTimes(1); + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'cancelled', + }); + }); + it('rechecks account deletion after acquiring the shared lease', async () => { const userId = 'lease-fence-gap-user'; const parentConversationId = randomUUID(); @@ -1297,6 +1490,999 @@ describe('SubagentThreadTaskStore', () => { expect(JSON.stringify(loggerErrorSpy.mock.calls)).not.toContain('provider-secret'); }); + it('routes live task polling and controls to the replica that owns the execution', async () => { + const userId = 'routed-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const requesterStore = new SubagentThreadTaskStore(methods); + const ownerTransport = new TestTaskControlTransport(hub); + await ownerStore.configureTaskControlTransport(ownerTransport); + await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => result, + }), + ); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + expect(ownerTransport.registrations).toContainEqual({ + scopeId: config.scopeId, + taskId, + ttlMs: 30_000, + }); + + await expect(requesterStore.hasTasks(config.scopeId)).resolves.toBe(true); + await expect(requesterStore.listTasks(config.scopeId)).resolves.toEqual([ + expect.objectContaining({ taskId, status: 'running' }), + ]); + await expect( + requesterStore.controlTask(config.scopeId, taskId, { + action: 'queue', + message: 'Verify the primary source too.', + }), + ).resolves.toMatchObject({ status: 'accepted' }); + + finish({ content: 'Cross-replica result.' }); + await waitForSettled(ownerStore, config.scopeId, started); + await expect(requesterStore.claimTask(config.scopeId, taskId)).resolves.toMatchObject({ + status: 'completed', + result: 'Cross-replica result.', + }); + await expect(requesterStore.claimTask(config.scopeId, taskId)).resolves.toMatchObject({ + status: 'claimed', + }); + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + requesterStore.destroyTaskControlTransport(), + ]); + }); + + it('applies one control invocation once whether it arrives locally or through routing', async () => { + const userId = 'invocation-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const requesterStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => result, + }), + ); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + + const steer = { action: 'queue' as const, message: 'Verify the primary source too.' }; + const routed = await requesterStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'); + expect(routed).toMatchObject({ status: 'accepted' }); + + /** The same invocation reaching the owner directly replays that result rather than + * queueing a second steer, so local and routed callers agree. */ + await expect( + ownerStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'), + ).resolves.toEqual(routed); + expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(1); + + /** Reusing one invocation id for different content is a caller error, not a retry. */ + await expect( + requesterStore.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Something else entirely.' }, + 'invocation-1', + ), + ).resolves.toMatchObject({ status: 'invalid' }); + expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(1); + + finish({ content: 'Cross-replica result.' }); + await waitForSettled(ownerStore, config.scopeId, started); + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + requesterStore.destroyTaskControlTransport(), + ]); + }); + + it('fails a child closed when its owner address cannot be published', async () => { + const userId = 'unregistered-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const transport = new TestTaskControlTransport(hub); + transport.registerTask = async () => { + throw new Error('registration failed'); + }; + const store = new SubagentThreadTaskStore(methods); + await store.configureTaskControlTransport(transport); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(async () => ({ content: 'never reached' })); + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + /** An unaddressable child cannot be polled, controlled, or cancelled, so no + * provider work may start behind a failed registration. */ + expect(run).not.toHaveBeenCalled(); + expect(store.get(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'error', + }); + await store.destroyTaskControlTransport(); + }); + + it('returns a lost result to the same poll invocation and refuses a different one', async () => { + const userId = 'durable-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const requesterStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Cross-replica result.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(ownerStore, config.scopeId, started); + + await expect(requesterStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { status: 'completed', result: 'Cross-replica result.' }, + ); + + /** The owner's one-shot result is gone, but the child's durable thread still holds + * it, so the invocation that already collected it recovers its own result. */ + await expect(requesterStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { status: 'completed', result: 'Cross-replica result.' }, + ); + + /** A different invocation is told it was collected rather than handed a copy. */ + await expect(requesterStore.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject( + { status: 'claimed' }, + ); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + requesterStore.destroyTaskControlTransport(), + ]); + }); + + it('bounds a result recovered from its durable child message', async () => { + const userId = 'large-result-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'x'.repeat(150_000) }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject({ + status: 'completed', + }); + + /** The durable message keeps the child's untruncated output, so recovering it + * must apply the same bound a routed response would have. */ + const recovered = await store.claimTask(config.scopeId, taskId, 'poll-1'); + expect(recovered.status).toBe('completed'); + if (recovered.status === 'completed') { + expect(recovered.result.length).toBeLessThanOrEqual(100_000); + } + }); + + it('recovers a durable result after the owning process and registration are gone', async () => { + const userId = 'restarted-owner-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const ownerStore = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Recovered without owner memory.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(ownerStore, config.scopeId, started); + + /** A fresh store has neither the in-memory task nor a Redis owner registration. */ + const restartedStore = new SubagentThreadTaskStore(methods); + const unrelatedParentConversationId = randomUUID(); + await saveParent(userId, unrelatedParentConversationId); + const unrelatedConfig = buildSubagentThreadTaskConfig(restartedStore, { + userId, + parentConversationId: unrelatedParentConversationId, + }); + await expect( + restartedStore.claimTask(unrelatedConfig.scopeId, taskId, 'wrong-parent-poll'), + ).resolves.toEqual({ status: 'not_found' }); + + await expect(restartedStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { + status: 'completed', + result: 'Recovered without owner memory.', + }, + ); + await expect(restartedStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { + status: 'completed', + result: 'Recovered without owner memory.', + }, + ); + await expect(restartedStore.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject( + { + status: 'claimed', + }, + ); + }); + + it('tells a second invocation a retained result was already collected', async () => { + const userId = 'duplicate-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Only one caller may hold this.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + const threadId = requireThreadId(started); + + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject({ + status: 'completed', + result: 'Only one caller may hold this.', + }); + + /** An owner replaying a retained response would hand the same terminal result to + * another invocation; the durable record decides, so that one is told it was + * already collected rather than being given a second copy. */ + const replayingStore = new SubagentThreadTaskStore(methods); + await replayingStore.configureTaskControlTransport( + replayTransport({ + status: 'completed', + task: { ...threadSnapshot(taskId), threadId, status: 'completed' }, + result: 'Only one caller may hold this.', + }), + ); + await expect(replayingStore.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject( + { status: 'claimed' }, + ); + + /** The invocation that already holds it still recovers its own result. */ + await expect(replayingStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { status: 'completed', result: 'Only one caller may hold this.' }, + ); + + await replayingStore.destroyTaskControlTransport(); + }); + + it('refuses to build a store the host wired without a required method', () => { + const { claimSubagentTaskResult: _omitted, ...incomplete } = methods; + + /** The host wires this from JavaScript, so a missing method has to fail at + * startup rather than as an unavailable result the first time a task settles. */ + expect(() => + createSubagentThreadTaskStore( + incomplete as unknown as Parameters[0], + ), + ).toThrow('claimSubagentTaskResult'); + expect(() => createSubagentThreadTaskStore(methods)).not.toThrow(); + }); + + it('renews its own fence while a long deletion is still running', async () => { + const userId = 'long-deletion-user'; + const renewals: string[] = []; + const store = new SubagentThreadTaskStore(methods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 30, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + renewOwnerAdmission: async (_userId: string, token: string) => { + renewals.push(token); + return true; + }, + releaseOwnerAdmission: async () => undefined, + }); + const listLeases = jest.spyOn(methods, 'listActiveSubagentThreadLeases').mockResolvedValue([]); + try { + await store.withOwnerDeletionFence(userId, undefined, async () => { + /** A deletion outlasting its 90ms fence window must not let the fence lapse. */ + await new Promise((resolve) => setTimeout(resolve, 300)); + return 'deleted'; + }); + } finally { + listLeases.mockRestore(); + } + + expect(renewals.length).toBeGreaterThan(0); + expect(new Set(renewals).size).toBe(1); + }); + + it('re-fences and drains again after a fence gap during deletion', async () => { + const userId = 'deletion-gap-user'; + let recoveryAllowed = false; + const fenceOwnerAdmission = jest.fn(async () => undefined); + const listActiveSubagentThreadLeases = jest.fn(async () => []); + const testMethods = { ...methods, listActiveSubagentThreadLeases }; + const renewOwnerAdmission = jest.fn(async () => { + if (!recoveryAllowed) { + throw new Error('database temporarily unavailable'); + } + return false; + }); + const store = new SubagentThreadTaskStore(testMethods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 30, + ownerFenceGraceMs: 60, + fenceOwnerAdmission, + renewOwnerAdmission, + releaseOwnerAdmission: async () => undefined, + }); + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + await new Promise((resolve) => setTimeout(resolve, 150)); + recoveryAllowed = true; + return 'deleted'; + }), + ).resolves.toBe('deleted'); + + expect(fenceOwnerAdmission).toHaveBeenCalledTimes(2); + expect(listActiveSubagentThreadLeases).toHaveBeenCalledTimes(2); + }); + + it('does not report deletion success when post-gap recovery fails', async () => { + const userId = 'deletion-gap-failure-user'; + const listActiveSubagentThreadLeases = jest.fn(async () => []); + const testMethods = { ...methods, listActiveSubagentThreadLeases }; + const store = new SubagentThreadTaskStore(testMethods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 30, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + renewOwnerAdmission: async () => { + throw new Error('database unavailable'); + }, + releaseOwnerAdmission: async () => undefined, + }); + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + await new Promise((resolve) => setTimeout(resolve, 150)); + return 'deleted'; + }), + ).rejects.toThrow('database unavailable'); + }); + + it('cancels a grandchild whose own conversation the cascade removed', async () => { + const userId = 'cascade-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods, { maxThreadDepth: 3 }); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const childConversationId = randomUUID(); + await saveParent(userId, childConversationId, { + subagentThread: { + rootConversationId: parentConversationId, + parentConversationId, + parentAgentId: 'parent-agent', + subagentType: 'researcher', + depth: 1, + }, + }); + /** The grandchild runs inside the child's scope, which a plan naming only the + * deleted root never covers. */ + const config = buildSubagentThreadTaskConfig(ownerStore, { + userId, + parentConversationId: childConversationId, + }); + let finish = (_value: { content: string }): void => undefined; + const running = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start(taskRequest(config.scopeId, { run: async () => running })); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + + const plan = await deletingStore.planCancellationForConversations(userId, [ + parentConversationId, + ]); + await expect( + deletingStore.cancelPlan(plan, [parentConversationId, childConversationId]), + ).resolves.toBeGreaterThanOrEqual(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + finish({ content: 'late' }); + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('reports a result as unavailable when its collection cannot be recorded', async () => { + const userId = 'unrecordable-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Recorded before it is handed over.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + + /** Handing the result over without recording its claimant would let another + * invocation collect the same one-shot output once the database recovers. */ + const claimResult = jest + .spyOn(methods, 'claimSubagentTaskResult') + .mockRejectedValueOnce(new Error('database unavailable')); + try { + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + } finally { + claimResult.mockRestore(); + } + + /** The result stays unclaimed, so a later poll still collects it exactly once. */ + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject({ + status: 'completed', + result: 'Recorded before it is handed over.', + }); + await expect(store.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject({ + status: 'claimed', + }); + }); + + it('keeps a live task’s control invocation when settled tasks fill the window', async () => { + const userId = 'invocation-eviction-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { + maxControlInvocations: 2, + completedTtlMs: 20, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const running = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const live = store.start(taskRequest(config.scopeId, { run: async () => running })); + const liveTaskId = requireAccepted(live).task.taskId; + const settled = store.start( + taskRequest(config.scopeId, { run: async () => ({ content: 'done' }) }), + ); + const settledTaskId = requireAccepted(settled).task.taskId; + await waitForSettled(store, config.scopeId, settled); + + const steer = { action: 'queue' as const, message: 'Verify the primary source too.' }; + const applied = store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-live'); + expect(applied).toMatchObject({ status: 'accepted' }); + store.controlInvocation(config.scopeId, settledTaskId, steer, 'invocation-settled'); + + for (let attempt = 0; attempt < 100; attempt += 1) { + if (store.get(config.scopeId, settledTaskId) == null) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + expect(store.get(config.scopeId, settledTaskId)).toBeUndefined(); + + /** The window is full, so admitting another invocation sweeps the records of tasks + * this store no longer holds. The live task's record survives, so a caller + * retrying it replays instead of steering that child a second time. */ + store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-later'); + expect(store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-live')).toEqual( + applied, + ); + expect(store.get(config.scopeId, liveTaskId)?.pendingControls).toBe(2); + + /** With every remaining record belonging to a live task, a further invocation is + * refused rather than displacing one: applying it unrecorded would let its own + * retry apply the command twice. */ + expect( + store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-third'), + ).toMatchObject({ status: 'invalid' }); + expect(store.get(config.scopeId, liveTaskId)?.pendingControls).toBe(2); + + finish({ content: 'done' }); + await waitForSettled(store, config.scopeId, live); + }); + + it('caps the merged local and remote task list the poll tool reads', async () => { + const userId = 'merged-list-cap-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + /** The base store caps concurrent runs twice over — ten per scope and a hundred + * across the store — and this test is about what the merge returns rather than + * about admission, so both are raised to admit every task it starts. */ + const store = new SubagentThreadTaskStore(methods, { + maxRunningPerScope: 150, + maxRunningTotal: 150, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const remote = Array.from({ length: 150 }, (_unused, index) => + threadSnapshot(`remote-task-${index + 1}`), + ); + await store.configureTaskControlTransport({ + ...replayTransport({ status: 'claimed', task: threadSnapshot('remote-task-1') }), + list: async () => remote, + }); + + const local = await Promise.all( + Array.from({ length: 150 }, () => store.start(taskRequest(config.scopeId))), + ); + await Promise.all(local.map((started) => waitForSettled(store, config.scopeId, started))); + expect(store.list(config.scopeId)).toHaveLength(150); + + /** Each owner's reply and the remote aggregation are bounded on their own, but the + * poll tool reads this merge — 300 distinct tasks must still arrive as 200. */ + await expect(store.listTasks(config.scopeId)).resolves.toHaveLength(200); + + await store.destroyTaskControlTransport(); + }); + + it('routes a control for a remote task while the local invocation window is full', async () => { + const userId = 'remote-control-under-load-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { maxControlInvocations: 1 }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const remoteResult: SubagentTaskControlResult = { + status: 'cancelled', + task: threadSnapshot('remote-task'), + }; + const routed = jest.fn(async () => remoteResult); + await store.configureTaskControlTransport({ + ...replayTransport({ status: 'claimed', task: threadSnapshot('remote-task') }), + control: routed, + }); + + let finish = (_value: { content: string }): void => undefined; + const running = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const live = store.start(taskRequest(config.scopeId, { run: async () => running })); + const liveTaskId = requireAccepted(live).task.taskId; + const steer = { action: 'queue' as const, message: 'Check the changelog as well.' }; + expect(store.controlInvocation(config.scopeId, liveTaskId, steer, 'local-1')).toMatchObject({ + status: 'accepted', + }); + + /** The window holds a live task's record and cannot be swept, but a task this + * replica never owned is the remote owner's to refuse or apply. */ + await expect( + store.controlTask(config.scopeId, 'remote-task', { action: 'cancel' }, 'remote-1'), + ).resolves.toEqual(remoteResult); + expect(routed).toHaveBeenCalledWith( + config.scopeId, + 'remote-task', + { action: 'cancel' }, + 'remote-1', + ); + + finish({ content: 'done' }); + await waitForSettled(store, config.scopeId, live); + await store.destroyTaskControlTransport(); + }); + + it('fails a deletion closed when the admission fence cannot be held', async () => { + const userId = 'fence-lapse-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { + ownerDrainTimeoutMs: 60, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + renewOwnerAdmission: async () => { + throw new Error('database unavailable'); + }, + releaseOwnerAdmission: async () => undefined, + }); + /** A drain that outlasts the 120ms fence window while every renewal rejects: the + * last confirmed `fencedUntil` passes and nothing is left holding admission shut. */ + const leases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockImplementationOnce(async () => { + await new Promise((resolve) => setTimeout(resolve, 200)); + return []; + }); + const deletion = jest.fn(async () => 'deleted'); + try { + await expect(store.withOwnerDeletionFence(userId, undefined, deletion)).rejects.toThrow( + 'admission fence expired', + ); + /** Nothing was removed, so the caller can retry once the fence holds again. */ + expect(deletion).not.toHaveBeenCalled(); + } finally { + leases.mockRestore(); + } + }); + + it('treats a renewal that lands after its own deadline as a lapse', async () => { + const userId = 'fence-late-renewal-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let renewals = 0; + const store = new SubagentThreadTaskStore(methods, { + ownerDrainTimeoutMs: 60, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + /** Succeeds, but the first write only lands well past the 120ms deadline it was + * meant to extend — admission stood open for the difference. */ + renewOwnerAdmission: async () => { + renewals += 1; + if (renewals === 1) { + await new Promise((resolve) => setTimeout(resolve, 150)); + } + return true; + }, + releaseOwnerAdmission: async () => undefined, + }); + const leases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockImplementationOnce(async () => { + await new Promise((resolve) => setTimeout(resolve, 250)); + return []; + }); + const deletion = jest.fn(async () => 'deleted'); + try { + await expect(store.withOwnerDeletionFence(userId, undefined, deletion)).rejects.toThrow( + 'admission fence expired', + ); + /** Every renewal reported success, so a deadline restored from the write's own + * start time would have read as continuously fenced. */ + expect(renewals).toBeGreaterThan(0); + expect(deletion).not.toHaveBeenCalled(); + } finally { + leases.mockRestore(); + } + }); + + it('releases the owner fence after an in-flight renewal instead of racing it', async () => { + const userId = 'fence-renewal-race-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let releaseRenewal = (): void => undefined; + const renewalBlocked = new Promise((resolve) => { + releaseRenewal = resolve; + }); + let markRenewing = (): void => undefined; + const renewing = new Promise((resolve) => { + markRenewing = resolve; + }); + const order: string[] = []; + const fenceOwnerAdmission = jest.fn(async () => { + order.push('fence'); + }); + /** The renewal is still waiting on the database when the deletion finishes, and it + * reports the fence lost — the shape that used to leave a fresh, unreleasable one. */ + let renewalAttempts = 0; + const renewOwnerAdmission = jest.fn(async () => { + renewalAttempts += 1; + markRenewing(); + await renewalBlocked; + order.push('renew'); + /** The in-flight renewal discovers the entry missing and re-takes it; the + * recovery renewal then confirms that replacement while the second drain runs. */ + return renewalAttempts > 1; + }); + const releaseOwnerAdmission = jest.fn(async () => { + order.push('release'); + }); + const testMethods = { + ...methods, + listActiveSubagentThreadLeases: jest.fn(async () => []), + }; + const store = new SubagentThreadTaskStore(testMethods, { + ownerDrainTimeoutMs: 60, + ownerFenceGraceMs: 60, + fenceOwnerAdmission, + renewOwnerAdmission, + releaseOwnerAdmission, + }); + + let releaseDeletion = (): void => undefined; + const deletionBlocked = new Promise((resolve) => { + releaseDeletion = resolve; + }); + const fenced = store.withOwnerDeletionFence(userId, undefined, async () => { + await deletionBlocked; + return 'deleted'; + }); + await renewing; + releaseDeletion(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(order).toEqual(['fence']); + + releaseRenewal(); + await expect(fenced).resolves.toBe('deleted'); + /** The lost entry is re-taken before the recovery drain and only released after + * the in-flight renewal and recovery renewal both settle. */ + expect(order).toEqual(['fence', 'renew', 'fence', 'renew', 'release']); + expect(fenceOwnerAdmission).toHaveBeenCalledTimes(2); + }); + + it('cancels each drained task once and retries only unconfirmed deliveries', async () => { + const userId = 'drain-user'; + const parentConversationId = randomUUID(); + const store = new SubagentThreadTaskStore(methods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 5_000, + }); + const lease = { taskId: 'task-1', parentConversationId, conversationId: randomUUID() }; + const listLeases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockResolvedValueOnce([lease]) + .mockResolvedValueOnce([lease]) + .mockResolvedValueOnce([lease]) + .mockResolvedValueOnce([lease]) + .mockResolvedValue([]); + const controlTask = jest + .spyOn(store, 'controlTask') + .mockRejectedValueOnce(new Error('owner unavailable')) + .mockResolvedValueOnce({ status: 'not_found' }) + .mockResolvedValue({ status: 'cancelled', task: threadSnapshot('task-1') }); + try { + await store.cancelAndDrainForOwner(userId); + + /** An unconfirmed delivery is retried under the same invocation — including a + * `not_found`, which means the owner's registration is missing while its lease + * is live — and once the owner confirms, the drain only waits for the lease. */ + expect(controlTask).toHaveBeenCalledTimes(3); + expect(new Set(controlTask.mock.calls.map((call) => call[3])).size).toBe(1); + expect(listLeases).toHaveBeenCalledTimes(5); + } finally { + listLeases.mockRestore(); + controlTask.mockRestore(); + } + }); + + it('fences owner admission around the deletion it drains for', async () => { + const userId = 'fenced-user'; + const order: string[] = []; + const tokens: string[] = []; + const renewed: string[] = []; + const released: string[] = []; + const store = new SubagentThreadTaskStore(methods, { + ownerDrainPollMs: 1, + fenceOwnerAdmission: async (_userId: string, token: string) => { + tokens.push(token); + order.push('fence'); + }, + renewOwnerAdmission: async (_userId: string, token: string) => { + renewed.push(token); + return true; + }, + releaseOwnerAdmission: async (_userId: string, token: string) => { + released.push(token); + order.push('release'); + }, + }); + const listLeases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockImplementation(async () => { + order.push('drain'); + return []; + }); + try { + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + order.push('delete'); + return 'deleted'; + }), + ).resolves.toBe('deleted'); + expect(order).toEqual(['fence', 'drain', 'delete', 'release']); + /** Only the fence this deletion took is lifted, so an overlapping deletion + * keeps admission closed until its own fence is released. */ + expect(released).toEqual(tokens); + expect(tokens[0]).toEqual(expect.any(String)); + + /** A failed deletion still lifts the fence, so one bad request cannot leave the + * account unable to run subagents. */ + order.length = 0; + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + throw new Error('deletion failed'); + }), + ).rejects.toThrow('deletion failed'); + expect(order).toEqual(['fence', 'drain', 'release']); + } finally { + listLeases.mockRestore(); + } + }); + + it('routes conversation-deletion cancellation to a remote task owner', async () => { + const userId = 'routed-delete-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + + const plan = await deletingStore.planCancellationForConversations(userId, [ + parentConversationId, + ]); + await expect(deletingStore.cancelPlan(plan)).resolves.toBe(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('routes cancellation for a deleted child thread to its remote owner', async () => { + const userId = 'routed-child-delete-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + for (let attempt = 0; attempt < 200; attempt += 1) { + if ((await methods.getConvo(userId, threadId)) != null) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + expect(await methods.getConvo(userId, threadId)).not.toBeNull(); + + /** The parent survives this deletion, so the child's own thread is the only target. */ + const plan = await deletingStore.planCancellationForConversations(userId, [threadId]); + await expect(deletingStore.cancelPlan(plan)).resolves.toBe(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('cancels a child admitted after the deletion snapshot from its durable lease', async () => { + const userId = 'lease-cancel-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + for (let attempt = 0; attempt < 200; attempt += 1) { + const leases = await methods.listActiveSubagentThreadLeases({ + user: userId, + now: new Date(), + }); + if (leases.length > 0) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + + /** Production ordering: the plan is resolved first, the cascade is deleted, and + * only then is the plan replayed against the owner directory. */ + const plan = await deletingStore.planCancellationForConversations(userId, [ + parentConversationId, + ]); + await methods.deleteConvos(userId, { conversationId: parentConversationId }); + await expect( + deletingStore.cancelPlan(plan, [parentConversationId, requireThreadId(started)]), + ).resolves.toBeGreaterThanOrEqual(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('drains only active lease addresses when deleting every conversation across replicas', async () => { + const userId = 'routed-owner-drain-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods, { ownerDrainPollMs: 5 }); + const deletingStore = new SubagentThreadTaskStore(methods, { ownerDrainPollMs: 5 }); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => { + markEntered(); + return new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }); + }, + }), + ); + await entered; + + await deletingStore.cancelAndDrainForOwner(userId); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'cancelled', + }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + it('bounds durable delegation depth to one by default', async () => { const userId = 'depth-user'; const rootConversationId = randomUUID(); diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index dc0f0477068..415ef9533f8 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -8,24 +8,36 @@ import { } from '@librechat/agents/langchain/messages'; import type { InMemorySubagentTaskStoreOptions, + SubagentTaskClaim, SubagentTaskConfig, SubagentTaskControlCommand, SubagentTaskControlResult, SubagentTaskRuntime, + SubagentTaskSnapshot, SubagentTaskStartRequest, SubagentTaskStartResult, } from '@librechat/agents'; import type { AllMethods, + IActiveSubagentThreadLease, IConversation, IMessage, MessageMethods, ConversationMethods, + SubagentTaskResultClaim, } from '@librechat/data-schemas'; import type { BaseMessage, StoredMessage } from '@librechat/agents/langchain/messages'; +import type { SubagentTaskControlTransport } from './subagentTaskRouting'; import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; +import { + boundedClaim, + boundedTaskList, + controlFingerprint, + SubagentTaskOwnerUnavailableError, +} from './subagentTaskRouting'; import { createSubagentAttemptKey, createSubagentThreadId } from './subagentThreadIds'; import { runWithDetachedSubagentUsage } from './subagentTaskContext'; +import { createConcurrencyLimiter } from '~/utils/promise'; import { aggregateEmittedUsage } from './usage'; const SCOPE_VERSION = 1; @@ -33,10 +45,29 @@ const DEFAULT_MAX_THREAD_DEPTH = 1; const DEFAULT_LEASE_TTL_MS = 30_000; const DEFAULT_LEASE_HEARTBEAT_MS = 10_000; const DEFAULT_OWNER_DRAIN_TIMEOUT_MS = 45_000; +/** Keeps the admission fence alive across the deletion that follows the drain. */ +const OWNER_FENCE_GRACE_MS = 5 * 60_000; const DEFAULT_OWNER_DRAIN_POLL_MS = 100; +/** Matches the deletion drain batch so cancellation cannot burst Redis. */ +const DELETION_CANCEL_CONCURRENCY = 32; +/** Bounds retained control invocations; one entry per applied command. */ +const MAX_CONTROL_INVOCATIONS = 4_096; + +/** A cancellation target set resolved before the conversations are removed. */ +export interface SubagentCancellationPlan { + userId: string; + tenantId?: string; + conversationIds: string[]; + scopes: Array<{ scopeId: string; threadIds: string[] | null }>; + leases: IActiveSubagentThreadLease[]; +} +/** Three missed 10-second transport heartbeats retire a crashed owner. */ +const DEFAULT_TASK_ROUTING_TTL_MS = 30_000; const MAX_TRANSCRIPT_BYTES = 12 * 1024 * 1024; const TRANSCRIPT_SELECT = 'messageId parentMessageId text createdAt +subagentTranscript +subagentTask'; +const DURABLE_RESULT_SELECT = + 'messageId conversationId sender text createdAt updatedAt +subagentTask'; class SubagentThreadPublicError extends Error {} class SubagentThreadDeletedError extends SubagentThreadPublicError {} @@ -44,11 +75,13 @@ class SubagentThreadDeletedError extends SubagentThreadPublicError {} type SubagentThreadMethods = Pick< AllMethods, | 'acquireSubagentThreadLease' + | 'claimSubagentTaskResult' | 'countActiveSubagentThreadLeases' | 'deleteConvos' | 'deleteMessages' | 'getConvo' | 'getMessages' + | 'listActiveSubagentThreadLeases' | 'reserveSubagentThread' | 'releaseSubagentThreadLease' | 'renewSubagentThreadLease' @@ -88,6 +121,8 @@ interface TaskThreadLease { shared?: { token: string; lost: boolean; + /** Epoch ms this lease is durable until, advanced only by a confirmed renewal. */ + expiresAt: number; heartbeat?: ReturnType; heartbeatInFlight?: Promise; }; @@ -99,7 +134,13 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor leaseHeartbeatMs?: number; ownerDrainTimeoutMs?: number; ownerDrainPollMs?: number; + taskRoutingTtlMs?: number; isOwnerActive?: (userId: string) => Promise; + maxControlInvocations?: number; + ownerFenceGraceMs?: number; + fenceOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; + renewOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; + releaseOwnerAdmission?: (userId: string, token: string) => Promise; } function positiveInteger(value: number | undefined, fallback: number): number { @@ -142,6 +183,10 @@ function parseScope(scopeId: string): SubagentThreadScope { }; } +function serializeScope(scope: Omit): string { + return JSON.stringify({ version: SCOPE_VERSION, ...scope }); +} + function matchesTenant(actual: string | undefined, expected: string | undefined): boolean { return actual === expected; } @@ -284,21 +329,65 @@ function publicFailureDetail(error: unknown): string { : 'The child run could not be completed.'; } +/** Rebuilds the terminal claim a recovered durable result stands for. */ +function recoveredClaim( + message: IMessage, + claim: Extract, +): SubagentTaskClaim | undefined { + const status = message.subagentTask?.status; + const content = message.text ?? ''; + /** A durable child message keeps the untruncated output, so recovering one applies + * the same bounds a routed response would have. */ + if (status === 'completed') { + return boundedClaim({ status: 'completed', task: claim.task, result: content }); + } + if (status === 'error' || status === 'cancelled') { + return boundedClaim({ status, task: claim.task, error: content }); + } + return undefined; +} + +function drainKey(parentConversationId: string, taskId: string): string { + return `${parentConversationId}\u0000${taskId}`; +} + function safeErrorMessage(error: unknown): string { return `Subagent task failed: ${publicFailureDetail(error).slice(0, 2_000)}`; } -/** Persists view-only logical child threads with process-local controls and a shared execution fence. */ +/** Persists view-only logical child threads with owner-routed controls and a shared execution fence. */ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { readonly supportsThreadContinuation = true; private readonly activeThreads = new Map(); + private readonly controlInvocations = new Map< + string, + { scopeId: string; taskId: string; fingerprint: string; result: SubagentTaskControlResult } + >(); + private readonly parentPersistence = new Map>(); private readonly maxThreadDepth: number; private readonly leaseTtlMs: number; private readonly leaseHeartbeatMs: number; private readonly ownerDrainTimeoutMs: number; private readonly ownerDrainPollMs: number; + private readonly taskRoutingTtlMs: number; + private readonly maxControlInvocations: number; + private readonly ownerFenceGraceMs: number; private readonly isOwnerActive: (userId: string) => Promise; + private readonly fenceOwnerAdmission?: ( + userId: string, + token: string, + fencedUntil: Date, + ) => Promise; + + private readonly renewOwnerAdmission?: ( + userId: string, + token: string, + fencedUntil: Date, + ) => Promise; + + private readonly releaseOwnerAdmission?: (userId: string, token: string) => Promise; + private taskControlTransport?: SubagentTaskControlTransport; constructor( private readonly methods: SubagentThreadMethods, @@ -319,7 +408,37 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { DEFAULT_OWNER_DRAIN_TIMEOUT_MS, ); this.ownerDrainPollMs = positiveInteger(options.ownerDrainPollMs, DEFAULT_OWNER_DRAIN_POLL_MS); + this.taskRoutingTtlMs = positiveInteger(options.taskRoutingTtlMs, DEFAULT_TASK_ROUTING_TTL_MS); + this.maxControlInvocations = positiveInteger( + options.maxControlInvocations, + MAX_CONTROL_INVOCATIONS, + ); + this.ownerFenceGraceMs = positiveInteger(options.ownerFenceGraceMs, OWNER_FENCE_GRACE_MS); this.isOwnerActive = options.isOwnerActive ?? (async () => true); + this.fenceOwnerAdmission = options.fenceOwnerAdmission; + this.renewOwnerAdmission = options.renewOwnerAdmission; + this.releaseOwnerAdmission = options.releaseOwnerAdmission; + } + + /** Enables optional cross-replica lookup after the host's Redis service is ready. */ + async configureTaskControlTransport(transport: SubagentTaskControlTransport): Promise { + if (this.taskControlTransport != null) { + throw new Error('Subagent task control transport is already configured.'); + } + await transport.bind({ + claim: (scopeId, taskId) => super.claim(scopeId, taskId), + control: (scopeId, taskId, command, invocationId) => + this.controlInvocation(scopeId, taskId, command, invocationId), + list: (scopeId) => super.list(scopeId), + cancelScope: (scopeId, threadIds) => this.cancelForScope(scopeId, threadIds), + }); + this.taskControlTransport = transport; + } + + async destroyTaskControlTransport(): Promise { + const transport = this.taskControlTransport; + this.taskControlTransport = undefined; + await transport?.destroy(); } /** Gates child creation on the ordinary parent write without retaining request state. */ @@ -381,6 +500,15 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { if (runtime.signal.aborted) { throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); } + /** Publish the owner address before any provider work: a child running + * while unaddressable cannot be polled, controlled, or cancelled, and its + * side effects would already have happened by the time a heartbeat + * republished it. A failed registration fails the task closed instead. */ + await this.taskControlTransport?.registerTask( + request.scopeId, + runtime.taskId, + this.taskRoutingTtlMs, + ); await parentReady; const prepared = await this.prepareThread( request.scopeId, @@ -490,6 +618,285 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { return started; } + /** + * Claims locally when possible, otherwise asks the registered owning replica. + * + * A child's terminal result is durable in its own thread, so collection is recorded + * there against the polling invocation rather than kept alive in the owner's memory. + * The invocation that lost a response re-acquires its own result on the next poll; + * a different invocation is told the result was already collected. Owner-side + * retention stays a fast path, free to expire, instead of the only copy. + */ + async claimTask( + scopeId: string, + taskId: string, + invocationId?: string, + ): Promise { + const local = super.claim(scopeId, taskId); + const claim = + local.status !== 'not_found' + ? local + : ((await this.taskControlTransport?.claim(scopeId, taskId)) ?? local); + if (invocationId == null || claim.status === 'running') { + return claim; + } + if (claim.status === 'not_found') { + return this.claimDurableTaskResult(scopeId, taskId, invocationId); + } + const threadId = claim.task.threadId; + if (threadId == null || threadId === '') { + return claim; + } + /** The durable record decides who holds this one-shot result. The invocation that + * already consumed it re-acquires and is handed it again, a second invocation is + * told it was collected instead of being given a duplicate, and a task with no + * durable record to arbitrate keeps whatever the owner just answered. */ + const collected = await this.assignResultClaim( + parseScope(scopeId).userId, + threadId, + claim.task.taskId, + invocationId, + ); + if (collected.status === 'claimed') { + return { status: 'claimed', task: claim.task }; + } + if (collected.status === 'not_found') { + return claim; + } + return claim.status === 'claimed' ? (recoveredClaim(collected.message, claim) ?? claim) : claim; + } + + /** + * Recovers a terminal task after its owning process and Redis registration are gone. + * The task id locates only a candidate; durable child lineage re-establishes the + * trusted parent scope before the one-shot result is claimed. + */ + private async claimDurableTaskResult( + scopeId: string, + taskId: string, + invocationId: string, + ): Promise { + const scope = parseScope(scopeId); + let message: IMessage | undefined; + try { + [message] = await this.methods.getMessages( + { + user: scope.userId, + messageId: `${taskId}:assistant`, + 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + }, + DURABLE_RESULT_SELECT, + { limit: 1, sort: false }, + ); + } catch (error) { + logger.warn('[subagentThreads] Failed to locate a durable child result', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const threadId = message?.conversationId; + const status = message?.subagentTask?.status; + if ( + message == null || + !isNonEmptyString(threadId) || + !isNonEmptyString(message.sender) || + (status !== 'completed' && status !== 'error' && status !== 'cancelled') + ) { + return { status: 'not_found' }; + } + + let parent: IConversation | null; + let conversation: IConversation | null; + try { + [parent, conversation] = await Promise.all([ + this.methods.getConvo(scope.userId, scope.parentConversationId), + this.methods.getConvo(scope.userId, threadId), + ]); + } catch (error) { + logger.warn('[subagentThreads] Failed to verify durable child lineage', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const lineage = conversation?.subagentThread; + if ( + parent == null || + conversation == null || + lineage == null || + conversation.endpoint !== EModelEndpoint.agents || + lineage.parentConversationId !== scope.parentConversationId || + lineage.subagentType !== message.sender || + lineage.depth > this.maxThreadDepth || + !matchesTenant(parent.tenantId, scope.tenantId) || + !matchesTenant(conversation.tenantId, scope.tenantId) + ) { + return { status: 'not_found' }; + } + + const createdAt = message.createdAt?.getTime(); + const updatedAt = message.updatedAt?.getTime() ?? createdAt; + if (createdAt == null || updatedAt == null) { + return { status: 'not_found' }; + } + const task: SubagentTaskSnapshot = { + taskId, + threadId, + subagentType: lineage.subagentType, + status, + createdAt, + updatedAt, + resultAvailable: true, + resultClaimed: true, + pendingControls: 0, + ...(status === 'completed' ? {} : { error: message.text ?? '' }), + }; + const collected = await this.assignResultClaim(scope.userId, threadId, taskId, invocationId); + if (collected.status === 'not_found') { + return { status: 'not_found' }; + } + if (collected.status === 'claimed') { + return { status: 'claimed', task }; + } + return ( + recoveredClaim(collected.message, { status: 'claimed', task }) ?? { + status: 'not_found', + } + ); + } + + /** + * Assigns one durable terminal result to the invocation collecting it. A failed + * write is not an absent record: handing the result over without recording its + * claimant would let another invocation acquire the same one-shot output once the + * database recovers, so this reports the retryable path and leaves the result + * unclaimed for a later poll. + */ + private async assignResultClaim( + userId: string, + threadId: string, + taskId: string, + invocationId: string, + ): Promise { + try { + return await this.methods.claimSubagentTaskResult({ + userId, + conversationId: threadId, + taskId, + claimId: invocationId, + }); + } catch (error) { + logger.warn('[subagentThreads] Failed to record a collected child result', error); + throw new SubagentTaskOwnerUnavailableError(); + } + } + + /** + * Controls locally when possible, otherwise asks the registered owning replica. + * `invocationId` identifies one caller invocation: a routed retransmission of that + * invocation replays the owner's result, while a fresh invocation applies again even + * when its action and message are identical. + */ + async controlTask( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string = randomUUID(), + ): Promise { + const local = this.controlInvocation(scopeId, taskId, command, invocationId); + if (local.status !== 'not_found') { + return local; + } + return ( + (await this.taskControlTransport?.control(scopeId, taskId, command, invocationId)) ?? local + ); + } + + /** + * Applies one logical control exactly once for its owning task. Idempotency lives + * here rather than in the transport so a local and a routed caller of the same + * invocation agree, and it is keyed by task as well as invocation because provider + * tool-call ids repeat across runs and agents. + */ + controlInvocation( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): SubagentTaskControlResult { + const key = `${scopeId}\u0000${taskId}\u0000${invocationId}`; + const fingerprint = controlFingerprint(command); + const applied = this.controlInvocations.get(key); + if (applied != null) { + /** One invocation is one command; reusing its id for different content is a + * caller error rather than a retry, so it is refused instead of applied. */ + return applied.fingerprint === fingerprint + ? applied.result + : { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + } + if (this.get(scopeId, taskId) == null) { + /** Not this replica's task. Refusing here would keep the command from ever + * reaching its owner, so local load cannot veto a remote cancellation: the + * owner applies its own window to the routed request. */ + return this.control(scopeId, taskId, command); + } + if (!this.makeRoomForInvocation()) { + /** Every tracked invocation belongs to a task this store still holds. Applying + * this command without room to record it would let a caller retry apply it a + * second time, so it is refused before the child is touched at all. */ + logger.warn('[subagentThreads] Refused a control; live invocation records are full'); + return { + status: 'invalid', + message: 'Too many control invocations are in flight for this process; retry shortly.', + }; + } + const result = this.control(scopeId, taskId, command); + if (result.status === 'not_found') { + return result; + } + this.controlInvocations.set(key, { scopeId, taskId, fingerprint, result }); + return result; + } + + /** + * Frees invocation slots by dropping records whose task the store no longer holds: + * a settled task cannot be controlled again, so its record is worthless, while a + * live one is exactly what a caller retry needs to replay instead of applying its + * command twice. The sweep runs only when the window is full and clears every dead + * record at once, so it is amortized rather than repeated per control. + */ + private makeRoomForInvocation(): boolean { + if (this.controlInvocations.size < this.maxControlInvocations) { + return true; + } + for (const [key, invocation] of this.controlInvocations) { + if (this.get(invocation.scopeId, invocation.taskId) == null) { + this.controlInvocations.delete(key); + } + } + return this.controlInvocations.size < this.maxControlInvocations; + } + + /** Returns this process's tasks plus tasks reported by registered remote owners. */ + async listTasks(scopeId: string): Promise { + const local = super.list(scopeId); + const remote = (await this.taskControlTransport?.list(scopeId)) ?? []; + const byId = new Map(local.map((task) => [task.taskId, task])); + for (const task of remote) { + byId.set(task.taskId, task); + } + /** The remote aggregation and each owner's reply carry their own bound, but this + * merge is what the poll tool reads: without a cap here the list the model sees is + * that bound plus however many children this replica happens to own. */ + return boundedTaskList([...byId.values()]); + } + + /** Fast capability probe used while deciding whether a later turn needs the poll tool. */ + async hasTasks(scopeId: string): Promise { + if (super.list(scopeId).length > 0) { + return true; + } + return (await this.taskControlTransport?.hasTasks(scopeId)) ?? false; + } + override control( scopeId: string, taskId: string, @@ -544,6 +951,148 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ); } + /** + * Resolves every cancellation target while the conversations still exist. The plan is + * replayed after deletion, when those rows can no longer be read back, so the second + * pass only has to reach registered owners through Redis. + */ + async planCancellationForConversations( + userId: string, + conversationIds: Iterable, + tenantId?: string, + ): Promise { + const targetIds = [...new Set(conversationIds)]; + const plan: SubagentCancellationPlan = { + userId, + ...(tenantId == null ? {} : { tenantId }), + conversationIds: targetIds, + scopes: [], + leases: [], + }; + if (targetIds.length === 0 || this.taskControlTransport == null) { + return plan; + } + const targets = new Set(targetIds); + const scopeIdFor = (parentConversationId: string): string => + serializeScope({ + userId, + parentConversationId, + ...(tenantId ? { tenantId } : {}), + }); + /** Deleting a conversation takes its whole scope; a deleted child only cancels its + * own thread inside a parent scope that survives. */ + const conversations = await Promise.all( + targetIds.map((conversationId) => this.methods.getConvo(userId, conversationId)), + ); + const threadTargetsByParent = new Map>(); + for (const [index, conversation] of conversations.entries()) { + const parentConversationId = conversation?.subagentThread?.parentConversationId; + if ( + parentConversationId == null || + targets.has(parentConversationId) || + !matchesTenant(conversation?.tenantId, tenantId) + ) { + continue; + } + const threadIds = threadTargetsByParent.get(parentConversationId) ?? new Set(); + threadIds.add(targetIds[index]); + threadTargetsByParent.set(parentConversationId, threadIds); + } + plan.scopes = [ + ...targetIds.map((parentConversationId) => ({ + scopeId: scopeIdFor(parentConversationId), + threadIds: null, + })), + ...[...threadTargetsByParent].map(([parentConversationId, threadIds]) => ({ + scopeId: scopeIdFor(parentConversationId), + threadIds: [...threadIds], + })), + ]; + /** Captured now so descendants removed by the cascade stay reachable afterwards. */ + plan.leases = await this.methods.listActiveSubagentThreadLeases({ + user: userId, + now: new Date(), + ...(tenantId == null ? {} : { tenantId }), + }); + return plan; + } + + /** + * Cancels local children and replays a plan against registered remote owners. + * `removedConversationIds` extends it with the cascade a deletion reported, matched + * against leases captured before those rows were removed. + */ + async cancelPlan( + plan: SubagentCancellationPlan, + removedConversationIds: Iterable = [], + ): Promise { + const { userId, tenantId } = plan; + const planned = new Set(plan.conversationIds); + const removed = new Set(removedConversationIds); + /** A cascade can remove descendants the plan never named — a grandchild lives in + * its own parent's scope, not the deleted root's — so every removed conversation + * is cancelled as a scope of its own. */ + const targets = [...new Set([...planned, ...removed])]; + let cancelled = this.cancelForConversations(userId, targets, tenantId); + const transport = this.taskControlTransport; + if (transport == null) { + return cancelled; + } + const cancelSlot = createConcurrencyLimiter(DELETION_CANCEL_CONCURRENCY); + const cascadeScopes = [...removed] + .filter((conversationId) => !planned.has(conversationId)) + .map((parentConversationId) => ({ + scopeId: serializeScope({ + userId, + parentConversationId, + ...(tenantId ? { tenantId } : {}), + }), + threadIds: null, + })); + const scopeCancellations = [...plan.scopes, ...cascadeScopes].map((scope) => + cancelSlot(() => transport.cancelScope(scope.scopeId, scope.threadIds)), + ); + const leaseCancellations = plan.leases + .filter( + (lease) => removed.has(lease.parentConversationId) || removed.has(lease.conversationId), + ) + .map((lease) => + cancelSlot(() => + this.controlTask( + serializeScope({ + userId, + parentConversationId: lease.parentConversationId, + ...(tenantId ? { tenantId } : {}), + }), + lease.taskId, + { action: 'cancel' }, + ), + ), + ); + for (const count of await Promise.all(scopeCancellations)) { + cancelled += count; + } + for (const result of await Promise.all(leaseCancellations)) { + if (result.status === 'cancelled') { + cancelled += 1; + } + } + return cancelled; + } + + /** Cancels this process's live children for one scope, optionally narrowed to threads. */ + private cancelForScope(scopeId: string, threadIds: string[] | null): number { + const scope = parseScope(scopeId); + const targets = threadIds == null ? null : new Set(threadIds); + return this.cancelMatchingThreads( + (candidate, threadId) => + candidate.userId === scope.userId && + candidate.parentConversationId === scope.parentConversationId && + matchesTenant(candidate.tenantId, scope.tenantId) && + (targets == null || targets.has(threadId)), + ); + } + /** Cancels every active child owned by a user before a delete-all operation. */ cancelForOwner(userId: string, tenantId?: string): number { return this.cancelMatchingThreads( @@ -551,26 +1100,202 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ); } - /** Cancels local work and waits for every replica's durable lease to drain. */ + /** + * Deletes an owner's conversations behind a durable admission fence. Draining alone + * cannot close the race: a child admitted on another replica after the drain read + * its leases would begin provider work against a parent that is about to disappear. + * Fencing first inverts that — the fence is written before any lease is read, and a + * child validates the fence after its own lease is written, so one of the two always + * observes the other. The fence expires by itself, so a process lost mid-deletion + * cannot leave the account unable to run subagents. + */ + async withOwnerDeletionFence( + userId: string, + tenantId: string | undefined, + deletion: () => Promise, + ): Promise { + const fenceWindowMs = this.ownerDrainTimeoutMs + this.ownerFenceGraceMs; + const token = randomUUID(); + /** Only a confirmed write moves this, so a run of failed renewals leaves it in the + * past and the deletion can tell that its fence is no longer guaranteed. */ + let fencedUntil = Date.now() + fenceWindowMs; + let fenceLapsed = false; + await this.fenceOwnerAdmission?.(userId, token, new Date(fencedUntil)); + /** A very large account, or a stalled database, can outlast one fence window, and + * a fence that expires mid-deletion lets another replica admit a child against + * conversations being deleted. It is renewed for as long as the work runs. */ + let releasing = false; + let inFlight: Promise | undefined; + const renewal = setInterval( + () => { + if (inFlight != null) { + return; + } + inFlight = (async () => { + const deadline = fencedUntil; + const renewedUntil = Date.now() + fenceWindowMs; + const held = await this.renewOwnerAdmission?.(userId, token, new Date(renewedUntil)); + if (held === false) { + /** The durable entry was absent, so admission may already have opened even + * when the local deadline has not passed. Reacquire for containment, but + * retain the lapse so the enclosing deletion re-drains before success. */ + fenceLapsed = true; + if (releasing) { + return; + } + /** The entry is gone — expired, or pruned by another deletion — so this + * deletion takes its fence again rather than running on unfenced. */ + await this.fenceOwnerAdmission?.(userId, token, new Date(renewedUntil)); + } + if (Date.now() >= deadline) { + /** The write only landed after the deadline it was meant to extend, so + * admission stood open in between and a child could have taken a lease the + * drain had already read past. A fence cannot be restored backwards over + * that gap, so the lapse is recorded rather than papered over. */ + fenceLapsed = true; + return; + } + fencedUntil = renewedUntil; + })() + .catch((error) => { + logger.warn('[subagentThreads] Failed to hold the owner admission fence', error); + }) + .finally(() => { + inFlight = undefined; + }); + }, + Math.max(1, Math.floor(fenceWindowMs / 3)), + ); + renewal.unref?.(); + const stopRenewal = async (): Promise => { + clearInterval(renewal); + await inFlight; + }; + const fenceHeld = (): boolean => + this.fenceOwnerAdmission == null || (!fenceLapsed && Date.now() < fencedUntil); + try { + await this.cancelAndDrainForOwner(userId, tenantId); + /** The drain can outlast the fence window when the database is unreachable, and + * renewals that keep failing leave the account open to admitting a child against + * conversations about to disappear. Nothing has been removed yet, so this fails + * closed and the caller retries once the fence can be held again. */ + if (!fenceHeld()) { + throw new Error('The subagent admission fence expired before this deletion began.'); + } + const deleted = await deletion(); + /** Settle a renewal already in flight before deciding whether deletion crossed a + * gap. Otherwise a late write can report the lapse only after this check and the + * finally block would release the fence without re-draining. */ + await stopRenewal(); + if (!fenceHeld()) { + /** The rows are gone, but the gap can leave a child another replica admitted + * while the fence was down. Re-take the fence and drain that work before this + * operation may report success. */ + logger.error( + '[subagentThreads] Owner deletion outlived its admission fence; draining children admitted in the gap', + ); + const recoveryUntil = Date.now() + fenceWindowMs; + const reheld = await this.renewOwnerAdmission?.(userId, token, new Date(recoveryUntil)); + if (reheld !== true) { + await this.fenceOwnerAdmission?.(userId, token, new Date(recoveryUntil)); + } + if (Date.now() >= recoveryUntil) { + throw new Error('The subagent admission fence expired while it was being restored.'); + } + fencedUntil = recoveryUntil; + fenceLapsed = false; + await this.cancelAndDrainForOwner(userId, tenantId); + if (!fenceHeld()) { + throw new Error('The subagent admission fence expired while recovering this deletion.'); + } + } + return deleted; + } finally { + releasing = true; + await stopRenewal(); + /** `clearInterval` stops only future passes. A renewal still waiting on the + * database would otherwise find its fence released, read that as expiry, and + * write a fresh one that nothing is left to lift. */ + /** Only this deletion's own fence is lifted: an overlapping deletion that took a + * later one keeps admission closed until it finishes. */ + await this.releaseOwnerAdmission?.(userId, token).catch((error) => { + logger.warn('[subagentThreads] Failed to release the owner admission fence', error); + }); + } + } + + /** + * Cancels local work and waits for every replica's durable lease to drain. Each task + * is cancelled under one invocation held for the whole drain and only while its + * owner has not answered: a fresh invocation per poll would retain a replay entry on + * the owner for every pass, and a task already reported cancelled needs no second + * command, only its lease to disappear. + */ async cancelAndDrainForOwner(userId: string, tenantId?: string): Promise { this.cancelForOwner(userId, tenantId); const deadline = Date.now() + this.ownerDrainTimeoutMs; + const invocations = new Map(); + const answered = new Set(); while (true) { - const active = await this.methods.countActiveSubagentThreadLeases({ + const activeLeases = await this.methods.listActiveSubagentThreadLeases({ user: userId, now: new Date(), ...(tenantId == null ? {} : { tenantId }), }); - if (active === 0) { + if (activeLeases.length === 0) { return; } if (Date.now() >= deadline) { throw new Error('Timed out draining detached subagent tasks for account deletion.'); } + const unanswered = activeLeases.filter( + ({ parentConversationId, taskId }) => !answered.has(drainKey(parentConversationId, taskId)), + ); + for (let index = 0; index < unanswered.length; index += DELETION_CANCEL_CONCURRENCY) { + await Promise.all( + unanswered + .slice(index, index + DELETION_CANCEL_CONCURRENCY) + .map(({ parentConversationId, taskId }) => + this.cancelDrainedTask( + { userId, parentConversationId, taskId, tenantId }, + invocations, + answered, + ), + ), + ); + } await new Promise((resolve) => setTimeout(resolve, this.ownerDrainPollMs)); } } + /** Sends one drained task's cancellation, retrying only unconfirmed deliveries. */ + private async cancelDrainedTask( + target: { userId: string; parentConversationId: string; taskId: string; tenantId?: string }, + invocations: Map, + answered: Set, + ): Promise { + const { userId, parentConversationId, taskId, tenantId } = target; + const key = drainKey(parentConversationId, taskId); + const invocationId = invocations.get(key) ?? randomUUID(); + invocations.set(key, invocationId); + const scopeId = serializeScope({ + userId, + parentConversationId, + ...(tenantId == null ? {} : { tenantId }), + }); + try { + const result = await this.controlTask(scopeId, taskId, { action: 'cancel' }, invocationId); + /** Only the owner confirming the task is stopped ends the commands for it. A + * `not_found` means its registration is missing while its lease is still live — + * an unconfirmed delivery, retried once the owner republishes itself. */ + if (result.status === 'cancelled' || result.status === 'not_running') { + answered.add(key); + } + } catch (error) { + logger.warn('[subagentThreads] Retrying an unconfirmed child cancellation', error); + } + } + private startSharedLeaseHeartbeat( scopeId: string, scope: SubagentThreadScope, @@ -645,19 +1370,32 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { return false; } try { + const deadline = shared.expiresAt; const now = new Date(); + const renewedUntil = now.getTime() + this.leaseTtlMs; const renewed = await this.methods.renewSubagentThreadLease({ user: scope.userId, conversationId: threadId, token: shared.token, now, - expiresAt: new Date(now.getTime() + this.leaseTtlMs), + expiresAt: new Date(renewedUntil), ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), }); if (!renewed) { shared.lost = true; + return false; } - return renewed; + if (Date.now() >= deadline) { + /** The renewal filter compares against the `now` captured before the call, so a + * write that only lands after this lease had expired still succeeds and moves + * the row forward. An owner drain reading active leases in that gap saw this + * thread as free, so the executor stops rather than run past a deletion that + * may already have stepped over it. */ + shared.lost = true; + return false; + } + shared.expiresAt = renewedUntil; + return true; } catch (error) { shared.lost = true; logger.warn('[subagentThreads] Lost the shared child-thread lease', error); @@ -779,7 +1517,11 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { 'This child thread is already being continued by another run.', ); } - lease.shared = { token: sharedToken, lost: false }; + lease.shared = { + token: sharedToken, + lost: false, + expiresAt: now.getTime() + this.leaseTtlMs, + }; this.startSharedLeaseHeartbeat(scopeId, scope, threadId, lease); /** Account deletion can fence the owner after the optimistic probe but before * this lease exists. Once the lease is visible, revalidate so deletion either @@ -1226,6 +1968,22 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { } } +const REQUIRED_THREAD_METHODS = [ + 'acquireSubagentThreadLease', + 'claimSubagentTaskResult', + 'countActiveSubagentThreadLeases', + 'deleteConvos', + 'deleteMessages', + 'getConvo', + 'getMessages', + 'listActiveSubagentThreadLeases', + 'releaseSubagentThreadLease', + 'renewSubagentThreadLease', + 'reserveSubagentThread', + 'saveConvo', + 'saveMessage', +] as const; + export function createSubagentThreadTaskStore( methods: Pick< ConversationMethods, @@ -1233,14 +1991,27 @@ export function createSubagentThreadTaskStore( | 'countActiveSubagentThreadLeases' | 'deleteConvos' | 'getConvo' + | 'listActiveSubagentThreadLeases' | 'releaseSubagentThreadLease' | 'reserveSubagentThread' | 'renewSubagentThreadLease' | 'saveConvo' > & - Pick, + Pick< + MessageMethods, + 'claimSubagentTaskResult' | 'deleteMessages' | 'getMessages' | 'saveMessage' + >, options?: SubagentThreadTaskStoreOptions, ): SubagentThreadTaskStore { + /** The host wires this from JavaScript, where the parameter type checks nothing. A + * method missing there would otherwise surface as a routed failure at claim time, + * long after startup, so the omission is caught here instead. */ + const missing = REQUIRED_THREAD_METHODS.filter( + (name) => typeof (methods as Record)[name] !== 'function', + ); + if (missing.length > 0) { + throw new Error(`Subagent thread task store is missing methods: ${missing.join(', ')}`); + } return new SubagentThreadTaskStore(methods, options); } @@ -1250,6 +2021,6 @@ export function buildSubagentThreadTaskConfig( ): SubagentTaskConfig { return { store, - scopeId: JSON.stringify({ version: SCOPE_VERSION, ...scope }), + scopeId: serializeScope(scope), }; } diff --git a/packages/api/src/cache/redisUtils.spec.ts b/packages/api/src/cache/redisUtils.spec.ts new file mode 100644 index 00000000000..08abbd2753a --- /dev/null +++ b/packages/api/src/cache/redisUtils.spec.ts @@ -0,0 +1,81 @@ +import IoRedis from 'ioredis'; +import { duplicateIoRedisClient } from './redisUtils'; + +describe('duplicateIoRedisClient', () => { + it('applies overrides to a single-node duplicate', () => { + const client = new IoRedis({ host: '127.0.0.1', port: 6379, lazyConnect: true }); + const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false }); + try { + expect(duplicate.options.enableOfflineQueue).toBe(false); + expect(client.options.enableOfflineQueue).not.toBe(false); + } finally { + duplicate.disconnect(); + client.disconnect(); + } + }); + + it('applies overrides to a cluster duplicate, whose options come second', () => { + const client = new IoRedis.Cluster([{ host: '127.0.0.1', port: 6379 }], { + lazyConnect: true, + }); + const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false }); + try { + /** `Cluster.duplicate` reads its first argument as startup nodes, so passing the + * overrides positionally silently keeps the original's queueing behaviour. */ + expect(duplicate.options.enableOfflineQueue).toBe(false); + expect(client.options.enableOfflineQueue).not.toBe(false); + } finally { + duplicate.disconnect(); + client.disconnect(); + } + }); + + it('disables the offline queue only after a cluster node is ready', () => { + const client = new IoRedis.Cluster([{ host: '127.0.0.1', port: 6379 }], { + lazyConnect: true, + }); + const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false }); + try { + /** ioredis emits from its private pool and synchronously forwards `+node` from + * `Cluster`; drive that real discovery path so the test cannot pass merely + * because a synthetic event happened to share the public event name. */ + const pool = ( + duplicate as unknown as { + connectionPool: { + findOrCreate(options: { host: string; port: number }): InstanceType; + }; + } + ).connectionPool; + const node = pool.findOrCreate({ host: '127.0.0.1', port: 6380 }); + /** Topology discovery needs the node queue until this connection is ready. */ + expect(node.options.enableOfflineQueue).toBe(true); + node.emit('ready'); + expect(node.options.enableOfflineQueue).toBe(false); + } finally { + duplicate.disconnect(); + client.disconnect(); + } + }); + + it('disables the offline queue immediately on nodes discovered after cluster readiness', () => { + const client = new IoRedis.Cluster([{ host: '127.0.0.1', port: 6379 }], { + lazyConnect: true, + }); + const duplicate = duplicateIoRedisClient(client, { enableOfflineQueue: false }); + try { + duplicate.emit('ready'); + const pool = ( + duplicate as unknown as { + connectionPool: { + findOrCreate(options: { host: string; port: number }): InstanceType; + }; + } + ).connectionPool; + const replacement = pool.findOrCreate({ host: '127.0.0.1', port: 6381 }); + expect(replacement.options.enableOfflineQueue).toBe(false); + } finally { + duplicate.disconnect(); + client.disconnect(); + } + }); +}); diff --git a/packages/api/src/cache/redisUtils.ts b/packages/api/src/cache/redisUtils.ts index de37c8ba5cd..7c824d4ebe8 100644 --- a/packages/api/src/cache/redisUtils.ts +++ b/packages/api/src/cache/redisUtils.ts @@ -1,7 +1,49 @@ -import type { RedisClientType, RedisClusterType } from '@redis/client'; import { logger } from '@librechat/data-schemas'; +import type { ClusterOptions, RedisOptions, Cluster, Redis } from 'ioredis'; +import type { RedisClientType, RedisClusterType } from '@redis/client'; import { cacheConfig } from './cacheConfig'; +/** + * Duplicates an ioredis connection with option overrides. `Cluster.duplicate` reads its + * first argument as an optional startup-node list and its second as the overrides, + * unlike `Redis.duplicate`, so options passed positionally to a cluster are silently + * dropped and the duplicate quietly inherits the original's behaviour. + */ +export function duplicateIoRedisClient( + client: Redis | Cluster, + options: RedisOptions & ClusterOptions = {}, +): Redis | Cluster { + if (client.isCluster) { + const duplicate = (client as Cluster).duplicate([], options); + if (options.enableOfflineQueue !== false) { + return duplicate; + } + let clusterHasBeenReady = duplicate.status === 'ready'; + duplicate.once('ready', () => { + clusterHasBeenReady = true; + }); + /** ioredis deliberately forces `enableOfflineQueue: true` on every Cluster node + * after applying `redisOptions`. It needs that queue while a new node discovers + * topology, so changing it at `+node` prevents the cluster from ever becoming + * ready. Initial nodes switch once connected; nodes discovered after the cluster + * was usable fail fast immediately, including during a slot-owner replacement. */ + const disableNodeOfflineQueue = (node: Redis): void => { + const disable = (): void => { + node.options.enableOfflineQueue = false; + }; + if (node.status === 'ready' || clusterHasBeenReady) { + disable(); + } else { + node.once('ready', disable); + } + }; + duplicate.on('+node', disableNodeOfflineQueue); + duplicate.nodes('all').forEach(disableNodeOfflineQueue); + return duplicate; + } + return (client as Redis).duplicate(options); +} + /** * Efficiently deletes multiple Redis keys with support for both cluster and single-node modes. * diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index 1c69d79d6d3..b87de135c78 100644 --- a/packages/data-schemas/src/methods/conversation.spec.ts +++ b/packages/data-schemas/src/methods/conversation.spec.ts @@ -3338,6 +3338,15 @@ describe('Conversation Operations', () => { const winner = claims[0] ? 'token-a' : 'token-b'; const loser = winner === 'token-a' ? 'token-b' : 'token-a'; expect(await methods.countActiveSubagentThreadLeases({ user: 'lease-user', now })).toBe(1); + await expect( + methods.listActiveSubagentThreadLeases({ user: 'lease-user', now }), + ).resolves.toEqual([ + { + conversationId, + parentConversationId: 'parent', + taskId: `task-${winner}`, + }, + ]); expect(await methods.getConvo('lease-user', conversationId)).not.toHaveProperty( 'subagentThreadLease', ); diff --git a/packages/data-schemas/src/methods/conversation.ts b/packages/data-schemas/src/methods/conversation.ts index 79b9d84ad93..0703e5129c8 100644 --- a/packages/data-schemas/src/methods/conversation.ts +++ b/packages/data-schemas/src/methods/conversation.ts @@ -4,6 +4,7 @@ import type { DeleteResult } from 'mongoose'; import type { AppConfig, IChatProjectDocument, + IActiveSubagentThreadLease, IConversation, ISharedLink, ISubagentThreadReservation, @@ -183,6 +184,11 @@ export interface ConversationMethods { now: Date; tenantId?: string; }): Promise; + listActiveSubagentThreadLeases(input: { + user: string; + now: Date; + tenantId?: string; + }): Promise; getConvoOwnership( user: string, conversationId: string, @@ -389,6 +395,37 @@ export function createConversationMethods( }); } + /** Resolves only live task addresses so account-wide cancellation stays O(active tasks). */ + async function listActiveSubagentThreadLeases(input: { + user: string; + now: Date; + tenantId?: string; + }): Promise { + const Conversation = mongoose.models.Conversation as Model; + const conversations = await Conversation.find({ + user: input.user, + ...subagentLeaseTenantFilter(input.tenantId), + 'subagentThreadLease.expiresAt': { $gt: input.now }, + }) + .select('conversationId subagentThread.parentConversationId +subagentThreadLease') + .lean< + Array> + >(); + return conversations.flatMap((conversation) => { + const { conversationId } = conversation; + const parentConversationId = conversation.subagentThread?.parentConversationId; + const taskId = conversation.subagentThreadLease?.taskId; + return typeof conversationId === 'string' && + conversationId !== '' && + typeof parentConversationId === 'string' && + parentConversationId !== '' && + typeof taskId === 'string' && + taskId !== '' + ? [{ conversationId, parentConversationId, taskId }] + : []; + }); + } + /** * Ownership probe for request validation: resolves only the owning user id * instead of materializing the full conversation document (preset spread + @@ -1525,6 +1562,7 @@ export function createConversationMethods( renewSubagentThreadLease, releaseSubagentThreadLease, countActiveSubagentThreadLeases, + listActiveSubagentThreadLeases, getConvoOwnership, getConvoRetention, getConvoTitle, diff --git a/packages/data-schemas/src/methods/index.ts b/packages/data-schemas/src/methods/index.ts index 772938d9fac..d9f16f327b7 100644 --- a/packages/data-schemas/src/methods/index.ts +++ b/packages/data-schemas/src/methods/index.ts @@ -44,7 +44,12 @@ import { createCategoriesMethods, type CategoriesMethods } from './categories'; import { createPresetMethods, type PresetMethods } from './preset'; /* Tier 2 — Moderate (service deps injected) */ import { createConversationTagMethods, type ConversationTagMethods } from './conversationTag'; -import { createMessageMethods, CLIENT_MESSAGE_SELECT, type MessageMethods } from './message'; +import { + createMessageMethods, + CLIENT_MESSAGE_SELECT, + type MessageMethods, + type SubagentTaskResultClaim, +} from './message'; import { createConversationMethods, type ConversationMethods } from './conversation'; import { createChatProjectMethods, type ChatProjectMethods } from './chatProject'; export type { @@ -368,6 +373,7 @@ export type { PresetMethods, ConversationTagMethods, MessageMethods, + SubagentTaskResultClaim, ConversationMethods, ChatProjectMethods, TxMethods, diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 097fe998104..add99d0fc63 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -28,6 +28,7 @@ let bulkSaveMessages: ReturnType['bulkSaveMessages' let updateMessageText: ReturnType['updateMessageText']; let deleteMessagesSince: ReturnType['deleteMessagesSince']; let recordMessage: ReturnType['recordMessage']; +let claimSubagentTaskResult: ReturnType['claimSubagentTaskResult']; beforeAll(async () => { mongoServer = await MongoMemoryServer.create(); @@ -47,6 +48,7 @@ beforeAll(async () => { updateMessageText = methods.updateMessageText; deleteMessagesSince = methods.deleteMessagesSince; recordMessage = methods.recordMessage; + claimSubagentTaskResult = methods.claimSubagentTaskResult; await mongoose.connect(mongoUri); }); @@ -1549,4 +1551,81 @@ describe('Message Operations', () => { expect(doc?.tenantId).toBeUndefined(); }); }); + describe('claimSubagentTaskResult', () => { + const terminalResult = async (taskId: string, conversationId: string, status: string) => + saveMessage({ userId: 'user123' }, { + messageId: `${taskId}:assistant`, + conversationId, + text: 'child result', + subagentTask: { attemptKey: `${taskId}:attempt`, status }, + } as Partial); + + it('hands one terminal result to a single polling invocation', async () => { + const taskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(taskId, conversationId, 'completed'); + + const first = await claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + claimId: 'poll-1', + }); + expect(first.status).toBe('acquired'); + expect(first.status === 'acquired' && first.message.text).toBe('child result'); + + /** The same invocation retrying recovers the result it never received. */ + const retried = await claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + claimId: 'poll-1', + }); + expect(retried.status).toBe('acquired'); + + /** Another invocation is told it was collected instead of handed a copy. */ + await expect( + claimSubagentTaskResult({ userId: 'user123', conversationId, taskId, claimId: 'poll-2' }), + ).resolves.toEqual({ status: 'claimed' }); + }); + + it('reports a result that is missing or still running as not found', async () => { + const runningTaskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(runningTaskId, conversationId, 'running'); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: runningTaskId, + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: uuidv4(), + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + }); + + it('never hands one owner’s result to another user', async () => { + const taskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(taskId, conversationId, 'completed'); + + await expect( + claimSubagentTaskResult({ + userId: 'other-user', + conversationId, + taskId, + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + }); + }); }); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index a864acc5766..6f427bb8207 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -48,6 +48,11 @@ interface MessageQueryOptions { sort?: Record | false; } +export type SubagentTaskResultClaim = + | { status: 'not_found' } + | { status: 'claimed' } + | { status: 'acquired'; message: IMessage }; + export interface MessageMethods { saveMessage( ctx: { userId: string; isTemporary?: boolean; interfaceConfig?: AppConfig['interfaceConfig'] }, @@ -82,6 +87,12 @@ export interface MessageMethods { message: Partial & { newMessageId?: string }, metadata?: { context?: string }, ): Promise>; + claimSubagentTaskResult(params: { + userId: string; + conversationId: string; + taskId: string; + claimId: string; + }): Promise; deleteMessagesSince( userId: string, params: { messageId: string; conversationId: string }, @@ -518,6 +529,62 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } } + /** + * Assigns one durable terminal child result to the polling invocation that collects + * it. The same invocation may re-acquire, so a poll whose response was lost recovers + * the result it never received; a different invocation is told it was already + * collected rather than handed a second copy. + */ + async function claimSubagentTaskResult({ + userId, + conversationId, + taskId, + claimId, + }: { + userId: string; + conversationId: string; + taskId: string; + claimId: string; + }): Promise { + if ( + taskId.length === 0 || + taskId.length > 256 || + conversationId.length === 0 || + conversationId.length > 256 || + claimId.length === 0 || + claimId.length > 128 + ) { + throw new TypeError('Invalid subagent task result claim'); + } + const Message = mongoose.models.Message as Model; + const filter = { + user: userId, + conversationId, + messageId: `${taskId}:assistant`, + 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + }; + const acquired = await Message.findOneAndUpdate( + { + ...filter, + $or: [ + { 'subagentTask.resultClaim': { $exists: false } }, + { 'subagentTask.resultClaim.claimId': claimId }, + ], + }, + { $set: { 'subagentTask.resultClaim': { claimId, claimedAt: new Date() } } }, + { + new: true, + timestamps: false, + projection: { messageId: 1, conversationId: 1, text: 1, subagentTask: 1 }, + }, + ).lean(); + if (acquired != null) { + return { status: 'acquired', message: acquired }; + } + const existing = await Message.exists(filter); + return existing == null ? { status: 'not_found' } : { status: 'claimed' }; + } + /** * Deletes messages in a conversation since a specific message. */ @@ -655,6 +722,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateMessageText, updateToolCallResult, updateMessage, + claimSubagentTaskResult, deleteMessagesSince, getMessages, getMessage, diff --git a/packages/data-schemas/src/methods/user.methods.spec.ts b/packages/data-schemas/src/methods/user.methods.spec.ts index 9540026232f..d7be36e814c 100644 --- a/packages/data-schemas/src/methods/user.methods.spec.ts +++ b/packages/data-schemas/src/methods/user.methods.spec.ts @@ -707,6 +707,159 @@ describe('User Methods - Database Tests', () => { }); }); + describe('subagent admission fence', () => { + test('closes admission until the deletion that took the fence releases it', async () => { + const user = await User.create({ + name: 'Subagent Fence', + email: 'subagent-fence@example.com', + provider: 'local', + }); + const userId = user._id.toString(); + const fencedUntil = new Date(Date.now() + 60_000); + + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true); + await methods.fenceSubagentAdmission(userId, 'deletion-a', fencedUntil); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false); + + /** Each overlapping deletion holds its own fence, so admission reopens only + * once the last one finishes — in either completion order. */ + await methods.fenceSubagentAdmission(userId, 'deletion-b', fencedUntil); + await methods.releaseSubagentAdmission(userId, 'deletion-a'); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false); + + await methods.releaseSubagentAdmission(userId, 'deletion-b'); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true); + }); + + test('keeps admission closed when the later deletion finishes first', async () => { + const user = await User.create({ + name: 'Reverse Fence', + email: 'reverse-fence@example.com', + provider: 'local', + }); + const userId = user._id.toString(); + const fencedUntil = new Date(Date.now() + 60_000); + + await methods.fenceSubagentAdmission(userId, 'deletion-a', fencedUntil); + await methods.fenceSubagentAdmission(userId, 'deletion-b', fencedUntil); + + /** The deletion that started second finishes first; the first is still running. */ + await methods.releaseSubagentAdmission(userId, 'deletion-b'); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false); + + await methods.releaseSubagentAdmission(userId, 'deletion-a'); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true); + }); + + test('prunes an expired fence when the next deletion takes one', async () => { + const user = await User.create({ + name: 'Pruned Fence', + email: 'pruned-fence@example.com', + provider: 'local', + }); + const userId = user._id.toString(); + + await methods.fenceSubagentAdmission(userId, 'abandoned', new Date(Date.now() - 1)); + await methods.fenceSubagentAdmission(userId, 'deletion-a', new Date(Date.now() + 60_000)); + + const stored = await User.findById(userId).select('+subagentAdmissionFences').lean(); + expect(stored?.subagentAdmissionFences).toHaveLength(1); + expect(stored?.subagentAdmissionFences?.[0]?.token).toBe('deletion-a'); + }); + + test('reopens admission once an abandoned fence expires', async () => { + const user = await User.create({ + name: 'Expired Fence', + email: 'expired-fence@example.com', + provider: 'local', + }); + const userId = user._id.toString(); + + await methods.fenceSubagentAdmission(userId, 'deletion-a', new Date(Date.now() - 1)); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(true); + }); + + test('refuses an excess deletion instead of discarding an active fence', async () => { + const user = await User.create({ + name: 'Saturated Fence', + email: 'saturated-fence@example.com', + provider: 'local', + }); + const userId = user._id.toString(); + const fencedUntil = new Date(Date.now() + 60_000); + + for (let index = 0; index < 32; index += 1) { + await methods.fenceSubagentAdmission(userId, `deletion-${index}`, fencedUntil); + } + await expect( + methods.fenceSubagentAdmission(userId, 'deletion-overflow', fencedUntil), + ).rejects.toThrow('Too many concurrent bulk deletions'); + + /** The first deletion still owns its fence, so admission stays closed for it. */ + const stored = await User.findById(userId).select('+subagentAdmissionFences').lean(); + expect(stored?.subagentAdmissionFences).toHaveLength(32); + expect(stored?.subagentAdmissionFences?.[0]?.token).toBe('deletion-0'); + await expect(methods.isSubagentOwnerAdmissible(userId)).resolves.toBe(false); + }); + + test('invalidates the cached auth document when a refused fence still pruned', async () => { + enableAuthUserDocCache(); + const user = await User.create({ + name: 'Refused Fence', + email: 'refused-fence@example.com', + provider: 'local', + }); + const userId = user._id?.toString() ?? ''; + const indexKey = `${AUTH_USER_DOC_BY_ID_PREFIX}:${userId}`; + const fencedUntil = new Date(Date.now() + 60_000); + /** A saturated owner that has since abandoned one fence: the next attempt prunes + * the expired entry and is then refused by the cap, so the two writes disagree. */ + await User.updateOne( + { _id: userId }, + { + $set: { + subagentAdmissionFences: [ + ...Array.from({ length: 32 }, (_unused, index) => ({ + token: `deletion-${index}`, + expiresAt: fencedUntil, + })), + { token: 'abandoned', expiresAt: new Date(Date.now() - 1) }, + ], + }, + }, + ); + + const cache = { + get: jest.fn().mockResolvedValue(['auth-cache-key-a']), + delete: jest.fn().mockResolvedValue(true), + }; + const methodsWithCache = createUserMethods(mongoose, { + getCache: jest.fn().mockReturnValue(cache), + }); + + await expect( + methodsWithCache.fenceSubagentAdmission(userId, 'deletion-overflow', fencedUntil), + ).rejects.toThrow('Too many concurrent bulk deletions'); + const stored = await User.findById(userId).select('+subagentAdmissionFences').lean(); + expect(stored?.subagentAdmissionFences).toHaveLength(32); + /** The prune committed, so leaving the cached document in place would serve the + * pruned fence until its own TTL expired. */ + expect(cache.delete).toHaveBeenCalledWith('auth-cache-key-a'); + expect(cache.delete).toHaveBeenCalledWith(indexKey); + }); + + test('refuses an unbounded or invalid fence', async () => { + const userId = new mongoose.Types.ObjectId().toString(); + + await expect( + methods.fenceSubagentAdmission(userId, 'deletion-a', new Date(Number.NaN)), + ).rejects.toThrow('fencedUntil must be a valid Date'); + await expect( + methods.fenceSubagentAdmission(userId, '', new Date(Date.now() + 60_000)), + ).rejects.toThrow('bounded owner token'); + }); + }); + describe('countUsers', () => { test('should count all users', async () => { await User.create([ diff --git a/packages/data-schemas/src/methods/user.ts b/packages/data-schemas/src/methods/user.ts index 2d324bfcdc3..2840c458b39 100644 --- a/packages/data-schemas/src/methods/user.ts +++ b/packages/data-schemas/src/methods/user.ts @@ -14,6 +14,8 @@ import { signPayload } from '~/crypto'; export const DEFAULT_SESSION_EXPIRY: number = 1000 * 60 * 15; /** Minimum age before an explicitly offline operator may recover an abandoned deletion fence. */ export const USER_DELETION_FENCE_STALE_MS: number = 15 * 60_000; +/** Bounds concurrent bulk deletions held for one owner at any moment. */ +const MAX_SUBAGENT_ADMISSION_FENCES = 32; interface UserMethodDeps { getCache?: (key: string) => CacheStore | undefined; @@ -129,6 +131,10 @@ export function createUserMethods( ) => Promise<'acquired' | 'in_progress' | 'missing'>; cancelAgentTriggerUserDeletion: (userId: string, startedAt: Date) => Promise; isAgentTriggerPrincipalActive: (userId: string) => Promise; + fenceSubagentAdmission: (userId: string, token: string, fencedUntil: Date) => Promise; + renewSubagentAdmission: (userId: string, token: string, fencedUntil: Date) => Promise; + releaseSubagentAdmission: (userId: string, token: string) => Promise; + isSubagentOwnerAdmissible: (userId: string) => Promise; deleteUserById: (userId: string) => Promise; updateUserPlugins: ( userId: string, @@ -452,6 +458,101 @@ export function createUserMethods( ); } + /** + * Closes subagent admission for one owner while a bulk conversation deletion drains + * its live children. Every concurrent deletion holds its own fence, so admission + * reopens only once the last one finishes, in whatever order they complete. Each + * fence expires on its own, so a process that dies mid-delete cannot lock the + * account out of running subagents, and expired fences are pruned as new ones + * arrive rather than accumulating. + */ + async function fenceSubagentAdmission( + userId: string, + token: string, + fencedUntil: Date, + ): Promise { + if (!(fencedUntil instanceof Date) || !Number.isFinite(fencedUntil.getTime())) { + throw new TypeError('fencedUntil must be a valid Date'); + } + if (token.length === 0 || token.length > 128) { + throw new TypeError('A subagent admission fence needs a bounded owner token'); + } + const User = mongoose.models.User; + /** Plain update operators only: DocumentDB rejects pipeline-form updates, and + * this runs before any deletion, so using one would fail the whole endpoint. */ + await User.updateOne( + { _id: userId }, + { $pull: { subagentAdmissionFences: { expiresAt: { $lte: new Date() } } } }, + { timestamps: false }, + ); + try { + /** Admitted only while the owner is under the concurrent-deletion cap. Dropping + * an active fence to make room would reopen admission for a deletion that is + * still running, so an excess deletion is refused instead. */ + const fenced = await User.updateOne( + { + _id: userId, + [`subagentAdmissionFences.${MAX_SUBAGENT_ADMISSION_FENCES - 1}`]: { $exists: false }, + }, + { $push: { subagentAdmissionFences: { token, expiresAt: fencedUntil } } }, + { timestamps: false }, + ); + if (fenced.matchedCount !== 1) { + throw new Error('Too many concurrent bulk deletions are already fencing this owner.'); + } + } finally { + /** The prune above commits on its own, so a refused or failed fence still leaves + * the cached document describing entries the collection no longer holds. */ + await invalidateAuthUserDocCache(userId); + } + } + + /** Extends only this deletion's own fence while its work is still running. */ + async function renewSubagentAdmission( + userId: string, + token: string, + fencedUntil: Date, + ): Promise { + if (!(fencedUntil instanceof Date) || !Number.isFinite(fencedUntil.getTime())) { + throw new TypeError('fencedUntil must be a valid Date'); + } + const User = mongoose.models.User; + const result = await User.updateOne( + { _id: userId, 'subagentAdmissionFences.token': token }, + { $set: { 'subagentAdmissionFences.$.expiresAt': fencedUntil } }, + { timestamps: false }, + ); + if (result.modifiedCount === 1) { + await invalidateAuthUserDocCache(userId); + } + return result.matchedCount === 1; + } + + /** Lifts only this deletion's fence, so an overlapping one keeps admission closed. */ + async function releaseSubagentAdmission(userId: string, token: string): Promise { + const User = mongoose.models.User; + const result = await User.updateOne( + { _id: userId }, + { $pull: { subagentAdmissionFences: { token } } }, + { timestamps: false }, + ); + if (result.modifiedCount === 1) { + await invalidateAuthUserDocCache(userId); + } + } + + /** True while this owner may admit a new child: no account deletion, no live fence. */ + async function isSubagentOwnerAdmissible(userId: string): Promise { + const User = mongoose.models.User; + return ( + (await User.exists({ + _id: userId, + agentTriggerDeletionStartedAt: { $exists: false }, + subagentAdmissionFences: { $not: { $elemMatch: { expiresAt: { $gt: new Date() } } } }, + })) != null + ); + } + /** * Generates a JWT token for a given user. * @param user - The user object @@ -707,6 +808,10 @@ export function createUserMethods( recoverStaleAgentTriggerUserDeletion, cancelAgentTriggerUserDeletion, isAgentTriggerPrincipalActive, + fenceSubagentAdmission, + renewSubagentAdmission, + releaseSubagentAdmission, + isSubagentOwnerAdmissible, deleteUserById, updateUserPlugins, toggleUserMemories, diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index 890f1fe2cdd..8282cbf0597 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -144,6 +144,14 @@ const messageSchema: Schema = new Schema( enum: ['running', 'completed', 'error', 'cancelled'], required: true, }, + resultClaim: { + type: { + claimId: { type: String, required: true }, + claimedAt: { type: Date, required: true }, + }, + _id: false, + default: undefined, + }, }, _id: false, select: false, diff --git a/packages/data-schemas/src/schema/user.ts b/packages/data-schemas/src/schema/user.ts index 79a9a641957..8d13c1ec79f 100644 --- a/packages/data-schemas/src/schema/user.ts +++ b/packages/data-schemas/src/schema/user.ts @@ -135,6 +135,17 @@ const userSchema: Schema = new Schema( type: Date, select: false, }, + subagentAdmissionFences: { + type: [ + { + token: { type: String, required: true }, + expiresAt: { type: Date, required: true }, + }, + ], + _id: false, + select: false, + default: undefined, + }, personalization: { type: { memories: { diff --git a/packages/data-schemas/src/types/convo.ts b/packages/data-schemas/src/types/convo.ts index 5fa5fbf0a90..387720cf9f1 100644 --- a/packages/data-schemas/src/types/convo.ts +++ b/packages/data-schemas/src/types/convo.ts @@ -7,6 +7,12 @@ export interface ISubagentThreadLease { expiresAt: Date; } +export interface IActiveSubagentThreadLease { + conversationId: string; + parentConversationId: string; + taskId: string; +} + export interface ISubagentThreadReservation { conversation: IConversation; created: boolean; diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index 9b1d4492f41..541a5f1d1f2 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -53,6 +53,11 @@ export interface IMessage extends Document { attemptKey: string; requestFingerprint?: string; status: 'running' | 'completed' | 'error' | 'cancelled'; + /** Records which polling invocation collected this terminal result. */ + resultClaim?: { + claimId: string; + claimedAt: Date; + }; }; contextMeta?: { calibrationRatio?: number; diff --git a/packages/data-schemas/src/types/user.ts b/packages/data-schemas/src/types/user.ts index 169209cf4a8..8370fb63b65 100644 --- a/packages/data-schemas/src/types/user.ts +++ b/packages/data-schemas/src/types/user.ts @@ -55,6 +55,11 @@ export interface IUser extends Document { termsAcceptedAt?: Date | null; /** Internal fence that prevents agent-trigger admission during account deletion. */ agentTriggerDeletionStartedAt?: Date; + /** Expiring fences closing subagent admission while bulk deletions drain. */ + subagentAdmissionFences?: Array<{ + token: string; + expiresAt: Date; + }>; personalization?: { memories?: boolean; statefulCodeEnvironment?: StatefulCodeEnvironment; From d175741010ec7f20ec76162524e3e49ad943e5f5 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 10:21:12 -0400 Subject: [PATCH 07/15] =?UTF-8?q?=F0=9F=AA=A2=20ci:=20Prevent=20Playwright?= =?UTF-8?q?=20Apt=20Lock=20Leakage=20(#14993)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/playwright-mock.yml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/playwright-mock.yml b/.github/workflows/playwright-mock.yml index 6c57b4fe102..5f984eb6bb3 100644 --- a/.github/workflows/playwright-mock.yml +++ b/.github/workflows/playwright-mock.yml @@ -251,12 +251,17 @@ jobs: continue-on-error: true run: timeout -k 10 90 npx playwright install ffmpeg + # This job deliberately skips the optional font install: its bounded + # Playwright apt process can outlive the wrapper on a slow mirror and + # retain the package-manager lock needed by the required Redis install. + # The MCP suite does not enable visual snapshot assertions. + # Redis is a hard requirement for this job, so this step stays fatal. - name: Install Redis runtime dependencies timeout-minutes: 5 run: | - sudo apt-get update - sudo apt-get install -y redis-server redis-tools + sudo apt-get -o DPkg::Lock::Timeout=300 update + sudo apt-get -o DPkg::Lock::Timeout=300 install -y redis-server redis-tools - name: Start standalone Redis and Redis Cluster run: | From 5e3c6807616bcc1e7cb90019875ea91c437852b2 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 12:18:31 -0400 Subject: [PATCH 08/15] =?UTF-8?q?=F0=9F=AA=83=20feat:=20Wake=20Parent=20Ag?= =?UTF-8?q?ents=20on=20Child=20Completion=20(#14975)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: wake parent agents on child completion * wip: harden child completion wakeup lifecycle * fix: close the completion-wakeup static failures Type the durable-claim store fixture, the continue-envelope test helper, and the terminal message's task metadata so the wakeup suites compile against the shapes they actually exercise. Replace `Array.prototype.at`, which the package target library does not provide. Capture the prepared child thread in a non-optional local before the provider callback closes over it, and narrow the trigger envelope itself on `mode === 'continue'` rather than a separately copied mode, so reading the continue target is sound. Lift the parent-message fallback out of a nested ternary into a named resolver. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * test: cover the active-predecessor admission fence The Redis job-creation call gained a thirteenth scalar argument, so the spec helper reconstructed the HSET pairs one slot early and rebuilt an invalid job hash; three creation tests failed on that alone. Give the fence itself direct coverage in both store adapters, which it had none of despite deciding whether an automatic continuation may replace a live parent turn. Each proves a running and a requires_action predecessor are refused with the state a controller needs for a finite 409, that an absent or settled predecessor is admitted, and that an ordinary user turn without the policy still replaces its predecessor. The Redis case also asserts a refused continuation leaves the parent's durable job and chunks untouched. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H1cCMDrTWaRNkmtKjpWELZ * fix: harden completion wakeup rollout and claims * fix: close completion wakeup race windows * test: keep the child store fixture exact * fix: close final subagent wakeup gaps * fix: preserve ambiguous completion claims * fix: release pre-admission wakeup claims * fix: stabilize subagent completion recovery --------- Co-authored-by: Claude --- .env.example | 8 +- CONTEXT.md | 1 + .../__tests__/request.resumeMetadata.spec.js | 114 +++++ api/server/controllers/agents/request.js | 157 +++++- api/server/services/Agents/triggers.js | 10 +- .../Endpoints/agents/subagentThreadStore.js | 12 + packages/api/src/agents/index.ts | 1 + .../agents/subagentCompletionWakeup.spec.ts | 483 ++++++++++++++++++ .../src/agents/subagentCompletionWakeup.ts | 425 +++++++++++++++ .../api/src/agents/subagentTaskRouting.ts | 7 +- .../api/src/agents/subagentThreads.spec.ts | 92 +++- packages/api/src/agents/subagentThreads.ts | 102 +++- packages/api/src/agents/triggers/README.md | 5 +- packages/api/src/agents/triggers/delivery.ts | 2 +- .../api/src/agents/triggers/dispatch.spec.ts | 58 ++- packages/api/src/agents/triggers/dispatch.ts | 16 +- .../api/src/agents/triggers/engine.spec.ts | 35 ++ packages/api/src/agents/triggers/engine.ts | 16 +- .../api/src/agents/triggers/envelope.spec.ts | 24 +- packages/api/src/agents/triggers/envelope.ts | 52 +- packages/api/src/agents/triggers/host.spec.ts | 276 ++++++++++ packages/api/src/agents/triggers/host.ts | 245 +++++++-- packages/api/src/agents/triggers/service.ts | 2 + .../api/src/stream/GenerationJobManager.ts | 10 + .../stream/__tests__/RedisJobStore.spec.ts | 4 +- .../stream/__tests__/predecessorFence.spec.ts | 93 ++++ ...redecessorFence.stream_integration.spec.ts | 96 ++++ .../implementations/InMemoryJobStore.ts | 18 + .../stream/implementations/RedisJobStore.ts | 16 +- .../api/src/stream/interfaces/IJobStore.ts | 1 + .../data-schemas/src/methods/message.spec.ts | 134 ++++- packages/data-schemas/src/methods/message.ts | 131 ++++- packages/data-schemas/src/schema/message.ts | 2 + packages/data-schemas/src/types/message.ts | 4 +- 34 files changed, 2526 insertions(+), 126 deletions(-) create mode 100644 packages/api/src/agents/subagentCompletionWakeup.spec.ts create mode 100644 packages/api/src/agents/subagentCompletionWakeup.ts diff --git a/.env.example b/.env.example index eb00215d0ef..7091d3f9ad0 100644 --- a/.env.example +++ b/.env.example @@ -1211,11 +1211,17 @@ OPENWEATHER_API_KEY= # Agent Trigger Delivery # #===========================# -# Base URL used by trusted in-process event producers to dispatch agent fires and steers. +# Base URL used by trusted in-process event producers to dispatch agent fires, continuations, +# and steers. # Defaults to this process's bound listener. Set only when internal trigger admission must # traverse another trusted HTTP origin, such as a TLS front door. # AGENT_TRIGGERS_SELF_URL=http://127.0.0.1:3080 +# Automatically continue a saved parent agent after a detached subagent settles. +# Rolling-deploy safety: deploy support with this disabled first, wait until every API +# replica is upgraded, then enable it in a subsequent rollout. +# ENABLE_SUBAGENT_COMPLETION_WAKEUPS=false + # Trusted event adapters enqueue through the shared durable trigger service. Mongo-backed # leases make its workers safe across replicas; successful delivery records expire after # 90 days, while dead letters remain available for explicit operator requeue. diff --git a/CONTEXT.md b/CONTEXT.md index 6e411689c07..276b1f8ecab 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -3,4 +3,5 @@ - **Agent run envelope**: the versioned, JSON-safe request contract created after ingress authentication and protocol validation but before agent, provider, tool, or MCP initialization. It carries only the validated protocol payload and the minimum trusted principal identifiers. The execution host rehydrates all runtime state from those identifiers. - **Subagent thread**: a durable, view-only child conversation owned by one parent conversation and subagent identity. A parent agent may continue it by stable `threadId`; each continuation uses a fresh execution lease restored from the canonical child transcript. It is not an ordinary human-writable chat. - **Live subagent task owner**: the one API process holding a detached child execution, its abort controller, and its bounded control queue. Redis may route trusted poll/control envelopes to that owner, but it does not migrate or persist the executor; Mongo persists only the logical child thread and its continuation fence. +- **Subagent completion wakeup**: a durable internal `continue` trigger pre-registered before detached child execution so a process crash cannot lose the wakeup. Delivery defers until the child's terminal transcript is persisted, targets the initiating agent and exact parent response branch, carries task metadata rather than child output, waits for the parent generation to settle, and starts the parent turn that collects the result through the existing task store. - **Theme definition**: a versioned, data-only description of LibreChat semantic colors and shared appearance roles, optionally specialized by light or dark mode. The theme module validates and resolves partial definitions against bundled defaults before adapters apply them. A theme definition does not contain arbitrary CSS, application behavior, or alternate feature layouts. diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index 2ddc02abacd..e6592429772 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -519,6 +519,120 @@ describe('ResumableAgentController resume metadata', () => { ); }); + it('defers a trusted trigger resume while its parent generation is still active', async () => { + const conversationId = 'conversation-123'; + mockGetMessages.mockResolvedValue([{ _id: 'persisted-parent' }]); + mockGenerationJobManager.getJob.mockResolvedValue({ + status: 'running', + metadata: { userId: 'user-123' }, + }); + const initializeClient = jest.fn(); + const req = { + _isAgentTrigger: true, + user: { id: 'user-123' }, + body: { + text: 'Collect the completed child.', + messageId: 'wakeup-user-message', + parentMessageId: 'persisted-response_', + conversationId, + clientRequestId: 'trigger_resume_1', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(res.status).toHaveBeenCalledWith(409); + expect(res.json).toHaveBeenCalledWith(expect.objectContaining({ code: 'PARENT_NOT_READY' })); + expect(mockGenerationJobManager.claimGeneration).not.toHaveBeenCalled(); + expect(mockCheckAndIncrementPendingRequest).not.toHaveBeenCalled(); + expect(initializeClient).not.toHaveBeenCalled(); + }); + + it('labels a trigger parent-state lookup failure as provably pre-admission', async () => { + const conversationId = 'conversation-123'; + mockGetMessages.mockResolvedValue([{ _id: 'persisted-parent' }]); + mockGenerationJobManager.getJob.mockRejectedValue(new Error('redis unavailable')); + const initializeClient = jest.fn(); + const req = { + _isAgentTrigger: true, + user: { id: 'user-123' }, + body: { + text: 'Collect the completed child.', + messageId: 'wakeup-user-message', + parentMessageId: 'persisted-response_', + conversationId, + clientRequestId: 'trigger_resume_1', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(res.set).toHaveBeenCalledWith('Retry-After', '1'); + expect(res.status).toHaveBeenCalledWith(503); + expect(res.json).toHaveBeenCalledWith( + expect.objectContaining({ code: 'PARENT_STATE_UNAVAILABLE' }), + ); + expect(mockGenerationJobManager.claimGeneration).not.toHaveBeenCalled(); + expect(mockCheckAndIncrementPendingRequest).not.toHaveBeenCalled(); + expect(initializeClient).not.toHaveBeenCalled(); + }); + + it('deduplicates the active continuation whose admission response was lost', async () => { + const conversationId = 'conversation-123'; + mockGetMessages.mockResolvedValue([{ _id: 'persisted-parent' }]); + mockGenerationJobManager.getJob.mockResolvedValue({ + createdAt: 1000, + status: 'requires_action', + metadata: { + userId: 'user-123', + idempotencyClientRequestId: 'trigger_resume_1', + }, + }); + mockGenerationJobManager.claimGeneration.mockResolvedValue({ + claimed: false, + existing: { + streamId: conversationId, + conversationId, + claimedAt: 100, + claimToken: 'existing-token', + startedAt: 1000, + }, + }); + const req = { + _isAgentTrigger: true, + user: { id: 'user-123' }, + body: { + text: 'Collect the completed child.', + messageId: 'wakeup-user-message', + parentMessageId: 'persisted-response_', + conversationId, + clientRequestId: 'trigger_resume_1', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.status).not.toHaveBeenCalledWith(409); + expect(res.json).toHaveBeenCalledWith({ + streamId: conversationId, + conversationId, + generationCreatedAt: 1000, + status: 'resumed', + generationProtocolVersion: 1, + }); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + expect(mockCheckAndIncrementPendingRequest).not.toHaveBeenCalled(); + }); + it('creates the job with the in-flight turn before MCP initialization can emit OAuth', async () => { const conversationId = 'conversation-123'; const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index fea97639406..580b545f4cc 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -301,6 +301,7 @@ function rejectPreliminaryParentMessageId(res, generationProtocolVersion) { res, 409, { + code: 'PARENT_NOT_READY', error: 'Cannot submit a follow-up while the selected parent response is still being saved. Please wait and try again.', }, @@ -308,6 +309,18 @@ function rejectPreliminaryParentMessageId(res, generationProtocolVersion) { ); } +function rejectMissingTriggerParentMessageId(res, generationProtocolVersion) { + return sendGenerationJson( + res, + 404, + { + code: 'PARENT_NOT_FOUND', + error: 'The selected parent response is no longer available.', + }, + generationProtocolVersion, + ); +} + /** * Resumable Agent Controller - Generation runs independently of HTTP connection. * Returns streamId immediately, client subscribes separately via SSE. @@ -449,6 +462,9 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit : undefined, }); + const isTriggerContinuation = + req._isAgentTrigger === true && !isNewConvo && parentMessageId !== Constants.NO_PARENT; + if ( await isUnpersistedPreliminaryParent({ userId, @@ -457,6 +473,38 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit getMessages, }) ) { + if (isTriggerContinuation) { + let parentJob; + try { + parentJob = await GenerationJobManager.getJob(conversationId); + } catch (error) { + logger.warn('[ResumableAgentController] Trigger parent lookup failed', error); + res.set('Retry-After', '1'); + startupTelemetry?.end('rejected'); + return sendGenerationJson( + res, + 503, + { code: 'PARENT_STATE_UNAVAILABLE', error: 'Parent generation state is unavailable.' }, + generationProtocolVersion, + ); + } + if ( + parentJob != null && + liveJobBelongsToRequester(parentJob, req.user) && + (parentJob.status === 'running' || + parentJob.status === 'requires_action' || + parentJob.metadata?.terminalPersistencePending === true) && + !( + typeof clientRequestId === 'string' && + parentJob.metadata?.idempotencyClientRequestId === clientRequestId + ) + ) { + startupTelemetry?.end('rejected'); + return rejectPreliminaryParentMessageId(res, generationProtocolVersion); + } + startupTelemetry?.end('rejected'); + return rejectMissingTriggerParentMessageId(res, generationProtocolVersion); + } startupTelemetry?.end('rejected'); return rejectPreliminaryParentMessageId(res, generationProtocolVersion); } @@ -472,6 +520,51 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit const streamId = conversationId; req.body.conversationId = conversationId; + /** A durable continuation trigger appends below a completed parent response. If + * that response belongs to a still-running or paused generation, admitting + * another generation on the same conversation stream would replace it. + * Defer without claiming the continuation idempotency key so the delivery engine + * can retry after the parent reaches a terminal state. */ + if (isTriggerContinuation) { + let parentJob; + try { + parentJob = await GenerationJobManager.getJob(streamId); + } catch (error) { + logger.warn('[ResumableAgentController] Trigger continuation parent lookup failed', error); + res.set('Retry-After', '1'); + startupTelemetry?.end('rejected'); + return sendGenerationJson( + res, + 503, + { + code: 'PARENT_STATE_UNAVAILABLE', + error: 'Parent generation state is temporarily unavailable.', + }, + generationProtocolVersion, + ); + } + if ( + parentJob != null && + liveJobBelongsToRequester(parentJob, req.user) && + (parentJob.status === 'running' || + parentJob.status === 'requires_action' || + parentJob.metadata?.terminalPersistencePending === true) && + !( + typeof clientRequestId === 'string' && + parentJob.metadata?.idempotencyClientRequestId === clientRequestId + ) + ) { + res.set('Retry-After', '1'); + startupTelemetry?.end('rejected'); + return sendGenerationJson( + res, + 409, + { code: 'PARENT_NOT_READY', error: 'The parent generation has not settled yet.' }, + generationProtocolVersion, + ); + } + } + // Idempotency: a lost/reset start-generation response makes the client re-POST the // identical payload, which would otherwise start a second fully-billed generation. // Claim the submission's clientRequestId before creating the job so a retry attaches @@ -881,6 +974,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit ...(recoveredSteerId && { recoveredSteerId }), ...(recoveredSteerPayload && { recoveredSteerPayload }), ...(expectedPredecessorCreatedAt != null && { expectedPredecessorCreatedAt }), + ...(isTriggerContinuation && { rejectActivePredecessor: true }), ...(ownedIdempotencyClaim?.claimToken && { idempotencyClientRequestId: clientRequestId, idempotencyClaimToken: ownedIdempotencyClaim.claimToken, @@ -1908,31 +2002,44 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit if (error?.code === 'GENERATION_PREDECESSOR_MISMATCH') { const currentJob = error.currentJob; const currentStatus = currentJob?.status; - const predecessorVerified = - currentJob != null && - Number.isSafeInteger(currentJob.createdAt) && - currentJob.createdAt >= 0 && - currentJob.verified !== false; - sendGenerationJson( - res, - 409, - { - status: 'predecessor_mismatch', - code: 'GENERATION_PREDECESSOR_MISMATCH', - error: predecessorVerified - ? 'A newer generation became current before this request could start.' - : 'The prior generation could not be verified. Please retry.', - streamId, - conversationId: currentJob?.conversationId ?? conversationId, - generationCreatedAt: currentJob?.createdAt, - predecessorVerified, - active: - typeof currentJob?.active === 'boolean' - ? currentJob.active - : currentStatus === 'running' || currentStatus === 'requires_action', - }, - generationProtocolVersion, - ); + if (isTriggerContinuation && currentJob?.active === true) { + res.set('Retry-After', '1'); + sendGenerationJson( + res, + 409, + { + code: 'PARENT_NOT_READY', + error: 'Another generation became active before the continuation could start.', + }, + generationProtocolVersion, + ); + } else { + const predecessorVerified = + currentJob != null && + Number.isSafeInteger(currentJob.createdAt) && + currentJob.createdAt >= 0 && + currentJob.verified !== false; + sendGenerationJson( + res, + 409, + { + status: 'predecessor_mismatch', + code: 'GENERATION_PREDECESSOR_MISMATCH', + error: predecessorVerified + ? 'A newer generation became current before this request could start.' + : 'The prior generation could not be verified. Please retry.', + streamId, + conversationId: currentJob?.conversationId ?? conversationId, + generationCreatedAt: currentJob?.createdAt, + predecessorVerified, + active: + typeof currentJob?.active === 'boolean' + ? currentJob.active + : currentStatus === 'running' || currentStatus === 'requires_action', + }, + generationProtocolVersion, + ); + } } else if (error?.code === 'RECOVERY_PAYLOAD_MISMATCH') { sendGenerationJson( res, diff --git a/api/server/services/Agents/triggers.js b/api/server/services/Agents/triggers.js index 7beb92f8eeb..2279a0da2ad 100644 --- a/api/server/services/Agents/triggers.js +++ b/api/server/services/Agents/triggers.js @@ -1,9 +1,17 @@ -const { createAgentTriggerService } = require('@librechat/api'); +const { + createAgentTriggerService, + createSubagentCompletionWakeupResolver, + GenerationJobManager, +} = require('@librechat/api'); const methods = require('~/models'); const service = createAgentTriggerService({ methods, isPrincipalActive: methods.isAgentTriggerPrincipalActive, + prepareContinue: createSubagentCompletionWakeupResolver({ + methods, + getGenerationJob: (conversationId) => GenerationJobManager.getJob(conversationId), + }), }); module.exports = { diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.js b/api/server/services/Endpoints/agents/subagentThreadStore.js index f11c6a22670..5f7ff966f6c 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -1,12 +1,20 @@ const { cacheConfig, ioredisClient, + isEnabled, registerShutdownTask, duplicateIoRedisClient, createSubagentThreadTaskStore, + createSubagentCompletionWakeupHandler, RedisSubagentTaskControlTransport, } = require('@librechat/api'); const db = require('~/models'); +const { enqueueAgentTrigger } = require('../../Agents/triggers'); + +/** Keep producers off for the first rollout so older trigger workers cannot + * permanently reject the new `continue` envelope. Enable only after every API + * replica runs a release that understands completion wakeups. */ +const completionWakeupsEnabled = isEnabled(process.env.ENABLE_SUBAGENT_COMPLETION_WAKEUPS); /** Durable logical threads use normal LibreChat conversations/messages. Mongo * fences continuation; optional Redis routing reaches the live owning process. */ @@ -14,6 +22,7 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( { acquireSubagentThreadLease: db.acquireSubagentThreadLease, claimSubagentTaskResult: db.claimSubagentTaskResult, + releaseSubagentTaskResultClaim: db.releaseSubagentTaskResultClaim, countActiveSubagentThreadLeases: db.countActiveSubagentThreadLeases, deleteConvos: db.deleteConvos, deleteMessages: db.deleteMessages, @@ -31,6 +40,9 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( fenceOwnerAdmission: db.fenceSubagentAdmission, renewOwnerAdmission: db.renewSubagentAdmission, releaseOwnerAdmission: db.releaseSubagentAdmission, + ...(completionWakeupsEnabled && { + onTaskPrepared: createSubagentCompletionWakeupHandler(enqueueAgentTrigger), + }), }, ); diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index f79b560f11a..cd77218f36b 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -32,6 +32,7 @@ export * from './skills'; export * from './phases'; export * from './startup'; export * from './subagentThreads'; +export * from './subagentCompletionWakeup'; export * from './subagentTaskRouting'; export * from './skillConfigurable'; export * from './skillFiles'; diff --git a/packages/api/src/agents/subagentCompletionWakeup.spec.ts b/packages/api/src/agents/subagentCompletionWakeup.spec.ts new file mode 100644 index 00000000000..9bd0f2df189 --- /dev/null +++ b/packages/api/src/agents/subagentCompletionWakeup.spec.ts @@ -0,0 +1,483 @@ +import type { IMessage } from '@librechat/data-schemas'; +import type { AgentContinueTriggerEnvelope } from './triggers/envelope'; +import type { SubagentTaskWakeupRegistration } from './subagentThreads'; +import type { EnqueueAgentTrigger } from './subagentCompletionWakeup'; +import { + createAgentTriggerEnvelope, + getAgentTriggerIdempotencyKey, + parseAgentTriggerEnvelope, +} from './triggers/envelope'; +import { + createSubagentCompletionWakeupHandler, + createSubagentCompletionWakeupResolver, +} from './subagentCompletionWakeup'; + +const NOW = 1_775_000_000_000; + +function enqueueMock(): jest.MockedFunction { + return jest.fn, Parameters>(async () => ({ + id: 'delivery-1', + })); +} + +function registration( + overrides: Partial = {}, +): SubagentTaskWakeupRegistration { + return { + userId: 'user-1', + tenantId: 'tenant-1', + parentConversationId: 'conversation-1', + parentMessageId: 'response-1', + parentAgentId: 'agent_parent_1', + taskId: 'task-1', + threadId: 'thread-1', + subagentType: 'researcher', + createdAt: NOW - 10, + ...overrides, + }; +} + +describe('createSubagentCompletionWakeupHandler', () => { + beforeEach(() => { + jest.useFakeTimers().setSystemTime(NOW); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('pre-registers a bounded continuation on the exact parent branch', async () => { + const enqueue = enqueueMock(); + const notify = createSubagentCompletionWakeupHandler(enqueue); + + await notify(registration()); + + expect(enqueue).toHaveBeenCalledTimes(1); + const [envelopeValue, options] = enqueue.mock.calls[0]!; + const envelope = parseAgentTriggerEnvelope(envelopeValue); + expect(envelope).toMatchObject({ + version: 1, + mode: 'continue', + principal: { userId: 'user-1', tenantId: 'tenant-1' }, + target: { + agentId: 'agent_parent_1', + conversationId: 'conversation-1', + parentMessageId: 'response-1', + }, + event: { + id: 'task-1', + type: 'subagent.completion', + source: { id: 'subagent-completion', type: 'internal' }, + payload: { + taskId: 'task-1', + threadId: 'thread-1', + subagentType: 'researcher', + }, + }, + }); + expect(envelope.input).toContain('waiting to complete'); + expect(options).toEqual({ + orderingKey: 'subagent-completion:conversation-1', + availableAt: new Date(NOW + 250), + }); + }); + + it('keeps one idempotency identity across duplicate registration callbacks', async () => { + const enqueue = enqueueMock(); + const notify = createSubagentCompletionWakeupHandler(enqueue); + const event = registration(); + + await notify(event); + await notify(event); + + const first = parseAgentTriggerEnvelope(enqueue.mock.calls[0]![0]); + const retry = parseAgentTriggerEnvelope(enqueue.mock.calls[1]![0]); + expect(first.requestId).not.toBe(retry.requestId); + expect(first.deliveryId).toBe('task-1'); + expect(getAgentTriggerIdempotencyKey(first)).toBe(getAgentTriggerIdempotencyKey(retry)); + expect(first.input).toContain('waiting to complete'); + }); + + it('does not enqueue without a stable initiating agent', async () => { + const enqueue = enqueueMock(); + const notify = createSubagentCompletionWakeupHandler(enqueue); + + await notify(registration({ parentAgentId: undefined })); + + expect(enqueue).not.toHaveBeenCalled(); + }); + + it('does not enqueue for an ephemeral initiating agent', async () => { + const enqueue = enqueueMock(); + const notify = createSubagentCompletionWakeupHandler(enqueue); + + await notify(registration({ parentAgentId: 'openAI__gpt-4o___GPT-4o____1' })); + + expect(enqueue).not.toHaveBeenCalled(); + }); +}); + +function wakeupEnvelope(): AgentContinueTriggerEnvelope { + const envelope = createAgentTriggerEnvelope({ + mode: 'continue', + requestId: 'request-1', + deliveryId: 'task-1', + receivedAt: NOW, + principal: { id: 'user-1', tenantId: 'tenant-1' }, + event: { + id: 'task-1', + type: 'subagent.completion', + occurredAt: NOW, + source: { id: 'subagent-completion', type: 'internal' }, + payload: { taskId: 'task-1', threadId: 'thread-1', subagentType: 'researcher' }, + }, + target: { + agentId: 'agent_parent_1', + conversationId: 'conversation-1', + parentMessageId: 'response-1', + }, + input: 'pending', + }); + if (envelope.mode !== 'continue') { + throw new Error('Expected a continue envelope.'); + } + return envelope; +} + +function resolverMethods() { + const subagentTask: IMessage['subagentTask'] = { + attemptKey: 'attempt-1', + parentRunId: 'response-1', + status: 'completed', + }; + const terminal = { + messageId: 'task-1:assistant', + conversationId: 'thread-1', + parentMessageId: 'task-1:user', + sender: 'researcher', + text: 'Child result', + isCreatedByUser: false, + createdAt: new Date(NOW), + updatedAt: new Date(NOW), + subagentTask, + }; + const methods = { + getConvo: jest.fn(async (_userId: string, conversationId: string) => + conversationId === 'conversation-1' + ? { conversationId, tenantId: 'tenant-1' } + : { + conversationId, + tenantId: 'tenant-1', + subagentThread: { + parentConversationId: 'conversation-1', + parentMessageId: 'response-1', + parentAgentId: 'agent_parent_1', + subagentType: 'researcher', + }, + }, + ), + getMessages: jest.fn(async (filter: { conversationId: string }) => + filter.conversationId === 'conversation-1' + ? [ + { + messageId: 'response-1', + parentMessageId: 'user-1', + isCreatedByUser: false, + createdAt: new Date(NOW - 30), + }, + { + messageId: 'wakeup-user', + parentMessageId: 'response-1', + isCreatedByUser: true, + createdAt: new Date(NOW - 20), + }, + { + messageId: 'wakeup-response', + parentMessageId: 'wakeup-user', + isCreatedByUser: false, + createdAt: new Date(NOW - 10), + }, + ] + : [ + { + messageId: 'task-1:user', + conversationId: 'thread-1', + isCreatedByUser: true, + }, + terminal, + ], + ), + claimSubagentTaskResult: jest.fn(async () => ({ status: 'acquired', message: terminal })), + releaseSubagentTaskResultClaim: jest.fn(async () => true), + }; + return { methods, terminal }; +} + +describe('createSubagentCompletionWakeupResolver', () => { + it('defers without claiming while the parent generation is active', async () => { + const { methods } = resolverMethods(); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ status: 'running' }), + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ + code: 'PARENT_NOT_READY', + retryable: true, + deferWithoutAttempt: true, + }); + expect(methods.claimSubagentTaskResult).not.toHaveBeenCalled(); + }); + + it('lets a lost-receipt retry reach HTTP dedup for its own active continuation', async () => { + const { methods } = resolverMethods(); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => ({ + status: 'requires_action', + metadata: { idempotencyClientRequestId: 'trigger_claim_1' }, + }), + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).resolves.toMatchObject({ status: 'ready' }); + }); + + it('bounds a persisted child result before rendering model input', async () => { + const { methods, terminal } = resolverMethods(); + terminal.text = 'x'.repeat(200_000); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + const prepared = await resolve(wakeupEnvelope(), { + idempotencyKey: 'trigger_claim_1', + } as never); + + expect(prepared).toMatchObject({ status: 'ready' }); + expect(prepared?.status === 'ready' && prepared.input.length).toBeLessThan(110_000); + }); + + it('dead-letters a child whose process disappeared after the task timeout grace', 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 fresh = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + now: () => NOW + 60_000, + }); + const stale = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + now: () => NOW + 36 * 60_000, + }); + + await expect( + fresh(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ + code: 'CHILD_NOT_READY', + retryable: true, + deferWithoutAttempt: true, + }); + await expect( + stale(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).rejects.toMatchObject({ code: 'CHILD_TASK_ABANDONED', retryable: false, status: 410 }); + expect(methods.claimSubagentTaskResult).not.toHaveBeenCalled(); + expect( + methods.getMessages.mock.calls.filter( + ([filter]) => filter.conversationId === 'conversation-1', + ), + ).toHaveLength(0); + }); + + it('resolves a crash-retry terminal by logical attempt without blocking its ordered lane', async () => { + const { methods, terminal } = resolverMethods(); + const supersedingTerminal = { + ...terminal, + messageId: 'task-2:assistant', + parentMessageId: 'task-1:user', + }; + methods.getMessages.mockImplementation(async (filter: Record) => { + if (filter.conversationId === 'conversation-1') { + return [ + { + messageId: 'response-1', + parentMessageId: 'user-1', + isCreatedByUser: false, + createdAt: new Date(NOW - 30), + }, + ]; + } + if (filter['subagentTask.attemptKey'] === 'attempt-1') { + return [supersedingTerminal]; + } + return [ + { + messageId: 'task-1:user', + conversationId: 'thread-1', + isCreatedByUser: true, + subagentTask: { + attemptKey: 'attempt-1', + parentRunId: 'response-1', + status: 'running', + }, + }, + ]; + }); + methods.claimSubagentTaskResult.mockResolvedValueOnce({ + status: 'acquired', + message: supersedingTerminal, + }); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + const prepared = await resolve(wakeupEnvelope(), { + idempotencyKey: 'trigger_claim_1', + } as never); + + expect(prepared).toMatchObject({ + status: 'ready', + input: expect.stringContaining('"background_task_id":"task-2"'), + }); + expect(methods.claimSubagentTaskResult).toHaveBeenCalledWith({ + userId: 'user-1', + conversationId: 'thread-1', + taskId: 'task-2', + kind: 'wakeup', + claimId: 'trigger_claim_1', + }); + }); + + it('validates a continued child against its per-task parent instead of original lineage', async () => { + const { methods } = resolverMethods(); + methods.getConvo.mockImplementation(async (_userId: string, conversationId: string) => + conversationId === 'conversation-1' + ? { conversationId, tenantId: 'tenant-1' } + : { + conversationId, + tenantId: 'tenant-1', + subagentThread: { + parentConversationId: 'conversation-1', + parentMessageId: 'original-response', + parentAgentId: 'agent_parent_1', + subagentType: 'researcher', + }, + }, + ); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).resolves.toMatchObject({ status: 'ready' }); + }); + + it('claims the durable result and chains onto the latest assistant descendant', async () => { + const { methods } = resolverMethods(); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).resolves.toMatchObject({ + status: 'ready', + parentMessageId: 'wakeup-response', + input: expect.stringContaining('Child result'), + }); + expect(methods.claimSubagentTaskResult).toHaveBeenCalledWith({ + userId: 'user-1', + conversationId: 'thread-1', + taskId: 'task-1', + kind: 'wakeup', + claimId: 'trigger_claim_1', + }); + + const prepared = await resolve(wakeupEnvelope(), { + idempotencyKey: 'trigger_claim_1', + } as never); + expect(prepared?.status).toBe('ready'); + if (prepared?.status === 'ready') { + await prepared.releaseOnDefiniteFailure?.(); + } + expect(methods.releaseSubagentTaskResultClaim).toHaveBeenCalledWith({ + userId: 'user-1', + conversationId: 'thread-1', + taskId: 'task-1', + kind: 'wakeup', + claimId: 'trigger_claim_1', + }); + }); + + it('settles without starting a turn when a manual poll already claimed the result', async () => { + const { methods, terminal } = resolverMethods(); + methods.claimSubagentTaskResult.mockResolvedValueOnce({ + status: 'claimed', + message: { + ...terminal, + subagentTask: { + ...terminal.subagentTask, + resultClaim: { kind: 'manual', claimId: 'poll-1', claimedAt: new Date(NOW) }, + }, + }, + }); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).resolves.toEqual({ status: 'settled' }); + }); + + it('releases a cancelled wakeup result for later explicit collection', async () => { + const { methods, terminal } = resolverMethods(); + terminal.subagentTask = { ...terminal.subagentTask!, status: 'cancelled' }; + methods.claimSubagentTaskResult.mockResolvedValueOnce({ + status: 'acquired', + message: terminal, + }); + const resolve = createSubagentCompletionWakeupResolver({ + methods: methods as never, + getGenerationJob: async () => null, + }); + + await expect( + resolve(wakeupEnvelope(), { idempotencyKey: 'trigger_claim_1' } as never), + ).resolves.toEqual({ status: 'settled' }); + expect(methods.releaseSubagentTaskResultClaim).toHaveBeenCalledWith({ + userId: 'user-1', + conversationId: 'thread-1', + taskId: 'task-1', + kind: 'wakeup', + claimId: 'trigger_claim_1', + }); + }); +}); diff --git a/packages/api/src/agents/subagentCompletionWakeup.ts b/packages/api/src/agents/subagentCompletionWakeup.ts new file mode 100644 index 00000000000..dd6329f37e4 --- /dev/null +++ b/packages/api/src/agents/subagentCompletionWakeup.ts @@ -0,0 +1,425 @@ +import { randomUUID } from 'node:crypto'; +import { isEphemeralAgentId } from 'librechat-data-provider'; +import type { ConversationMethods, IMessage, MessageMethods } from '@librechat/data-schemas'; +import type { + AgentTriggerContinuePreparation, + AgentTriggerExecutionHostDeps, +} from './triggers/host'; +import type { SubagentTaskWakeupRegistration } from './subagentThreads'; +import type { AgentContinueTriggerEnvelope } from './triggers/envelope'; +import type { AgentTriggerDispatchContext } from './triggers/dispatch'; +import type { AgentTriggerEnqueueOptions } from './triggers/delivery'; +import { boundedSubagentTaskResult } from './subagentTaskRouting'; +import { createAgentTriggerEnvelope } from './triggers/envelope'; +import { AgentTriggerExecutionError } from './triggers/host'; + +const WAKEUP_ADMISSION_DELAY_MS = 250; +/** SDK tasks time out after 30 minutes; this grace covers terminal persistence. */ +const CHILD_READY_WAIT_MS = 35 * 60_000; +const SOURCE_ID = 'subagent-completion'; +const EVENT_TYPE = 'subagent.completion'; +const MESSAGE_SELECT = 'messageId parentMessageId isCreatedByUser createdAt'; +const TASK_SELECT = + 'messageId conversationId parentMessageId sender text error createdAt updatedAt +subagentTask'; + +export type EnqueueAgentTrigger = ( + envelope: unknown, + options?: AgentTriggerEnqueueOptions, +) => Promise; + +type WakeupMethods = Pick & + Pick< + MessageMethods, + 'claimSubagentTaskResult' | 'getMessages' | 'releaseSubagentTaskResultClaim' + >; + +interface GenerationState { + status?: unknown; + metadata?: { + idempotencyClientRequestId?: unknown; + terminalPersistencePending?: unknown; + }; +} + +export interface SubagentCompletionWakeupResolverDeps { + methods: WakeupMethods; + getGenerationJob: (conversationId: string) => Promise; + now?: () => number; +} + +function payloadRegistration( + envelope: AgentContinueTriggerEnvelope, +): Pick | null | undefined { + if ( + envelope.event.source.type !== 'internal' || + envelope.event.source.id !== SOURCE_ID || + envelope.event.type !== EVENT_TYPE + ) { + return; + } + const payload = envelope.event.payload; + if (payload == null || typeof payload !== 'object' || Array.isArray(payload)) { + return null; + } + const { taskId, threadId, subagentType } = payload; + if ( + typeof taskId !== 'string' || + taskId.length === 0 || + taskId.length > 256 || + typeof threadId !== 'string' || + threadId.length === 0 || + threadId.length > 256 || + typeof subagentType !== 'string' || + subagentType.length === 0 || + subagentType.length > 256 + ) { + return null; + } + return { taskId, threadId, subagentType }; +} + +function executionError( + message: string, + options: { + code: string; + retryable: boolean; + deferWithoutAttempt?: boolean; + status?: number; + retryAfter?: string; + }, +): AgentTriggerExecutionError { + return new AgentTriggerExecutionError(message, { + mode: 'continue', + certainty: 'definite', + ...options, + }); +} + +function isParentActive(job: GenerationState | null): boolean { + return ( + job?.status === 'running' || + job?.status === 'requires_action' || + job?.metadata?.terminalPersistencePending === true + ); +} + +function sameTenant(actual: string | undefined, expected: string | undefined): boolean { + return actual === expected; +} + +function timestamp(message: Pick): number { + const value = message.createdAt; + if (value instanceof Date) { + return value.getTime(); + } + const parsed = value == null ? Number.NaN : new Date(value).getTime(); + return Number.isFinite(parsed) ? parsed : 0; +} + +/** Selects the newest persisted assistant on the branch below the original + * parent. Re-resolving for every ordered delivery serializes sibling child + * completions onto the branch produced by the preceding wakeup. */ +function latestAssistantDescendant(messages: IMessage[], anchorId: string): string | undefined { + const byId = new Map(messages.map((message) => [message.messageId, message])); + if (!byId.has(anchorId)) { + return; + } + const memo = new Map([[anchorId, true]]); + const reachesAnchor = (message: IMessage, visiting = new Set()): boolean => { + const known = memo.get(message.messageId); + if (known != null) { + return known; + } + if (visiting.has(message.messageId)) { + memo.set(message.messageId, false); + return false; + } + visiting.add(message.messageId); + const parent = + typeof message.parentMessageId === 'string' ? byId.get(message.parentMessageId) : undefined; + const reachable = parent != null && reachesAnchor(parent, visiting); + visiting.delete(message.messageId); + memo.set(message.messageId, reachable); + return reachable; + }; + const descendants = messages + .filter((message) => message.isCreatedByUser === false && reachesAnchor(message)) + .sort((left, right) => { + const time = timestamp(left) - timestamp(right); + return time === 0 ? left.messageId.localeCompare(right.messageId) : time; + }); + return descendants[descendants.length - 1]?.messageId; +} + +function renderWakeupInput( + registration: Pick, + resultTaskId: string, + terminal: IMessage, +): string { + const status = terminal.subagentTask?.status ?? 'error'; + return [ + `A detached subagent task has ${status}. Continue the parent task using its durable result below.`, + JSON.stringify({ + background_task_id: resultTaskId, + subagent_thread_id: registration.threadId, + subagent_type: registration.subagentType, + status, + result: boundedSubagentTaskResult(terminal.text ?? ''), + }), + ].join('\n'); +} + +/** Resolves a pre-registered completion delivery immediately before dispatch. + * The durable result claim elects exactly one consumer (manual poll or this + * delivery), while the branch lookup chains ordered sibling completions. */ +export function createSubagentCompletionWakeupResolver({ + methods, + getGenerationJob, + now = Date.now, +}: SubagentCompletionWakeupResolverDeps): NonNullable< + AgentTriggerExecutionHostDeps['prepareContinue'] +> { + return async ( + envelope: AgentContinueTriggerEnvelope, + context: AgentTriggerDispatchContext, + ): Promise => { + const registration = payloadRegistration(envelope); + if (registration === undefined) { + return; + } + if (registration === null) { + throw executionError('The subagent completion wakeup payload is invalid.', { + code: 'INVALID_SUBAGENT_WAKEUP', + retryable: false, + }); + } + + let parentJob: GenerationState | null; + try { + parentJob = await getGenerationJob(envelope.target.conversationId); + } catch (error) { + throw executionError( + `Parent generation state is temporarily unavailable: ${ + error instanceof Error ? error.message : String(error) + }`, + { code: 'PARENT_STATE_UNAVAILABLE', retryable: true }, + ); + } + if ( + isParentActive(parentJob) && + parentJob?.metadata?.idempotencyClientRequestId !== context.idempotencyKey + ) { + throw executionError('The parent generation has not settled yet.', { + code: 'PARENT_NOT_READY', + retryable: true, + status: 409, + retryAfter: '1', + deferWithoutAttempt: true, + }); + } + + const userId = envelope.principal.userId; + const tenantId = envelope.principal.tenantId; + const [parent, child, taskMessages] = await Promise.all([ + methods.getConvo(userId, envelope.target.conversationId), + methods.getConvo(userId, registration.threadId), + methods.getMessages( + { + user: userId, + conversationId: registration.threadId, + messageId: { $in: [`${registration.taskId}:user`, `${registration.taskId}:assistant`] }, + }, + TASK_SELECT, + { sort: { createdAt: 1, _id: 1 } }, + ), + ]); + if (parent == null || !sameTenant(parent.tenantId, tenantId)) { + throw executionError('The parent conversation is no longer available.', { + code: 'PARENT_NOT_FOUND', + retryable: false, + status: 404, + }); + } + const lineage = child?.subagentThread; + if ( + child == null || + !sameTenant(child.tenantId, tenantId) || + lineage?.parentConversationId !== envelope.target.conversationId || + lineage.parentAgentId !== envelope.target.agentId || + lineage.subagentType !== registration.subagentType + ) { + throw executionError('The child task lineage is no longer available.', { + code: 'CHILD_TASK_MISSING', + retryable: false, + status: 404, + }); + } + let resultTaskId = registration.taskId; + let terminal = taskMessages.find( + (message) => + message.messageId === `${registration.taskId}:assistant` && + message.subagentTask?.status !== 'running', + ); + const started = taskMessages.find( + (message) => message.messageId === `${registration.taskId}:user`, + ); + /** A worker can persist the input, lose its lease, and then have a retry + * close the same logical attempt under the retry's runtime task id. Resolve + * that terminal by the durable attempt identity so the earlier ordered + * delivery cannot block the repaired delivery behind it for the full + * abandonment grace period. */ + if (terminal == null && started?.subagentTask?.attemptKey != null) { + const [supersedingTerminal] = await methods.getMessages( + { + user: userId, + conversationId: registration.threadId, + 'subagentTask.attemptKey': started.subagentTask.attemptKey, + 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + }, + TASK_SELECT, + { sort: { createdAt: -1, _id: -1 }, limit: 1 }, + ); + if (supersedingTerminal?.messageId.endsWith(':assistant') === true) { + terminal = supersedingTerminal; + resultTaskId = supersedingTerminal.messageId.slice(0, -':assistant'.length); + } + } + if (terminal == null) { + if (started != null) { + if (now() - envelope.event.occurredAt > CHILD_READY_WAIT_MS) { + throw executionError('The child task owner disappeared before settlement.', { + code: 'CHILD_TASK_ABANDONED', + retryable: false, + status: 410, + }); + } + throw executionError('The child task has not settled yet.', { + code: 'CHILD_NOT_READY', + retryable: true, + status: 409, + retryAfter: '1', + deferWithoutAttempt: true, + }); + } + throw executionError('The child task no longer exists.', { + code: 'CHILD_TASK_MISSING', + retryable: false, + status: 404, + }); + } + if (terminal.subagentTask?.parentRunId !== envelope.target.parentMessageId) { + throw executionError('The child task lineage is no longer available.', { + code: 'CHILD_TASK_MISSING', + retryable: false, + status: 404, + }); + } + + const parentMessages = await methods.getMessages( + { user: userId, conversationId: envelope.target.conversationId }, + MESSAGE_SELECT, + { sort: { createdAt: 1, _id: 1 } }, + ); + + const parentMessageId = latestAssistantDescendant( + parentMessages, + envelope.target.parentMessageId, + ); + if (parentMessageId == null) { + throw executionError('The parent conversation branch is no longer available.', { + code: 'PARENT_NOT_FOUND', + retryable: false, + status: 404, + }); + } + + const claim = await methods.claimSubagentTaskResult({ + userId, + conversationId: registration.threadId, + taskId: resultTaskId, + kind: 'wakeup', + claimId: context.idempotencyKey, + }); + if (claim.status !== 'acquired') { + return { status: 'settled' }; + } + if (claim.message.subagentTask?.status === 'cancelled') { + const released = await methods.releaseSubagentTaskResultClaim({ + userId, + conversationId: registration.threadId, + taskId: resultTaskId, + kind: 'wakeup', + claimId: context.idempotencyKey, + }); + if (!released) { + throw executionError('The cancelled child result claim could not be released.', { + code: 'RESULT_CLAIM_RELEASE_FAILED', + retryable: true, + }); + } + return { status: 'settled' }; + } + return { + status: 'ready', + parentMessageId, + input: renderWakeupInput(registration, resultTaskId, claim.message), + releaseOnDefiniteFailure: async () => { + await methods.releaseSubagentTaskResultClaim({ + userId, + conversationId: registration.threadId, + taskId: resultTaskId, + kind: 'wakeup', + claimId: context.idempotencyKey, + }); + }, + }; + }; +} + +/** Pre-registers the idempotent delivery before child provider work starts. + * A process crash can therefore delay a wakeup but cannot lose it; dispatch + * simply defers until the terminal child message exists. */ +export function createSubagentCompletionWakeupHandler( + enqueue: EnqueueAgentTrigger, +): (registration: SubagentTaskWakeupRegistration) => Promise { + return async (registration) => { + const parentAgentId = registration.parentAgentId?.trim(); + if (parentAgentId == null || parentAgentId === '' || isEphemeralAgentId(parentAgentId)) { + return; + } + const eventId = registration.taskId; + const envelope = createAgentTriggerEnvelope({ + mode: 'continue', + requestId: randomUUID(), + deliveryId: eventId, + receivedAt: Date.now(), + principal: { + id: registration.userId, + ...(registration.tenantId == null ? {} : { tenantId: registration.tenantId }), + }, + event: { + id: eventId, + type: EVENT_TYPE, + occurredAt: registration.createdAt, + source: { id: SOURCE_ID, type: 'internal' }, + payload: { + taskId: registration.taskId, + threadId: registration.threadId, + subagentType: registration.subagentType, + }, + }, + target: { + agentId: parentAgentId, + conversationId: registration.parentConversationId, + parentMessageId: registration.parentMessageId, + }, + input: 'A detached subagent task is waiting to complete.', + }); + await enqueue(envelope, { + orderingKey: `subagent-completion:${registration.parentConversationId}`, + availableAt: new Date( + Math.max(Date.now(), registration.createdAt) + WAKEUP_ADMISSION_DELAY_MS, + ), + }); + }; +} diff --git a/packages/api/src/agents/subagentTaskRouting.ts b/packages/api/src/agents/subagentTaskRouting.ts index 7496bbdafd9..d4488776669 100644 --- a/packages/api/src/agents/subagentTaskRouting.ts +++ b/packages/api/src/agents/subagentTaskRouting.ts @@ -289,6 +289,11 @@ export function boundedTaskList(tasks: SubagentTaskSnapshot[]): SubagentTaskSnap return [...keptRunning, ...settled.slice(-remaining)].sort(byCreatedAt); } +/** Applies the shared model-facing bound to a durable child result. */ +export function boundedSubagentTaskResult(result: string): string { + return truncateMiddle(result, MAX_RESULT_CHARS); +} + /** Applies the routed result and snapshot bounds to a claim from any source. */ export function boundedClaim(claim: SubagentTaskClaim): SubagentTaskClaim { if (claim.status === 'not_found') { @@ -296,7 +301,7 @@ export function boundedClaim(claim: SubagentTaskClaim): SubagentTaskClaim { } const task = boundedSnapshot(claim.task); if (claim.status === 'completed') { - return { status: 'completed', task, result: truncateMiddle(claim.result, MAX_RESULT_CHARS) }; + return { status: 'completed', task, result: boundedSubagentTaskResult(claim.result) }; } if (claim.status === 'error' || claim.status === 'cancelled') { return { status: claim.status, task, error: truncateMiddle(claim.error, MAX_ERROR_CHARS) }; diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index a68b3bb0545..649bb8a0105 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -19,6 +19,7 @@ import type { SubagentTaskControlHandler, SubagentTaskControlTransport, } from './subagentTaskRouting'; +import type { SubagentTaskWakeupRegistration } from './subagentThreads'; import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; import { buildSubagentThreadTaskConfig, @@ -285,6 +286,68 @@ describe('SubagentThreadTaskStore', () => { }); }); + it('registers a host-safe wakeup before child provider work begins', async () => { + const userId = 'wakeup-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const run = jest.fn(taskRequest('').run); + const onTaskPrepared = jest.fn(async (registration: SubagentTaskWakeupRegistration) => { + const messages = await methods.getMessages({ + user: userId, + conversationId: registration.threadId, + messageId: `${registration.taskId}:assistant`, + }); + expect(messages).toHaveLength(0); + expect(run).not.toHaveBeenCalled(); + }); + const store = new SubagentThreadTaskStore(methods, { onTaskPrepared }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + parentRunId: 'parent-response-1', + parentAgentId: 'agent_parent_1', + run, + }), + ); + await waitForSettled(store, config.scopeId, started); + + const settledTask = store.get(config.scopeId, requireAccepted(started).task.taskId); + expect(settledTask?.error).toBeUndefined(); + expect(settledTask).toMatchObject({ + status: 'completed', + }); + expect(onTaskPrepared).toHaveBeenCalledWith({ + userId, + parentConversationId, + parentMessageId: 'parent-response-1', + parentAgentId: 'agent_parent_1', + taskId: requireAccepted(started).task.taskId, + threadId: requireThreadId(started), + subagentType: 'researcher-agent', + createdAt: expect.any(Number), + }); + }); + + it('fails before provider work and keeps the durable failure collectable when registration fails', async () => { + const userId = 'wakeup-failure-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const run = jest.fn(taskRequest('').run); + const store = new SubagentThreadTaskStore(methods, { + onTaskPrepared: async () => Promise.reject(new Error('trigger queue unavailable')), + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).not.toHaveBeenCalled(); + await expect( + store.claimTask(config.scopeId, requireAccepted(started).task.taskId), + ).resolves.toMatchObject({ + status: 'error', + }); + }); + it('waits for initial parent persistence before creating the first child', async () => { const userId = 'parent-gate-user'; const parentConversationId = randomUUID(); @@ -452,22 +515,35 @@ describe('SubagentThreadTaskStore', () => { const userId = 'durable-idempotency-user'; const parentConversationId = randomUUID(); await saveParent(userId, parentConversationId); - const firstWorker = new SubagentThreadTaskStore(methods); - const secondWorker = new SubagentThreadTaskStore(methods); + const firstWakeup = jest.fn(async (_registration: SubagentTaskWakeupRegistration) => undefined); + const replayWakeup = jest.fn( + async (_registration: SubagentTaskWakeupRegistration) => undefined, + ); + const firstWorker = new SubagentThreadTaskStore(methods, { onTaskPrepared: firstWakeup }); + const secondWorker = new SubagentThreadTaskStore(methods, { onTaskPrepared: replayWakeup }); const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); const firstRun = jest.fn(async () => ({ content: 'Original durable result.', messages: [new HumanMessage('Run once.'), new AIMessage('Original durable result.')], })); + const firstParentRunId = 'original-parent-response'; const first = firstWorker.start( taskRequest(config.scopeId, { idempotencyKey: 'cross-worker-attempt', + parentRunId: firstParentRunId, requestFingerprint: 'same-inputs', input: 'Run once.', run: firstRun, }), ); await waitForSettled(firstWorker, config.scopeId, first); + const durableAttempt = await methods.getMessages( + { user: userId, conversationId: requireThreadId(first) }, + '+subagentTask', + ); + expect(durableAttempt[durableAttempt.length - 1]?.subagentTask?.parentRunId).toBe( + firstParentRunId, + ); const replayRun = jest.fn(taskRequest(config.scopeId).run); const replay = secondWorker.start( @@ -483,6 +559,18 @@ describe('SubagentThreadTaskStore', () => { expect(firstRun).toHaveBeenCalledTimes(1); expect(replayRun).not.toHaveBeenCalled(); + expect(firstWakeup).toHaveBeenCalledTimes(1); + const firstRegistration = firstWakeup.mock.calls[0]?.[0]; + const replayRegistration = replayWakeup.mock.calls[0]?.[0]; + expect(firstRegistration?.createdAt).toBe(durableAttempt[0]?.createdAt?.getTime()); + expect(replayRegistration?.createdAt).toBe(firstRegistration?.createdAt); + expect(replayWakeup).toHaveBeenCalledWith( + expect.objectContaining({ + taskId: requireAccepted(first).task.taskId, + parentMessageId: firstParentRunId, + createdAt: firstRegistration?.createdAt, + }), + ); expect(secondWorker.claim(config.scopeId, requireAccepted(replay).task.taskId)).toMatchObject({ status: 'completed', result: 'Original durable result.', diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index 415ef9533f8..becac411041 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -101,10 +101,14 @@ interface PreparedThread { initialMessages: BaseMessage[]; initialStoredMessages: StoredMessage[]; attemptKey: string; + /** Stable source-occurrence time shared by first delivery and every replay. */ + taskCreatedAt: number; userMessageId?: string; replay?: { status: 'completed' | 'error' | 'cancelled'; content: string; + taskId: string; + parentRunId: string; }; } @@ -141,6 +145,19 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor fenceOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; renewOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; releaseOwnerAdmission?: (userId: string, token: string) => Promise; + onTaskPrepared?: (registration: SubagentTaskWakeupRegistration) => Promise | void; +} + +export interface SubagentTaskWakeupRegistration { + userId: string; + parentConversationId: string; + parentMessageId: string; + parentAgentId?: string; + tenantId?: string; + taskId: string; + threadId: string; + subagentType: string; + createdAt: number; } function positiveInteger(value: number | undefined, fallback: number): number { @@ -191,6 +208,14 @@ function matchesTenant(actual: string | undefined, expected: string | undefined) return actual === expected; } +function durableMessageTime(message: Pick, missingMessage: string): number { + const value = message.createdAt?.getTime(); + if (!Number.isSafeInteger(value) || value == null || value < 0) { + throw new Error(missingMessage); + } + return value; +} + function assertParentPersistence( value: unknown, scope: SubagentThreadScope, @@ -387,6 +412,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ) => Promise; private readonly releaseOwnerAdmission?: (userId: string, token: string) => Promise; + private readonly onTaskPrepared?: SubagentThreadTaskStoreOptions['onTaskPrepared']; private taskControlTransport?: SubagentTaskControlTransport; constructor( @@ -418,6 +444,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { this.fenceOwnerAdmission = options.fenceOwnerAdmission; this.renewOwnerAdmission = options.renewOwnerAdmission; this.releaseOwnerAdmission = options.releaseOwnerAdmission; + this.onTaskPrepared = options.onTaskPrepared; } /** Enables optional cross-replica lookup after the host's Redis service is ready. */ @@ -496,6 +523,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { lease.taskId = runtime.taskId; lease.running = true; const detachedUsage: UsageMetadata[] = []; + let prepared: PreparedThread | undefined; try { if (runtime.signal.aborted) { throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); @@ -510,7 +538,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { this.taskRoutingTtlMs, ); await parentReady; - const prepared = await this.prepareThread( + prepared = await this.prepareThread( request.scopeId, scope, threadId, @@ -519,6 +547,11 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { runtime.taskId, lease, ); + await this.registerTaskWakeup(scope, prepared.conversation.conversationId, request, { + taskId: prepared.replay?.taskId ?? runtime.taskId, + parentRunId: prepared.replay?.parentRunId ?? request.parentRunId, + createdAt: prepared.taskCreatedAt, + }); if (runtime.signal.aborted) { throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); } @@ -533,8 +566,9 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { 'This child thread is already being continued by another run.', ); } + const preparedThread = prepared; const result = await runWithDetachedSubagentUsage(detachedUsage, () => - request.run(runtime, prepared.initialMessages), + request.run(runtime, preparedThread.initialMessages), ); if (runtime.signal.aborted) { throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); @@ -555,6 +589,11 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ); return result; } catch (error) { + /** A replay is already terminal in Mongo. A temporary wakeup-queue + * outage must not overwrite that canonical result with a new error. */ + if (prepared?.replay != null) { + throw error; + } const mayPersist = lease.shared == null || (await this.renewSharedLease(scope, threadId, lease)); const terminalTask = this.get(request.scopeId, runtime.taskId); @@ -778,6 +817,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { userId, conversationId: threadId, taskId, + kind: 'manual', claimId: invocationId, }); } catch (error) { @@ -1553,13 +1593,29 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { .reverse() .find((message) => message.subagentTask?.status !== 'running'); if (terminal?.subagentTask != null) { + const canonicalTaskId = terminal.messageId.endsWith(':assistant') + ? terminal.messageId.slice(0, -':assistant'.length) + : ''; + if (canonicalTaskId === '') { + throw new Error('The prior subagent result has an invalid task identity.'); + } + const canonicalStart = priorAttempt.find( + (message) => message.messageId === `${canonicalTaskId}:user`, + ); + const taskCreatedAt = durableMessageTime( + canonicalStart ?? terminal, + 'The prior subagent result has no durable occurrence time.', + ); return { conversation, initialMessages: [], initialStoredMessages: [], attemptKey, + taskCreatedAt, replay: { status: terminal.subagentTask.status as 'completed' | 'error' | 'cancelled', + taskId: canonicalTaskId, + parentRunId: terminal.subagentTask.parentRunId ?? request.parentRunId, content: terminal.text ?? (terminal.subagentTask.status === 'completed' @@ -1588,6 +1644,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { error: true, subagentTask: { attemptKey, + parentRunId: request.parentRunId, ...(requestFingerprint == null ? {} : { requestFingerprint }), status: 'error', }, @@ -1604,7 +1661,16 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { initialMessages: [], initialStoredMessages: [], attemptKey, - replay: { status: 'error', content: abandonedMessage }, + taskCreatedAt: durableMessageTime( + savedAbandoned, + 'The abandoned subagent result has no durable occurrence time.', + ), + replay: { + status: 'error', + content: abandonedMessage, + taskId, + parentRunId: request.parentRunId, + }, }; } const branch = selectLatestBranch(allMessages); @@ -1631,6 +1697,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { isCreatedByUser: true, subagentTask: { attemptKey, + parentRunId: request.parentRunId, ...(requestFingerprint == null ? {} : { requestFingerprint }), status: 'running', }, @@ -1650,6 +1717,10 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { initialMessages, initialStoredMessages: mapChatMessagesToStoredMessages(initialMessages), attemptKey, + taskCreatedAt: durableMessageTime( + savedUserMessage, + 'The child-thread input has no durable occurrence time.', + ), userMessageId, }; } catch (error) { @@ -1732,6 +1803,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ...(subagentTranscript == null ? {} : { subagentTranscript }), subagentTask: { attemptKey: prepared.attemptKey, + parentRunId: request.parentRunId, ...(normalizedRequestFingerprint(request) == null ? {} : { requestFingerprint: normalizedRequestFingerprint(request) }), @@ -1775,6 +1847,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { error: true, subagentTask: { attemptKey: createSubagentAttemptKey(request.scopeId, request.idempotencyKey), + parentRunId: request.parentRunId, ...(normalizedRequestFingerprint(request) == null ? {} : { requestFingerprint: normalizedRequestFingerprint(request) }), @@ -1791,6 +1864,28 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { await this.touchAfterMessage(scope, threadId, taskId, 'failed'); } + private async registerTaskWakeup( + scope: SubagentThreadScope, + threadId: string, + request: SubagentTaskStartRequest, + task: { taskId: string; parentRunId: string; createdAt: number }, + ): Promise { + if (this.onTaskPrepared == null) { + return; + } + await this.onTaskPrepared({ + userId: scope.userId, + parentConversationId: scope.parentConversationId, + parentMessageId: task.parentRunId, + ...(request.parentAgentId == null ? {} : { parentAgentId: request.parentAgentId }), + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + taskId: task.taskId, + threadId, + subagentType: request.subagentType, + createdAt: task.createdAt, + }); + } + private async persistCancellation( scope: SubagentThreadScope, threadId: string, @@ -1816,6 +1911,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { unfinished: false, subagentTask: { attemptKey: createSubagentAttemptKey(request.scopeId, request.idempotencyKey), + parentRunId: request.parentRunId, ...(normalizedRequestFingerprint(request) == null ? {} : { requestFingerprint: normalizedRequestFingerprint(request) }), diff --git a/packages/api/src/agents/triggers/README.md b/packages/api/src/agents/triggers/README.md index 1da00fe38e2..f4da03ac9a1 100644 --- a/packages/api/src/agents/triggers/README.md +++ b/packages/api/src/agents/triggers/README.md @@ -11,6 +11,9 @@ envelope and calls `enqueueAgentTrigger`; the adapter does not invoke an agent r - Give each source event a stable `event.id`, and keep `deliveryId` stable for retries to one target. A retry may use a fresh `requestId` and `receivedAt`. - Render bounded model input on the host. Infrastructure and routing remain server-controlled. +- Use `continue` only with a persisted `conversationId` and exact `parentMessageId`. The host defers + that delivery while the parent generation is still running or paused, so it cannot replace the + generation it is meant to follow. - Use `orderingKey` only when deliveries must remain ordered across different event sources. Without an override, ordering is scoped to the user, source, mode, agent, and conversation. @@ -43,7 +46,7 @@ await enqueueAgentTrigger( - Mongo owns queue state, leases, retry history, and dead letters across restarts and replicas. - A fresh token fences every claim, including reclaims by the same process. -- A delivery is at-least-once. Fire and steer admission reuse the envelope's stable idempotency +- A delivery is at-least-once. Fire, continue, and steer admission reuse the envelope's stable idempotency identity, so ambiguous retries do not duplicate accepted work. - Retryable failures use bounded exponential backoff and honor `Retry-After`. Invalid envelopes, permanent authorization failures, and exhausted retries become durable dead letters. diff --git a/packages/api/src/agents/triggers/delivery.ts b/packages/api/src/agents/triggers/delivery.ts index 5d511b263ed..970ab945f1c 100644 --- a/packages/api/src/agents/triggers/delivery.ts +++ b/packages/api/src/agents/triggers/delivery.ts @@ -75,7 +75,7 @@ function orderingIdentity( envelope.event.source.id, envelope.mode, envelope.target.agentId, - envelope.mode === 'steer' ? envelope.target.conversationId : '', + envelope.mode === 'fire' ? '' : envelope.target.conversationId, ]; } return `trigger_lane_${digest([ diff --git a/packages/api/src/agents/triggers/dispatch.spec.ts b/packages/api/src/agents/triggers/dispatch.spec.ts index b86ae40b205..f8848ffa65d 100644 --- a/packages/api/src/agents/triggers/dispatch.spec.ts +++ b/packages/api/src/agents/triggers/dispatch.spec.ts @@ -24,11 +24,16 @@ describe('dispatchAgentTrigger', () => { it('routes fire deliveries with their stable idempotency identity and abort signal', async () => { const controller = new AbortController(); const fire = jest.fn(async () => ({ status: 'accepted' as const })); + const continueRun = jest.fn(async () => ({ status: 'accepted' as const })); const steer = jest.fn(async () => ({ status: 'accepted' as const })); const envelope = fireEnvelope(); await expect( - dispatchAgentTrigger(envelope, { fire, steer }, { signal: controller.signal }), + dispatchAgentTrigger( + envelope, + { continue: continueRun, fire, steer }, + { signal: controller.signal }, + ), ).resolves.toEqual({ status: 'accepted' }); expect(fire).toHaveBeenCalledWith(envelope, { @@ -36,10 +41,12 @@ describe('dispatchAgentTrigger', () => { signal: controller.signal, }); expect(steer).not.toHaveBeenCalled(); + expect(continueRun).not.toHaveBeenCalled(); }); it('routes steer deliveries without requiring a fire implementation detail', async () => { const fire = jest.fn(async () => 'fire'); + const continueRun = jest.fn(async () => 'continue'); const steer = jest.fn(async () => 'steer'); const envelope = createAgentTriggerEnvelope({ ...createFireInput(), @@ -51,64 +58,99 @@ describe('dispatchAgentTrigger', () => { }, }); - await expect(dispatchAgentTrigger(envelope, { fire, steer })).resolves.toBe('steer'); + await expect( + dispatchAgentTrigger(envelope, { continue: continueRun, fire, steer }), + ).resolves.toBe('steer'); expect(steer).toHaveBeenCalledWith(envelope, { idempotencyKey: expect.stringMatching(/^trigger_[a-f0-9]{64}$/), }); expect(fire).not.toHaveBeenCalled(); }); + it('routes continue deliveries to the exact existing conversation branch', async () => { + const fire = jest.fn(async () => 'fire'); + const continueRun = jest.fn(async () => 'continue'); + const steer = jest.fn(async () => 'steer'); + const envelope = createAgentTriggerEnvelope({ + ...createFireInput(), + mode: 'continue', + target: { + agentId: 'agent-1', + conversationId: 'conversation-1', + parentMessageId: 'response-1', + }, + }); + + await expect( + dispatchAgentTrigger(envelope, { continue: continueRun, fire, steer }), + ).resolves.toBe('continue'); + expect(continueRun).toHaveBeenCalledWith(envelope, { + idempotencyKey: expect.stringMatching(/^trigger_[a-f0-9]{64}$/), + }); + expect(fire).not.toHaveBeenCalled(); + expect(steer).not.toHaveBeenCalled(); + }); + it('propagates handler failures without falling back to another mode', async () => { const error = new Error('fire rejected'); const fire = jest.fn(async () => Promise.reject(error)); + const continueRun = jest.fn(async () => 'continue'); const steer = jest.fn(async () => 'steer'); - await expect(dispatchAgentTrigger(fireEnvelope(), { fire, steer })).rejects.toBe(error); + await expect( + dispatchAgentTrigger(fireEnvelope(), { continue: continueRun, fire, steer }), + ).rejects.toBe(error); expect(fire).toHaveBeenCalledTimes(1); expect(steer).not.toHaveBeenCalled(); }); it('rejects unknown modes before deriving identity or calling a handler', () => { const fire = jest.fn(async () => 'fire'); + const continueRun = jest.fn(async () => 'continue'); const steer = jest.fn(async () => 'steer'); const envelope = { ...fireEnvelope(), - mode: 'resume', + mode: 'launch', } as unknown as AgentTriggerEnvelope; - expect(() => dispatchAgentTrigger(envelope, { fire, steer })).toThrow( - new AgentTriggerDispatchError('Unsupported agent trigger mode: resume'), + expect(() => dispatchAgentTrigger(envelope, { continue: continueRun, fire, steer })).toThrow( + new AgentTriggerDispatchError('Unsupported agent trigger mode: launch'), ); expect(fire).not.toHaveBeenCalled(); expect(steer).not.toHaveBeenCalled(); + expect(continueRun).not.toHaveBeenCalled(); }); it('rejects unknown envelope versions before deriving identity or calling a handler', () => { const fire = jest.fn(async () => 'fire'); + const continueRun = jest.fn(async () => 'continue'); const steer = jest.fn(async () => 'steer'); const envelope = { ...fireEnvelope(), version: 2, } as unknown as AgentTriggerEnvelope; - expect(() => dispatchAgentTrigger(envelope, { fire, steer })).toThrow( + expect(() => dispatchAgentTrigger(envelope, { continue: continueRun, fire, steer })).toThrow( new AgentTriggerDispatchError('Unsupported agent trigger envelope version: 2'), ); expect(fire).not.toHaveBeenCalled(); expect(steer).not.toHaveBeenCalled(); + expect(continueRun).not.toHaveBeenCalled(); }); it('rejects malformed v1 envelopes before deriving identity or calling a handler', () => { const fire = jest.fn(async () => 'fire'); + const continueRun = jest.fn(async () => 'continue'); const steer = jest.fn(async () => 'steer'); const malformed = { ...fireEnvelope() }; Reflect.deleteProperty(malformed, 'target'); const envelope = malformed as unknown as AgentTriggerEnvelope; - expect(() => dispatchAgentTrigger(envelope, { fire, steer })).toThrow( + expect(() => dispatchAgentTrigger(envelope, { continue: continueRun, fire, steer })).toThrow( new AgentTriggerDispatchError('target must be an object'), ); expect(fire).not.toHaveBeenCalled(); expect(steer).not.toHaveBeenCalled(); + expect(continueRun).not.toHaveBeenCalled(); }); }); diff --git a/packages/api/src/agents/triggers/dispatch.ts b/packages/api/src/agents/triggers/dispatch.ts index 39b7a75107c..e8ceb196c9d 100644 --- a/packages/api/src/agents/triggers/dispatch.ts +++ b/packages/api/src/agents/triggers/dispatch.ts @@ -1,4 +1,5 @@ import type { + AgentContinueTriggerEnvelope, AgentFireTriggerEnvelope, AgentSteerTriggerEnvelope, AgentTriggerEnvelope, @@ -21,11 +22,15 @@ export class AgentTriggerDispatchError extends TypeError { * Host-owned execution adapters. Each handler must enforce current authorization, * limits, persistence, and the supplied idempotency identity before accepting work. */ -export interface AgentTriggerDispatchHandlers { +export interface AgentTriggerDispatchHandlers { fire: ( envelope: AgentFireTriggerEnvelope, context: AgentTriggerDispatchContext, ) => Promise; + continue: ( + envelope: AgentContinueTriggerEnvelope, + context: AgentTriggerDispatchContext, + ) => Promise; steer: ( envelope: AgentSteerTriggerEnvelope, context: AgentTriggerDispatchContext, @@ -33,11 +38,11 @@ export interface AgentTriggerDispatchHandlers { } /** Routes a normalized trigger without coupling its source to an execution transport. */ -export function dispatchAgentTrigger( +export function dispatchAgentTrigger( envelope: unknown, - handlers: AgentTriggerDispatchHandlers, + handlers: AgentTriggerDispatchHandlers, options?: { signal?: AbortSignal }, -): Promise { +): Promise { let normalized: AgentTriggerEnvelope; try { normalized = parseAgentTriggerEnvelope(envelope); @@ -51,5 +56,8 @@ export function dispatchAgentTrigger( if (normalized.mode === 'fire') { return handlers.fire(normalized, context); } + if (normalized.mode === 'continue') { + return handlers.continue(normalized, context); + } return handlers.steer(normalized, context); } diff --git a/packages/api/src/agents/triggers/engine.spec.ts b/packages/api/src/agents/triggers/engine.spec.ts index 4a0eaff1f54..522c3629703 100644 --- a/packages/api/src/agents/triggers/engine.spec.ts +++ b/packages/api/src/agents/triggers/engine.spec.ts @@ -280,6 +280,41 @@ describe('createAgentTriggerDeliveryEngine', () => { expect(store.dead).not.toHaveBeenCalled(); }); + it('defers a continuation until its parent generation settles without consuming an attempt', async () => { + const store = storeWith(); + const engine = createAgentTriggerDeliveryEngine( + { + store, + dispatch: async () => + Promise.reject( + new AgentTriggerExecutionError('parent generation is still running', { + mode: 'continue', + certainty: 'definite', + retryable: true, + deferWithoutAttempt: true, + code: 'PARENT_NOT_READY', + status: 409, + }), + ), + now: () => START, + workerId: 'worker-1', + }, + { concurrency: 1, maxAttempts: 1 }, + ); + + await engine.runTick(); + + expect(store.defer).toHaveBeenCalledWith({ + id: 'delivery-row-1', + workerId: 'worker-1', + claimToken: 'claim-1', + attempt: 1, + availableAt: new Date(START.getTime() + 5_000), + }); + expect(store.retry).not.toHaveBeenCalled(); + expect(store.dead).not.toHaveBeenCalled(); + }); + it('does not shorten Retry-After to the exponential backoff cap', async () => { const store = storeWith(); const error = new AgentTriggerExecutionError('maintenance', { diff --git a/packages/api/src/agents/triggers/engine.ts b/packages/api/src/agents/triggers/engine.ts index eb3e0a42a95..24bd0eff342 100644 --- a/packages/api/src/agents/triggers/engine.ts +++ b/packages/api/src/agents/triggers/engine.ts @@ -211,6 +211,10 @@ function isAccountDeletionDeferral(error: unknown): boolean { ); } +function isRuntimeReadinessDeferral(error: unknown): boolean { + return error instanceof AgentTriggerExecutionError && error.deferWithoutAttempt; +} + /** Durable, lease-fenced delivery runner shared by every trusted event source. */ export function createAgentTriggerDeliveryEngine( deps: AgentTriggerDeliveryEngineDeps, @@ -326,10 +330,12 @@ export function createAgentTriggerDeliveryEngine( const attemptedAt = now(); const deletionCancelled = controller.signal.aborted && cancelledUsers.has(userId); const deletionRejected = isAccountDeletionDeferral(error); + const runtimeNotReady = isRuntimeReadinessDeferral(error); if ( error instanceof AgentTriggerDeliveryDeferredError || deletionCancelled || - deletionRejected + deletionRejected || + runtimeNotReady ) { const delayMs = error instanceof AgentTriggerDeliveryDeferredError ? error.delayMs : DEFAULT_DEFER_MS; @@ -342,9 +348,15 @@ export function createAgentTriggerDeliveryEngine( availableAt, }); if (deferred) { + let reason = 'pre_dispatch'; + if (deletionCancelled || deletionRejected) { + reason = 'account_deletion'; + } else if (runtimeNotReady) { + reason = 'runtime_readiness'; + } logger.info('[agent-triggers] delivery deferred without consuming an attempt', { deliveryKey: delivery.deliveryKey, - reason: deletionCancelled || deletionRejected ? 'account_deletion' : 'pre_dispatch', + reason, availableAt: availableAt.toISOString(), }); } diff --git a/packages/api/src/agents/triggers/envelope.spec.ts b/packages/api/src/agents/triggers/envelope.spec.ts index ad93e632305..da869e1c638 100644 --- a/packages/api/src/agents/triggers/envelope.spec.ts +++ b/packages/api/src/agents/triggers/envelope.spec.ts @@ -93,6 +93,26 @@ describe('createAgentTriggerEnvelope', () => { }); }); + it('requires an exact existing branch for continue deliveries', () => { + const envelope = createAgentTriggerEnvelope({ + ...createFireInput(), + mode: 'continue', + target: { + agentId: 'agent-1', + conversationId: 'conversation-1', + parentMessageId: 'response-1', + }, + }); + + expect(envelope.mode).toBe('continue'); + expect(envelope.target).toEqual({ + agentId: 'agent-1', + conversationId: 'conversation-1', + parentMessageId: 'response-1', + }); + expect(parseAgentTriggerEnvelope(JSON.parse(JSON.stringify(envelope)))).toEqual(envelope); + }); + it('builds a stable generation-compatible idempotency key per delivery target', () => { const first = createAgentTriggerEnvelope(createFireInput()); const retry = createAgentTriggerEnvelope({ @@ -210,8 +230,8 @@ describe('createAgentTriggerEnvelope', () => { expect(() => createAgentTriggerEnvelope({ ...createFireInput(), - mode: 'resume', + mode: 'launch', } as unknown as CreateAgentTriggerEnvelopeInput), - ).toThrow('Unsupported agent trigger mode: resume'); + ).toThrow('Unsupported agent trigger mode: launch'); }); }); diff --git a/packages/api/src/agents/triggers/envelope.ts b/packages/api/src/agents/triggers/envelope.ts index bdc761e6046..2391fb0293d 100644 --- a/packages/api/src/agents/triggers/envelope.ts +++ b/packages/api/src/agents/triggers/envelope.ts @@ -6,7 +6,7 @@ import { cloneJsonValue } from '../json'; export const AGENT_TRIGGER_ENVELOPE_VERSION = 1 as const; export const AGENT_TRIGGER_IDEMPOTENCY_PREFIX = 'trigger_'; -export type AgentTriggerMode = 'fire' | 'steer'; +export type AgentTriggerMode = 'continue' | 'fire' | 'steer'; export interface AgentTriggerSource { /** Stable identity of the configured source, such as a webhook or schedule id. */ @@ -36,6 +36,13 @@ interface AgentTriggerTarget { */ export type AgentFireTarget = AgentTriggerTarget; +export interface AgentContinueTarget extends AgentTriggerTarget { + /** Existing conversation that receives a new host-authored turn. */ + conversationId: string; + /** Persisted branch leaf below which the new turn is appended. */ + parentMessageId: string; +} + export interface AgentSteerTarget extends AgentTriggerTarget { /** Existing conversation whose active generation receives the input. */ conversationId: string; @@ -63,12 +70,20 @@ export interface AgentFireTriggerEnvelope extends AgentTriggerEnvelopeBase { target: AgentFireTarget; } +export interface AgentContinueTriggerEnvelope extends AgentTriggerEnvelopeBase { + mode: 'continue'; + target: AgentContinueTarget; +} + export interface AgentSteerTriggerEnvelope extends AgentTriggerEnvelopeBase { mode: 'steer'; target: AgentSteerTarget; } -export type AgentTriggerEnvelope = AgentFireTriggerEnvelope | AgentSteerTriggerEnvelope; +export type AgentTriggerEnvelope = + | AgentContinueTriggerEnvelope + | AgentFireTriggerEnvelope + | AgentSteerTriggerEnvelope; interface CreateAgentTriggerEnvelopeBase { requestId: string; @@ -84,6 +99,10 @@ export type CreateAgentTriggerEnvelopeInput = mode: 'fire'; target: AgentFireTarget; }) + | (CreateAgentTriggerEnvelopeBase & { + mode: 'continue'; + target: AgentContinueTarget; + }) | (CreateAgentTriggerEnvelopeBase & { mode: 'steer'; target: AgentSteerTarget; @@ -191,6 +210,18 @@ export function createAgentTriggerEnvelope( }; } + if (input.mode === 'continue') { + return { + ...base, + mode: input.mode, + target: { + agentId: requireString(input.target?.agentId, 'target.agentId'), + conversationId: requireString(input.target?.conversationId, 'target.conversationId'), + parentMessageId: requireString(input.target?.parentMessageId, 'target.parentMessageId'), + }, + }; + } + throw error(`Unsupported agent trigger mode: ${receivedMode}`); } @@ -205,7 +236,7 @@ export function parseAgentTriggerEnvelope(input: unknown): AgentTriggerEnvelope } const mode = envelope.mode; - if (mode !== 'fire' && mode !== 'steer') { + if (mode !== 'continue' && mode !== 'fire' && mode !== 'steer') { throw error(`Unsupported agent trigger mode: ${String(mode)}`); } @@ -252,6 +283,18 @@ export function parseAgentTriggerEnvelope(input: unknown): AgentTriggerEnvelope }; } + if (mode === 'continue') { + return { + ...base, + mode, + target: { + agentId: requireString(target.agentId, 'target.agentId'), + conversationId: requireString(target.conversationId, 'target.conversationId'), + parentMessageId: requireString(target.parentMessageId, 'target.parentMessageId'), + }, + }; + } + if (target.preempt != null && typeof target.preempt !== 'boolean') { throw error('target.preempt must be a boolean'); } @@ -288,7 +331,8 @@ export function getAgentTriggerIdempotencyKey(envelope: AgentTriggerEnvelope): s envelope.deliveryId, envelope.mode, envelope.target.agentId, - envelope.mode === 'steer' ? envelope.target.conversationId : '', + envelope.mode === 'fire' ? '' : envelope.target.conversationId, + envelope.mode === 'continue' ? envelope.target.parentMessageId : '', ]), ) .digest('hex'); diff --git a/packages/api/src/agents/triggers/host.spec.ts b/packages/api/src/agents/triggers/host.spec.ts index 86f4ca6f702..c9c9624064c 100644 --- a/packages/api/src/agents/triggers/host.spec.ts +++ b/packages/api/src/agents/triggers/host.spec.ts @@ -44,6 +44,27 @@ const createSteerEnvelope = () => input: 'The opponent moved. Take your turn.', }); +const createContinueEnvelope = () => + createAgentTriggerEnvelope({ + mode: 'continue', + requestId: 'request-3', + deliveryId: 'delivery-3', + receivedAt: 35, + principal: { id: 'user-1', role: 'member', tenantId: 'tenant-1' }, + target: { + agentId: 'agent-1', + conversationId: 'conversation-1', + parentMessageId: 'response-1', + }, + event: { + id: 'event-3', + type: 'subagent.completed', + occurredAt: 31, + source: { id: 'subagent-completion', type: 'internal' }, + }, + input: 'Collect the completed child task.', + }); + function response(payload: unknown, init?: ResponseInit): Response { return new Response(JSON.stringify(payload), { status: 200, @@ -294,6 +315,48 @@ describe('createAgentTriggerExecutionHost fire adapter', () => { expect(fetcher).not.toHaveBeenCalled(); }); + it('releases a prepared result that arrives after setup has timed out', async () => { + let finishPreparation!: (value: { + status: 'ready'; + input: string; + parentMessageId: string; + releaseOnDefiniteFailure: () => Promise; + }) => void; + const preparation = new Promise<{ + status: 'ready'; + input: string; + parentMessageId: string; + releaseOnDefiniteFailure: () => Promise; + }>((resolve) => { + finishPreparation = resolve; + }); + const releaseOnDefiniteFailure = jest.fn(async () => undefined); + const fetcher = fetchMock(async () => response({})); + const host = createAgentTriggerExecutionHost( + deps(fetcher, { + prepareContinue: () => preparation, + timeoutMs: 10, + }), + ); + + await expect(host.dispatch(createContinueEnvelope())).rejects.toMatchObject({ + mode: 'continue', + certainty: 'definite', + retryable: true, + code: 'TIMEOUT', + }); + finishPreparation({ + status: 'ready', + input: 'late durable child result', + parentMessageId: 'response-1', + releaseOnDefiniteFailure, + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(releaseOnDefiniteFailure).toHaveBeenCalledTimes(1); + expect(fetcher).not.toHaveBeenCalled(); + }); + it('starts independent token, timezone, and origin setup concurrently', async () => { let resolveToken!: (value: string) => void; let resolveTimezone!: (value: string) => void; @@ -455,6 +518,219 @@ describe('createAgentTriggerExecutionHost fire adapter', () => { }); }); +describe('createAgentTriggerExecutionHost continue adapter', () => { + it('appends an idempotent turn to the exact existing conversation branch', async () => { + const envelope = createContinueEnvelope(); + const idempotencyKey = getAgentTriggerIdempotencyKey(envelope); + const fetcher = fetchMock(async () => + response({ + streamId: 'conversation-1', + conversationId: 'conversation-1', + generationCreatedAt: 50, + status: 'started', + }), + ); + const host = createAgentTriggerExecutionHost(deps(fetcher)); + + await expect(host.dispatch(envelope)).resolves.toEqual({ + mode: 'continue', + streamId: 'conversation-1', + conversationId: 'conversation-1', + generationCreatedAt: 50, + status: 'started', + }); + const [input, init] = fetcher.mock.calls[0]; + expect(String(input)).toBe('http://127.0.0.1:3080/api/agents/chat/agents'); + expect(JSON.parse(String(init?.body))).toEqual({ + text: envelope.input, + endpoint: EModelEndpoint.agents, + agent_id: 'agent-1', + parentMessageId: 'response-1', + conversationId: 'conversation-1', + isContinued: false, + isRegenerate: false, + clientRequestId: idempotencyKey, + generationProtocolVersion: 2, + }); + }); + + it('retries without consuming the logical delivery when the parent is not settled', async () => { + expect.hasAssertions(); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => + response( + { code: 'PARENT_NOT_READY', error: 'The parent is still running.' }, + { status: 409 }, + ), + ), + ), + ); + + await host.dispatch(createContinueEnvelope()).catch((error: unknown) => { + expectExecutionError(error, { + mode: 'continue', + certainty: 'definite', + retryable: true, + deferWithoutAttempt: true, + code: 'PARENT_NOT_READY', + status: 409, + }); + }); + }); + + it('releases a prepared durable result after a definite admission rejection', async () => { + const releaseOnDefiniteFailure = jest.fn(async () => undefined); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => response({ code: 'AGENT_NOT_FOUND' }, { status: 404 })), + { + prepareContinue: async () => ({ + status: 'ready', + input: 'durable child result', + parentMessageId: 'response-1', + releaseOnDefiniteFailure, + }), + }, + ), + ); + + await expect(host.dispatch(createContinueEnvelope())).rejects.toMatchObject({ + certainty: 'definite', + code: 'AGENT_NOT_FOUND', + }); + expect(releaseOnDefiniteFailure).toHaveBeenCalledTimes(1); + }); + + it('retains a prepared durable result after an ambiguous admission outcome', async () => { + const releaseOnDefiniteFailure = jest.fn(async () => undefined); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => Promise.reject(new Error('connection reset'))), + { + prepareContinue: async () => ({ + status: 'ready', + input: 'durable child result', + parentMessageId: 'response-1', + releaseOnDefiniteFailure, + }), + }, + ), + ); + + await expect(host.dispatch(createContinueEnvelope())).rejects.toMatchObject({ + certainty: 'ambiguous', + code: 'NETWORK_ERROR', + }); + expect(releaseOnDefiniteFailure).not.toHaveBeenCalled(); + }); + + it('retains a prepared durable result when a retry gets a definite 5xx response', async () => { + const releaseOnDefiniteFailure = jest.fn(async () => undefined); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => + response( + { code: 'SERVER_NOT_READY', error: 'Generation is finalizing.' }, + { status: 503, headers: { 'retry-after': '1' } }, + ), + ), + { + prepareContinue: async () => ({ + status: 'ready', + input: 'durable child result', + parentMessageId: 'response-1', + releaseOnDefiniteFailure, + }), + }, + ), + ); + + await expect(host.dispatch(createContinueEnvelope())).rejects.toMatchObject({ + certainty: 'definite', + retryable: true, + code: 'SERVER_NOT_READY', + status: 503, + }); + expect(releaseOnDefiniteFailure).not.toHaveBeenCalled(); + }); + + it('releases a prepared durable result when parent state fails before admission', async () => { + const releaseOnDefiniteFailure = jest.fn(async () => undefined); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => + response( + { code: 'PARENT_STATE_UNAVAILABLE', error: 'Parent state is unavailable.' }, + { status: 503, headers: { 'retry-after': '1' } }, + ), + ), + { + prepareContinue: async () => ({ + status: 'ready', + input: 'durable child result', + parentMessageId: 'response-1', + releaseOnDefiniteFailure, + }), + }, + ), + ); + + await expect(host.dispatch(createContinueEnvelope())).rejects.toMatchObject({ + certainty: 'definite', + retryable: true, + code: 'PARENT_STATE_UNAVAILABLE', + status: 503, + }); + expect(releaseOnDefiniteFailure).toHaveBeenCalledTimes(1); + }); + + it('retains a prepared durable result when an earlier admitted run was replaced', async () => { + const releaseOnDefiniteFailure = jest.fn(async () => undefined); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => response({ code: 'RUN_REPLACED' }, { status: 409 })), + { + prepareContinue: async () => ({ + status: 'ready', + input: 'durable child result', + parentMessageId: 'response-1', + releaseOnDefiniteFailure, + }), + }, + ), + ); + + await expect(host.dispatch(createContinueEnvelope())).rejects.toMatchObject({ + certainty: 'definite', + retryable: false, + code: 'RUN_REPLACED', + status: 409, + }); + expect(releaseOnDefiniteFailure).not.toHaveBeenCalled(); + }); + + it('rejects a mismatched continued conversation as an ambiguous outcome', async () => { + expect.hasAssertions(); + const host = createAgentTriggerExecutionHost( + deps( + fetchMock(async () => + response({ streamId: 'other', conversationId: 'other', status: 'started' }), + ), + ), + ); + + await host.dispatch(createContinueEnvelope()).catch((error: unknown) => { + expectExecutionError(error, { + mode: 'continue', + certainty: 'ambiguous', + retryable: true, + code: 'INVALID_RESPONSE', + }); + }); + }); +}); + describe('createAgentTriggerExecutionHost steer adapter', () => { it('steers through the authenticated admission route with a strict v2 receipt', async () => { const envelope = createSteerEnvelope(); diff --git a/packages/api/src/agents/triggers/host.ts b/packages/api/src/agents/triggers/host.ts index 730a75c10de..c3a33464d71 100644 --- a/packages/api/src/agents/triggers/host.ts +++ b/packages/api/src/agents/triggers/host.ts @@ -1,9 +1,11 @@ -import { tenantStorage } from '@librechat/data-schemas'; +import { logger, tenantStorage } from '@librechat/data-schemas'; import { Constants, EModelEndpoint } from 'librechat-data-provider'; import type { + AgentContinueTriggerEnvelope, AgentFireTriggerEnvelope, AgentSteerTriggerEnvelope, AgentTriggerEnvelope, + AgentTriggerMode, } from './envelope'; import type { AgentTriggerDispatchContext } from './dispatch'; import type { AgentRunPrincipal } from '../envelope'; @@ -30,12 +32,25 @@ type FireStatus = 'started' | 'resumed' | 'replaced' | 'settled'; export type AgentTriggerFetch = (input: string | URL, init?: RequestInit) => Promise; +export type AgentTriggerContinuePreparation = + | { + status: 'ready'; + input: string; + parentMessageId: string; + /** Compensates a durable pre-admission claim only when the host knows + * that no generation was admitted. Ambiguous outcomes retain the claim. */ + releaseOnDefiniteFailure?: () => MaybePromise; + } + | { status: 'settled' }; + export type AgentTriggerFailureCertainty = 'definite' | 'ambiguous'; export interface AgentTriggerExecutionErrorOptions { - mode: 'fire' | 'steer'; + mode: AgentTriggerMode; certainty: AgentTriggerFailureCertainty; retryable: boolean; + /** Release the delivery lease without consuming its logical retry budget. */ + deferWithoutAttempt?: boolean; code?: string; status?: number; retryAfter?: string; @@ -47,9 +62,10 @@ export interface AgentTriggerExecutionErrorOptions { * remains unchanged. */ export class AgentTriggerExecutionError extends Error { - readonly mode: 'fire' | 'steer'; + readonly mode: AgentTriggerMode; readonly certainty: AgentTriggerFailureCertainty; readonly retryable: boolean; + readonly deferWithoutAttempt: boolean; readonly code?: string; readonly status?: number; readonly retryAfter?: string; @@ -60,6 +76,7 @@ export class AgentTriggerExecutionError extends Error { this.mode = options.mode; this.certainty = options.certainty; this.retryable = options.retryable; + this.deferWithoutAttempt = options.deferWithoutAttempt === true; this.code = options.code; this.status = options.status; this.retryAfter = options.retryAfter; @@ -87,7 +104,18 @@ export interface AgentTriggerSteerResult { leftover?: boolean; } -export type AgentTriggerExecutionResult = AgentTriggerFireResult | AgentTriggerSteerResult; +export interface AgentTriggerContinueResult { + mode: 'continue'; + status: FireStatus; + conversationId: string; + streamId?: string; + generationCreatedAt?: number; +} + +export type AgentTriggerExecutionResult = + | AgentTriggerContinueResult + | AgentTriggerFireResult + | AgentTriggerSteerResult; export interface AgentTriggerExecutionHostDeps { /** Trusted root URL for this LibreChat server. */ @@ -97,8 +125,14 @@ export interface AgentTriggerExecutionHostDeps { /** Optional user-timezone resolver for dynamic date variables in a new run. */ getTimezone?: ( principal: AgentRunPrincipal, - envelope: AgentFireTriggerEnvelope, + envelope: AgentContinueTriggerEnvelope | AgentFireTriggerEnvelope, ) => MaybePromise; + /** Optional server-owned resolver for durable internal continuation inputs. + * External/source-neutral envelopes remain unchanged when this returns undefined. */ + prepareContinue?: ( + envelope: AgentContinueTriggerEnvelope, + context: AgentTriggerDispatchContext, + ) => MaybePromise; fetch?: AgentTriggerFetch; /** Total bound for setup, admission, and the bounded response read. */ timeoutMs?: number; @@ -163,7 +197,7 @@ function abortScope(parent: AbortSignal | undefined, timeoutMs: number): AbortSc } function abortError( - mode: 'fire' | 'steer', + mode: AgentTriggerMode, scope: AbortScope, parent: AbortSignal | undefined, stage: string, @@ -183,15 +217,18 @@ function abortError( function observeAbort( operation: () => MaybePromise, - mode: 'fire' | 'steer', + mode: AgentTriggerMode, scope: AbortScope, parent: AbortSignal | undefined, + onLateValue?: (value: T) => MaybePromise, ): Promise { if (scope.signal.aborted) { return Promise.reject(abortError(mode, scope, parent, 'before dispatch', 'definite')); } return new Promise((resolve, reject) => { + let abandoned = false; const onAbort = () => { + abandoned = true; scope.signal.removeEventListener('abort', onAbort); reject(abortError(mode, scope, parent, 'during setup', 'definite')); }; @@ -201,6 +238,12 @@ function observeAbort( .then( (value) => { scope.signal.removeEventListener('abort', onAbort); + if (abandoned) { + Promise.resolve(onLateValue?.(value)).catch((error: unknown) => { + logger.error('[agentTriggers] Failed to compensate late setup result', error); + }); + return; + } resolve(value); }, (error: unknown) => { @@ -213,12 +256,13 @@ function observeAbort( async function setupValue( operation: () => MaybePromise, - mode: 'fire' | 'steer', + mode: AgentTriggerMode, scope: AbortScope, parent: AbortSignal | undefined, + onLateValue?: (value: T) => MaybePromise, ): Promise { try { - return await observeAbort(operation, mode, scope, parent); + return await observeAbort(operation, mode, scope, parent, onLateValue); } catch (error) { if (error instanceof AgentTriggerExecutionError) { throw error; @@ -342,7 +386,7 @@ function abortCode(scope: AbortScope, parent: AbortSignal | undefined, fallback: return fallback; } -function requireToken(value: unknown, mode: 'fire' | 'steer'): string { +function requireToken(value: unknown, mode: AgentTriggerMode): string { if (typeof value !== 'string' || value.length === 0 || /\s/.test(value)) { throw executionError('Agent trigger token mint returned an invalid token', { mode, @@ -354,7 +398,7 @@ function requireToken(value: unknown, mode: 'fire' | 'steer'): string { return value; } -function triggerUrl(baseUrl: string, path: string, mode: 'fire' | 'steer'): string { +function triggerUrl(baseUrl: string, path: string, mode: AgentTriggerMode): string { let url: URL; try { url = new URL(baseUrl); @@ -384,6 +428,10 @@ function fireUrl(baseUrl: string): string { return triggerUrl(baseUrl, `/api/agents/chat/${EModelEndpoint.agents}`, 'fire'); } +function continueUrl(baseUrl: string): string { + return triggerUrl(baseUrl, `/api/agents/chat/${EModelEndpoint.agents}`, 'continue'); +} + function steerUrl(baseUrl: string): string { return triggerUrl(baseUrl, '/api/agents/chat/steer/deliver', 'steer'); } @@ -402,7 +450,10 @@ function fireStatus(value: unknown): FireStatus | undefined { : undefined; } -function parseFireResult(payload: unknown): AgentTriggerFireResult | undefined { +function parseStartResult( + payload: unknown, + mode: 'continue' | 'fire', +): AgentTriggerContinueResult | AgentTriggerFireResult | undefined { if (payload == null || typeof payload !== 'object') { return undefined; } @@ -421,7 +472,7 @@ function parseFireResult(payload: unknown): AgentTriggerFireResult | undefined { 'generationCreatedAt' in payload ? payload.generationCreatedAt : undefined, ); return { - mode: 'fire', + mode, status, conversationId, ...(streamId != null && { streamId }), @@ -429,33 +480,101 @@ function parseFireResult(payload: unknown): AgentTriggerFireResult | undefined { }; } -async function fire( +function resolveParentMessageId( + preparation: AgentTriggerContinuePreparation | undefined, + envelope: AgentContinueTriggerEnvelope | AgentFireTriggerEnvelope, +): string { + if (preparation?.status === 'ready') { + return preparation.parentMessageId; + } + if (envelope.mode === 'continue') { + return envelope.target.parentMessageId; + } + return Constants.NO_PARENT; +} + +/** A response can be definite at the HTTP layer while the idempotent logical + * generation is still outcome-ambiguous. In particular, a retry can receive a + * 5xx while the first request owns the generation claim or is finalizing its + * accepted run. Compensate the prepared durable result only for failures that + * prove admission did not happen. */ +function canReleasePreparedResult(error: AgentTriggerExecutionError): boolean { + if (error.certainty !== 'definite') { + return false; + } + if (error.code === 'START_ABORTED' || error.status == null) { + return true; + } + if (error.code === 'PARENT_NOT_READY' || error.code === 'PARENT_STATE_UNAVAILABLE') { + return true; + } + return error.status >= 400 && error.status < 500 && error.status !== 408 && error.status !== 409; +} + +function startRun( envelope: AgentFireTriggerEnvelope, context: AgentTriggerDispatchContext, deps: AgentTriggerExecutionHostDeps, timeoutMs: number, -): Promise { +): Promise; +function startRun( + envelope: AgentContinueTriggerEnvelope, + context: AgentTriggerDispatchContext, + deps: AgentTriggerExecutionHostDeps, + timeoutMs: number, +): Promise; +async function startRun( + envelope: AgentContinueTriggerEnvelope | AgentFireTriggerEnvelope, + context: AgentTriggerDispatchContext, + deps: AgentTriggerExecutionHostDeps, + timeoutMs: number, +): Promise { + const mode = envelope.mode; const scope = abortScope(context.signal, timeoutMs); + let preparation: AgentTriggerContinuePreparation | undefined; try { + preparation = + mode === 'continue' && deps.prepareContinue != null + ? await setupValue( + () => deps.prepareContinue?.(envelope, context), + mode, + scope, + context.signal, + async (latePreparation) => { + if (latePreparation?.status === 'ready') { + await latePreparation.releaseOnDefiniteFailure?.(); + } + }, + ) + : undefined; + if (preparation?.status === 'settled' && envelope.mode === 'continue') { + return { + mode: 'continue', + status: 'settled', + conversationId: envelope.target.conversationId, + }; + } + const input = preparation?.status === 'ready' ? preparation.input : envelope.input; + const parentMessageId = resolveParentMessageId(preparation, envelope); const [token, timezone, baseUrl] = await Promise.all([ setupValue( () => deps.mintToken(envelope.principal, envelope), - 'fire', + mode, scope, context.signal, - ).then((value) => requireToken(value, 'fire')), + ).then((value) => requireToken(value, mode)), setupValue( () => deps.getTimezone?.(envelope.principal, envelope), - 'fire', + mode, scope, context.signal, ), - setupValue(() => deps.getBaseUrl(), 'fire', scope, context.signal), + setupValue(() => deps.getBaseUrl(), mode, scope, context.signal), ]).catch((error: unknown) => { scope.abort(); throw error; }); - const url = fireUrl(baseUrl); + const url = mode === 'fire' ? fireUrl(baseUrl) : continueUrl(baseUrl); const fetcher: AgentTriggerFetch = deps.fetch ?? globalThis.fetch; let response: Response; try { @@ -472,10 +591,13 @@ async function fire( [GENERATION_PROTOCOL_HEADER]: '2', }, body: JSON.stringify({ - text: envelope.input, + text: input, endpoint: EModelEndpoint.agents, agent_id: envelope.target.agentId, - parentMessageId: Constants.NO_PARENT, + parentMessageId, + ...(envelope.mode === 'continue' && { + conversationId: envelope.target.conversationId, + }), isContinued: false, isRegenerate: false, clientRequestId: context.idempotencyKey, @@ -489,9 +611,9 @@ async function fire( const definite = isDefiniteConnectFailure(error); const message = error instanceof Error ? error.message : String(error); throw executionError( - `Agent trigger fire ${definite ? 'could not connect' : 'has an unknown outcome'}: ${message}`, + `Agent trigger ${mode} ${definite ? 'could not connect' : 'has an unknown outcome'}: ${message}`, { - mode: 'fire', + mode, certainty: definite ? 'definite' : 'ambiguous', retryable: true, code: abortCode(scope, context.signal, 'NETWORK_ERROR'), @@ -505,11 +627,11 @@ async function fire( } catch (error) { if (response.ok) { throw executionError( - `Agent trigger fire response has an unknown outcome: ${ + `Agent trigger ${mode} response has an unknown outcome: ${ error instanceof Error ? error.message : String(error) }`, { - mode: 'fire', + mode, certainty: 'ambiguous', retryable: true, code: abortCode(scope, context.signal, 'INVALID_RESPONSE'), @@ -522,11 +644,19 @@ async function fire( if (!response.ok) { const message = errorMessage(payload) ?? (boundedBody.text.slice(0, 300) || 'request rejected'); - throw executionError(`Agent trigger fire was rejected (${response.status}): ${message}`, { - mode: 'fire', + throw executionError(`Agent trigger ${mode} was rejected (${response.status}): ${message}`, { + mode, certainty: 'definite', - retryable: isRetryableStatus(response.status), - code: errorCode(payload) ?? 'FIRE_REJECTED', + retryable: + isRetryableStatus(response.status) || + (mode === 'continue' && + response.status === 409 && + errorCode(payload) === 'PARENT_NOT_READY'), + deferWithoutAttempt: + mode === 'continue' && + response.status === 409 && + errorCode(payload) === 'PARENT_NOT_READY', + code: errorCode(payload) ?? (mode === 'fire' ? 'FIRE_REJECTED' : 'CONTINUE_REJECTED'), status: response.status, ...(response.headers.get('retry-after') != null && { retryAfter: response.headers.get('retry-after') ?? undefined, @@ -534,8 +664,8 @@ async function fire( }); } if (boundedBody.truncated) { - throw executionError('Agent trigger fire returned an oversized success response', { - mode: 'fire', + throw executionError(`Agent trigger ${mode} returned an oversized success response`, { + mode, certainty: 'ambiguous', retryable: true, code: 'RESPONSE_TOO_LARGE', @@ -548,18 +678,27 @@ async function fire( 'status' in payload && payload.status === 'aborted' ) { - throw executionError('Agent trigger fire was aborted before generation started', { - mode: 'fire', + throw executionError(`Agent trigger ${mode} was aborted before generation started`, { + mode, certainty: 'definite', retryable: false, code: 'START_ABORTED', status: response.status, }); } - const result = parseFireResult(payload); + const result = parseStartResult(payload, mode); if (result == null) { - throw executionError('Agent trigger fire returned an invalid success response', { - mode: 'fire', + throw executionError(`Agent trigger ${mode} returned an invalid success response`, { + mode, + certainty: 'ambiguous', + retryable: true, + code: 'INVALID_RESPONSE', + status: response.status, + }); + } + if (mode === 'continue' && result.conversationId !== envelope.target.conversationId) { + throw executionError('Agent trigger continue returned a mismatched conversation', { + mode, certainty: 'ambiguous', retryable: true, code: 'INVALID_RESPONSE', @@ -567,6 +706,30 @@ async function fire( }); } return result; + } catch (error) { + if ( + preparation?.status === 'ready' && + preparation.releaseOnDefiniteFailure != null && + error instanceof AgentTriggerExecutionError && + canReleasePreparedResult(error) + ) { + try { + await preparation.releaseOnDefiniteFailure(); + } catch (releaseError) { + throw executionError( + `Agent trigger ${mode} could not release its rejected preparation: ${ + releaseError instanceof Error ? releaseError.message : String(releaseError) + }`, + { + mode, + certainty: 'definite', + retryable: true, + code: 'PREPARATION_RELEASE_FAILED', + }, + ); + } + } + throw error; } finally { scope.cleanup(); } @@ -800,7 +963,13 @@ export function createAgentTriggerExecutionHost( envelope, { fire: (normalized, context) => - runAsPrincipal(normalized, context, () => fire(normalized, context, deps, timeoutMs)), + runAsPrincipal(normalized, context, () => + startRun(normalized, context, deps, timeoutMs), + ), + continue: (normalized, context) => + runAsPrincipal(normalized, context, () => + startRun(normalized, context, deps, timeoutMs), + ), steer: (normalized, context) => runAsPrincipal(normalized, context, () => steer(normalized, context, deps, timeoutMs), diff --git a/packages/api/src/agents/triggers/service.ts b/packages/api/src/agents/triggers/service.ts index ad8e155821b..dc2694bd07d 100644 --- a/packages/api/src/agents/triggers/service.ts +++ b/packages/api/src/agents/triggers/service.ts @@ -34,6 +34,7 @@ export interface AgentTriggerServiceOptions { export interface AgentTriggerServiceDeps { fetch?: AgentTriggerExecutionHostDeps['fetch']; getTimezone?: AgentTriggerExecutionHostDeps['getTimezone']; + prepareContinue?: AgentTriggerExecutionHostDeps['prepareContinue']; mintToken?: AgentTriggerExecutionHostDeps['mintToken']; timeoutMs?: number; methods?: AgentTriggerDeliveryPersistence; @@ -201,6 +202,7 @@ export function createAgentTriggerService(deps: AgentTriggerServiceDeps = {}): A ((principal) => generateAgentTriggerToken(principal.userId, AGENT_TRIGGER_TOKEN_TTL)), ...(deps.fetch != null && { fetch: deps.fetch }), ...(deps.getTimezone != null && { getTimezone: deps.getTimezone }), + ...(deps.prepareContinue != null && { prepareContinue: deps.prepareContinue }), ...(deps.timeoutMs != null && { timeoutMs: deps.timeoutMs }), }); diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 5078051aa7c..b4a859d8964 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -465,6 +465,9 @@ export interface CreateGenerationJobOptions { * status result. Creation may proceed only if that exact epoch is still * current or the stream has no durable job. */ expectedPredecessorCreatedAt?: number; + /** Atomically refuse to replace a running/paused predecessor while allowing + * an absent or terminal predecessor. Used by automatic continuations. */ + rejectActivePredecessor?: boolean; } /** @@ -1973,6 +1976,12 @@ class GenerationJobManagerClass { ) { throw new Error('Invalid expected generation predecessor'); } + if ( + options.rejectActivePredecessor != null && + typeof options.rejectActivePredecessor !== 'boolean' + ) { + throw new Error('Invalid active generation predecessor policy'); + } const tenantId = getTenantId(); const safeTenantId = tenantId && tenantId !== SYSTEM_TENANT_ID ? tenantId : undefined; @@ -2010,6 +2019,7 @@ class GenerationJobManagerClass { options.recoveredSteerPayload, creationAttemptId, options.expectedPredecessorCreatedAt, + options.rejectActivePredecessor, ); } catch (error) { if (error instanceof JobPredecessorMismatchError) { diff --git a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts index 7197f88aba7..8506f13f5a0 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts @@ -32,8 +32,8 @@ async function waitFor(predicate: () => boolean): Promise { function jobHashFromCreationCall(call: unknown[]): Record { const keyCount = Number(call[1]); - // JOB_CREATE_LUA receives twelve scalar arguments before its HSET pairs. - const fields = call.slice(14 + keyCount); + // JOB_CREATE_LUA receives thirteen scalar arguments before its HSET pairs. + const fields = call.slice(15 + keyCount); const hash = Object.fromEntries( Array.from({ length: fields.length / 2 }, (_, index) => [ String(fields[index * 2]), diff --git a/packages/api/src/stream/__tests__/predecessorFence.spec.ts b/packages/api/src/stream/__tests__/predecessorFence.spec.ts index d5381dd2ecf..50e20695d0c 100644 --- a/packages/api/src/stream/__tests__/predecessorFence.spec.ts +++ b/packages/api/src/stream/__tests__/predecessorFence.spec.ts @@ -24,7 +24,100 @@ function createConditionalJob( ); } +/** An automatic continuation must never replace a parent turn that is still live. */ +function createWakeupJob(store: InMemoryJobStore, streamId: string) { + return store.createJob( + streamId, + 'owner-1', + streamId, + undefined, + { generationProtocolVersion: 2 }, + undefined, + undefined, + undefined, + undefined, + undefined, + 'wakeup-create-attempt', + undefined, + true, + ); +} + describe('generation predecessor create fence', () => { + test('in-memory admission refuses an active predecessor and admits a settled one', async () => { + const store = new InMemoryJobStore(); + const streamId = 'in-memory-active-predecessor-fence'; + try { + const parent = await store.createJob(streamId, 'owner-1', streamId, undefined, { + generationProtocolVersion: 2, + }); + + await expect(createWakeupJob(store, streamId)).rejects.toBeInstanceOf( + JobPredecessorMismatchError, + ); + /** The controller needs the live state to answer a finite PARENT_NOT_READY. */ + await expect(createWakeupJob(store, streamId)).rejects.toMatchObject({ + currentJob: { active: true, verified: true, status: 'running' }, + }); + + await expect( + store.transitionStatus(streamId, { + from: 'running', + to: 'requires_action', + expectCreatedAt: parent.createdAt, + }), + ).resolves.toBe(true); + await expect(createWakeupJob(store, streamId)).rejects.toMatchObject({ + currentJob: { active: true, status: 'requires_action' }, + }); + + await expect( + store.transitionStatus(streamId, { + from: 'requires_action', + to: 'aborted', + expectCreatedAt: parent.createdAt, + }), + ).resolves.toBe(true); + await store.updateJob(streamId, { terminalPersistencePending: true }, parent.createdAt); + await expect(createWakeupJob(store, streamId)).rejects.toMatchObject({ + currentJob: { active: true, status: 'aborted' }, + }); + await store.updateJob(streamId, { terminalPersistencePending: false }, parent.createdAt); + const wakeup = await createWakeupJob(store, streamId); + expect(wakeup.createdAt).toBeGreaterThanOrEqual(parent.createdAt); + } finally { + await store.destroy(); + } + }); + + test('in-memory admission accepts an absent predecessor', async () => { + const store = new InMemoryJobStore(); + const streamId = 'in-memory-absent-predecessor-fence'; + try { + const wakeup = await createWakeupJob(store, streamId); + expect(wakeup.createdAt).toEqual(expect.any(Number)); + } finally { + await store.destroy(); + } + }); + + test('in-memory ordinary turns still replace an active predecessor', async () => { + const store = new InMemoryJobStore(); + const streamId = 'in-memory-ordinary-replacement-unchanged'; + try { + const first = await store.createJob(streamId, 'owner-1', streamId, undefined, { + generationProtocolVersion: 2, + }); + /** Without the policy a user turn keeps replacing a running generation. */ + const second = await store.createJob(streamId, 'owner-1', streamId, undefined, { + generationProtocolVersion: 2, + }); + expect(second.createdAt).toBeGreaterThanOrEqual(first.createdAt); + } finally { + await store.destroy(); + } + }); + test('in-memory ordinary appends accept active states and reject retained terminal epochs', async () => { const store = new InMemoryJobStore(); const streamId = 'in-memory-terminal-append-fence'; diff --git a/packages/api/src/stream/__tests__/predecessorFence.stream_integration.spec.ts b/packages/api/src/stream/__tests__/predecessorFence.stream_integration.spec.ts index 0bf6c8033e8..d2854bfdd50 100644 --- a/packages/api/src/stream/__tests__/predecessorFence.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/predecessorFence.stream_integration.spec.ts @@ -26,6 +26,25 @@ function createConditionalJob( ); } +/** An automatic continuation must never replace a parent turn that is still live. */ +function createWakeupJob(store: RedisJobStore, streamId: string, attempt: string) { + return store.createJob( + streamId, + 'owner-1', + streamId, + undefined, + { generationProtocolVersion: 2 }, + undefined, + undefined, + undefined, + undefined, + undefined, + attempt, + undefined, + true, + ); +} + describe('Redis generation predecessor create fence', () => { const keyPrefix = `Predecessor-Fence-${process.pid}-${Date.now()}:`; let redis: RedisTestClient; @@ -50,6 +69,83 @@ describe('Redis generation predecessor create fence', () => { await redis.quit(); }); + test('admission refuses an active predecessor without mutating it, and admits a settled one', async () => { + const streamId = 'redis-active-predecessor-fence'; + const parent = await store.createJob(streamId, 'owner-1', streamId, undefined, { + generationProtocolVersion: 2, + }); + await expect( + store.appendChunk( + streamId, + { event: 'on_message_delta', data: { delta: 'parent still writing' } }, + parent.createdAt, + ), + ).resolves.toBe(true); + const chunksKey = `stream:{${streamId}}:chunks`; + const chunksBefore = await redis.xrange(chunksKey, '-', '+'); + + await expect(createWakeupJob(store, streamId, 'redis-wakeup-running')).rejects.toMatchObject({ + code: 'GENERATION_PREDECESSOR_MISMATCH', + currentJob: { createdAt: parent.createdAt, active: true, verified: true, status: 'running' }, + }); + /** The refused continuation must leave the live parent turn exactly as it was. */ + await expect(store.getJob(streamId)).resolves.toMatchObject({ + createdAt: parent.createdAt, + status: 'running', + }); + await expect(redis.xrange(chunksKey, '-', '+')).resolves.toEqual(chunksBefore); + + await expect( + store.transitionStatus(streamId, { + from: 'running', + to: 'requires_action', + expectCreatedAt: parent.createdAt, + }), + ).resolves.toBe(true); + await expect( + createWakeupJob(store, streamId, 'redis-wakeup-requires-action'), + ).rejects.toMatchObject({ + code: 'GENERATION_PREDECESSOR_MISMATCH', + currentJob: { active: true, status: 'requires_action' }, + }); + + await expect( + store.transitionStatus(streamId, { + from: 'requires_action', + to: 'aborted', + expectCreatedAt: parent.createdAt, + }), + ).resolves.toBe(true); + await store.updateJob(streamId, { terminalPersistencePending: true }, parent.createdAt); + await expect( + createWakeupJob(store, streamId, 'redis-wakeup-terminal-persistence'), + ).rejects.toMatchObject({ + currentJob: { active: true, status: 'aborted' }, + }); + await store.updateJob(streamId, { terminalPersistencePending: false }, parent.createdAt); + const wakeup = await createWakeupJob(store, streamId, 'redis-wakeup-settled'); + expect(wakeup.createdAt).toBeGreaterThanOrEqual(parent.createdAt); + }); + + test('admission accepts an absent predecessor', async () => { + const streamId = 'redis-absent-predecessor-fence'; + const wakeup = await createWakeupJob(store, streamId, 'redis-wakeup-absent'); + expect(wakeup.createdAt).toEqual(expect.any(Number)); + }); + + test('ordinary turns still replace an active predecessor', async () => { + const streamId = 'redis-ordinary-replacement-unchanged'; + const first = await store.createJob(streamId, 'owner-1', streamId, undefined, { + generationProtocolVersion: 2, + }); + /** Without the policy a user turn keeps replacing a running generation. */ + const second = await store.createJob(streamId, 'owner-1', streamId, undefined, { + generationProtocolVersion: 2, + }); + expect(second.createdAt).toBeGreaterThanOrEqual(first.createdAt); + await expect(store.getJob(streamId)).resolves.toMatchObject({ createdAt: second.createdAt }); + }); + test('mismatch returns the exact current generation without mutating durable state', async () => { const streamId = 'redis-predecessor-fence-no-mutation'; const predecessor = await store.createJob(streamId, 'owner-1', streamId, undefined, { diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index a1404915971..9b9653d00c9 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -264,6 +264,7 @@ export class InMemoryJobStore implements IJobStoreV2 { recoveredSteerPayload?: RecoveredSteerPayload, creationAttemptId?: string, expectedPredecessorCreatedAt?: number, + rejectActivePredecessor?: boolean, ): Promise { if (typeof userId !== 'string' || userId.length === 0) { throw new Error('Generation job requires a non-empty user id'); @@ -285,6 +286,9 @@ export class InMemoryJobStore implements IJobStoreV2 { ) { throw new Error('Invalid expected generation predecessor'); } + if (rejectActivePredecessor != null && typeof rejectActivePredecessor !== 'boolean') { + throw new Error('Invalid active generation predecessor policy'); + } const providerExecutionId = initialMetadata.providerExecutionId; if ( providerExecutionId != null && @@ -341,6 +345,20 @@ export class InMemoryJobStore implements IJobStoreV2 { const assertExpectedPredecessorCompatible = (): void => { const current = this.jobs.get(streamId); const currentCreatedAt = current?.createdAt ?? this.getRetainedGenerationEpoch(streamId); + if ( + rejectActivePredecessor === true && + (current?.status === 'running' || + current?.status === 'requires_action' || + current?.terminalPersistencePending === true) + ) { + throw new JobPredecessorMismatchError({ + createdAt: current.createdAt, + active: true, + verified: true, + status: current.status, + ...(current.conversationId !== undefined && { conversationId: current.conversationId }), + }); + } if ( expectedPredecessorCreatedAt == null || currentCreatedAt === expectedPredecessorCreatedAt diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index ec39812290b..b218d9d4acd 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -325,7 +325,7 @@ const REPLACEMENT_RECEIPT_ACK_LUA = * recoveredSteerPayloadJson | "", * generationProtocolVersion, * creationAttemptId | "", - * expectedPredecessorCreatedAt | "", + * expectedPredecessorCreatedAt | "", rejectActivePredecessor ("1" | "0"), * ...hsetPairs] * Returns: [previousUserId | "", previousTenantId | "", createdAt, "", * replacedCreatedAt | "", replacedStatus | "", replacedConversationId | "", @@ -354,6 +354,7 @@ const JOB_CREATE_LUA = 'local replacedProviderAbortReady = redis.call("HGET", KEYS[1], "providerAbortReady") ' + 'local replacedProviderExecutionId = redis.call("HGET", KEYS[1], "providerExecutionId") ' + 'local replacedProviderDrained = redis.call("HGET", KEYS[1], "providerDrained") ' + + 'local replacedTerminalPersistencePending = redis.call("HGET", KEYS[1], "terminalPersistencePending") ' + 'local replacedProtocol = redis.call("HGET", KEYS[1], "generationProtocolVersion") ' + 'local MAX_SAFE_EPOCH = 9007199254740991 ' + 'local function isSafeEpoch(value) return type(value) == "number" and value >= 0 ' + @@ -375,10 +376,14 @@ const JOB_CREATE_LUA = 'local observedCreatedAt = replacedCreatedAt local observedStatus = replacedStatus ' + 'local observedConversationId = replacedConversationId ' + 'local observedActive = previousJobExists == 1 and ' + - '(replacedStatus == "running" or replacedStatus == "requires_action") ' + + '(replacedStatus == "running" or replacedStatus == "requires_action" ' + + 'or replacedTerminalPersistencePending == "1") ' + 'if retainedEpoch and (not previousCreatedAt or retainedEpoch > previousCreatedAt) then ' + 'previousCreatedAt = retainedEpoch observedCreatedAt = retainedEpochRaw ' + 'observedStatus = nil observedConversationId = nil observedActive = false end ' + + 'if ARGV[13] == "1" and observedActive then ' + + 'return { previousUserId or "", previousTenantId or "", "0", "predecessor_mismatch", ' + + 'observedCreatedAt, observedStatus or "", observedConversationId or "", "1", "1" } end ' + 'if ARGV[12] ~= "" and (not observedCreatedAt or observedCreatedAt ~= ARGV[12]) then ' + 'return { previousUserId or "", previousTenantId or "", "0", "predecessor_mismatch", ' + 'observedCreatedAt or ARGV[12], observedStatus or "", observedConversationId or "", ' + @@ -526,7 +531,7 @@ const JOB_CREATE_LUA = 'local ttl = tonumber(ARGV[1]) ' + 'local generationEpochGraceTtl = tonumber(ARGV[3]) ' + 'local hset = {} ' + - 'for i = 13, #ARGV do hset[#hset + 1] = ARGV[i] end ' + + 'for i = 14, #ARGV do hset[#hset + 1] = ARGV[i] end ' + 'redis.call("HSET", KEYS[1], unpack(hset)) ' + 'redis.call("HSET", KEYS[1], "createdAt", tostring(createdAt)) ' + 'if ARGV[11] ~= "" then redis.call("HSET", KEYS[1], "__creationAttemptId", ARGV[11]) end ' + @@ -1769,6 +1774,7 @@ export class RedisJobStore implements IJobStoreV2 { recoveredSteerPayload?: RecoveredSteerPayload, creationAttemptId?: string, expectedPredecessorCreatedAt?: number, + rejectActivePredecessor?: boolean, ): Promise { if (typeof userId !== 'string' || userId.length === 0) { throw new Error('Generation job requires a non-empty user id'); @@ -1790,6 +1796,9 @@ export class RedisJobStore implements IJobStoreV2 { ) { throw new Error('Invalid expected generation predecessor'); } + if (rejectActivePredecessor != null && typeof rejectActivePredecessor !== 'boolean') { + throw new Error('Invalid active generation predecessor policy'); + } const providerExecutionId = initialMetadata.providerExecutionId; if ( providerExecutionId != null && @@ -1868,6 +1877,7 @@ export class RedisJobStore implements IJobStoreV2 { String(job.generationProtocolVersion), creationAttemptId ?? '', expectedPredecessorCreatedAt == null ? '' : String(expectedPredecessorCreatedAt), + rejectActivePredecessor === true ? '1' : '0', ...hsetPairs, ); if (Array.isArray(previousOwner) && previousOwner[3] === 'claim_lost') { diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index de13becfd26..10143c62cde 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -741,6 +741,7 @@ export interface IJobStoreV2 extends IJobStore { recoveredSteerPayload?: RecoveredSteerPayload, creationAttemptId?: string, expectedPredecessorCreatedAt?: number, + rejectActivePredecessor?: boolean, ): Promise; /** Remove transaction-time predecessor receipts after their handoff was diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index add99d0fc63..b94dd885286 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -29,6 +29,9 @@ let updateMessageText: ReturnType['updateMessageTex let deleteMessagesSince: ReturnType['deleteMessagesSince']; let recordMessage: ReturnType['recordMessage']; let claimSubagentTaskResult: ReturnType['claimSubagentTaskResult']; +let releaseSubagentTaskResultClaim: ReturnType< + typeof createMessageMethods +>['releaseSubagentTaskResultClaim']; beforeAll(async () => { mongoServer = await MongoMemoryServer.create(); @@ -49,6 +52,7 @@ beforeAll(async () => { deleteMessagesSince = methods.deleteMessagesSince; recordMessage = methods.recordMessage; claimSubagentTaskResult = methods.claimSubagentTaskResult; + releaseSubagentTaskResultClaim = methods.releaseSubagentTaskResultClaim; await mongoose.connect(mongoUri); }); @@ -1569,6 +1573,7 @@ describe('Message Operations', () => { userId: 'user123', conversationId, taskId, + kind: 'manual', claimId: 'poll-1', }); expect(first.status).toBe('acquired'); @@ -1579,14 +1584,136 @@ describe('Message Operations', () => { userId: 'user123', conversationId, taskId, + kind: 'manual', claimId: 'poll-1', }); expect(retried.status).toBe('acquired'); /** Another invocation is told it was collected instead of handed a copy. */ await expect( - claimSubagentTaskResult({ userId: 'user123', conversationId, taskId, claimId: 'poll-2' }), - ).resolves.toEqual({ status: 'claimed' }); + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + kind: 'manual', + claimId: 'poll-2', + }), + ).resolves.toMatchObject({ status: 'claimed' }); + }); + + it('elects either a manual poll or one idempotent automatic wakeup', async () => { + const manualTaskId = uuidv4(); + const wakeupTaskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(manualTaskId, conversationId, 'completed'); + await terminalResult(wakeupTaskId, conversationId, 'completed'); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: manualTaskId, + kind: 'manual', + claimId: 'poll-1', + }), + ).resolves.toMatchObject({ status: 'acquired' }); + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: manualTaskId, + kind: 'wakeup', + claimId: 'delivery-1', + }), + ).resolves.toMatchObject({ status: 'claimed' }); + + const wakeupClaim = { + userId: 'user123', + conversationId, + taskId: wakeupTaskId, + kind: 'wakeup' as const, + claimId: 'delivery-2', + }; + await expect(claimSubagentTaskResult(wakeupClaim)).resolves.toMatchObject({ + status: 'acquired', + }); + await expect(claimSubagentTaskResult(wakeupClaim)).resolves.toMatchObject({ + status: 'acquired', + }); + await expect( + claimSubagentTaskResult({ ...wakeupClaim, claimId: 'delivery-3' }), + ).resolves.toMatchObject({ status: 'claimed' }); + await expect( + claimSubagentTaskResult({ ...wakeupClaim, kind: 'manual', claimId: 'poll-2' }), + ).resolves.toMatchObject({ status: 'claimed' }); + }); + + it('preserves and upgrades retries of legacy manual claims without a kind', async () => { + const taskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(taskId, conversationId, 'completed'); + await claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + kind: 'manual', + claimId: 'legacy-poll', + }); + await Message.collection.updateOne( + { user: 'user123', conversationId, messageId: `${taskId}:assistant` }, + { $unset: { 'subagentTask.resultClaim.kind': '' } }, + ); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + kind: 'manual', + claimId: 'legacy-poll', + }), + ).resolves.toMatchObject({ + status: 'acquired', + message: { subagentTask: { resultClaim: { kind: 'manual', claimId: 'legacy-poll' } } }, + }); + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + kind: 'wakeup', + claimId: 'legacy-poll', + }), + ).resolves.toMatchObject({ status: 'claimed' }); + }); + + it('releases only the exact rejected wakeup so manual collection can take over', async () => { + const taskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(taskId, conversationId, 'completed'); + const wakeup = { + userId: 'user123', + conversationId, + taskId, + kind: 'wakeup' as const, + claimId: 'delivery-1', + }; + await expect(claimSubagentTaskResult(wakeup)).resolves.toMatchObject({ + status: 'acquired', + }); + await expect( + releaseSubagentTaskResultClaim({ ...wakeup, claimId: 'another-delivery' }), + ).resolves.toBe(false); + await expect(releaseSubagentTaskResultClaim(wakeup)).resolves.toBe(true); + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + kind: 'manual', + claimId: 'poll-after-rejection', + }), + ).resolves.toMatchObject({ status: 'acquired' }); }); it('reports a result that is missing or still running as not found', async () => { @@ -1599,6 +1726,7 @@ describe('Message Operations', () => { userId: 'user123', conversationId, taskId: runningTaskId, + kind: 'manual', claimId: 'poll-1', }), ).resolves.toEqual({ status: 'not_found' }); @@ -1608,6 +1736,7 @@ describe('Message Operations', () => { userId: 'user123', conversationId, taskId: uuidv4(), + kind: 'manual', claimId: 'poll-1', }), ).resolves.toEqual({ status: 'not_found' }); @@ -1623,6 +1752,7 @@ describe('Message Operations', () => { userId: 'other-user', conversationId, taskId, + kind: 'manual', claimId: 'poll-1', }), ).resolves.toEqual({ status: 'not_found' }); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 6f427bb8207..500fb05ea60 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -50,7 +50,7 @@ interface MessageQueryOptions { export type SubagentTaskResultClaim = | { status: 'not_found' } - | { status: 'claimed' } + | { status: 'claimed'; message: IMessage } | { status: 'acquired'; message: IMessage }; export interface MessageMethods { @@ -91,8 +91,16 @@ export interface MessageMethods { userId: string; conversationId: string; taskId: string; + kind: 'manual' | 'wakeup'; claimId: string; }): Promise; + releaseSubagentTaskResultClaim(params: { + userId: string; + conversationId: string; + taskId: string; + kind: 'manual' | 'wakeup'; + claimId: string; + }): Promise; deleteMessagesSince( userId: string, params: { messageId: string; conversationId: string }, @@ -529,21 +537,19 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } } - /** - * Assigns one durable terminal child result to the polling invocation that collects - * it. The same invocation may re-acquire, so a poll whose response was lost recovers - * the result it never received; a different invocation is told it was already - * collected rather than handed a second copy. - */ + /** Atomically assigns one durable terminal child result to either its + * explicit poller or one idempotent automatic wakeup delivery. */ async function claimSubagentTaskResult({ userId, conversationId, taskId, + kind, claimId, }: { userId: string; conversationId: string; taskId: string; + kind: 'manual' | 'wakeup'; claimId: string; }): Promise { if ( @@ -551,38 +557,112 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa taskId.length > 256 || conversationId.length === 0 || conversationId.length > 256 || + (kind !== 'manual' && kind !== 'wakeup') || claimId.length === 0 || claimId.length > 128 ) { throw new TypeError('Invalid subagent task result claim'); } const Message = mongoose.models.Message as Model; - const filter = { - user: userId, - conversationId, - messageId: `${taskId}:assistant`, - 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + const messageId = `${taskId}:assistant`; + const terminal = ['completed', 'error', 'cancelled']; + const claim = { + kind, + claimId, + claimedAt: new Date(), + }; + const claimable = { + $or: [ + { 'subagentTask.resultClaim': { $exists: false } }, + { + 'subagentTask.resultClaim.kind': kind, + 'subagentTask.resultClaim.claimId': claimId, + }, + ...(kind === 'manual' + ? [ + { + 'subagentTask.resultClaim.kind': { $exists: false }, + 'subagentTask.resultClaim.claimId': claimId, + }, + ] + : []), + ], + }; + const projection = { + messageId: 1, + conversationId: 1, + parentMessageId: 1, + sender: 1, + text: 1, + error: 1, + createdAt: 1, + updatedAt: 1, + subagentTask: 1, }; const acquired = await Message.findOneAndUpdate( { - ...filter, - $or: [ - { 'subagentTask.resultClaim': { $exists: false } }, - { 'subagentTask.resultClaim.claimId': claimId }, - ], - }, - { $set: { 'subagentTask.resultClaim': { claimId, claimedAt: new Date() } } }, - { - new: true, - timestamps: false, - projection: { messageId: 1, conversationId: 1, text: 1, subagentTask: 1 }, + user: userId, + conversationId, + messageId, + 'subagentTask.status': { $in: terminal }, + ...claimable, }, + { $set: { 'subagentTask.resultClaim': claim } }, + { new: true, projection }, ).lean(); if (acquired != null) { return { status: 'acquired', message: acquired }; } - const existing = await Message.exists(filter); - return existing == null ? { status: 'not_found' } : { status: 'claimed' }; + const existing = await Message.findOne({ + user: userId, + conversationId, + messageId, + 'subagentTask.status': { $in: terminal }, + }) + .select(projection) + .lean(); + return existing == null ? { status: 'not_found' } : { status: 'claimed', message: existing }; + } + + /** Releases only the exact consumer assignment. This is used when a + * pre-admission automatic continuation is definitively rejected, allowing a + * later manual poll (or the same delivery retry) to claim the durable result. */ + async function releaseSubagentTaskResultClaim({ + userId, + conversationId, + taskId, + kind, + claimId, + }: { + userId: string; + conversationId: string; + taskId: string; + kind: 'manual' | 'wakeup'; + claimId: string; + }): Promise { + if ( + taskId.length === 0 || + taskId.length > 256 || + conversationId.length === 0 || + conversationId.length > 256 || + (kind !== 'manual' && kind !== 'wakeup') || + claimId.length === 0 || + claimId.length > 128 + ) { + throw new TypeError('Invalid subagent task result claim release'); + } + const Message = mongoose.models.Message as Model; + const result = await Message.updateOne( + { + user: userId, + conversationId, + messageId: `${taskId}:assistant`, + 'subagentTask.resultClaim.kind': kind, + 'subagentTask.resultClaim.claimId': claimId, + }, + { $unset: { 'subagentTask.resultClaim': 1 } }, + ); + return result.modifiedCount === 1; } /** @@ -723,6 +803,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateToolCallResult, updateMessage, claimSubagentTaskResult, + releaseSubagentTaskResultClaim, deleteMessagesSince, getMessages, getMessage, diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index 8282cbf0597..ad1f59db9e6 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -138,6 +138,7 @@ const messageSchema: Schema = new Schema( subagentTask: { type: { attemptKey: { type: String, required: true }, + parentRunId: { type: String }, requestFingerprint: { type: String }, status: { type: String, @@ -146,6 +147,7 @@ const messageSchema: Schema = new Schema( }, resultClaim: { type: { + kind: { type: String, enum: ['manual', 'wakeup'], required: true }, claimId: { type: String, required: true }, claimedAt: { type: Date, required: true }, }, diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index 541a5f1d1f2..bd3da444ee5 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -51,10 +51,12 @@ export interface IMessage extends Document { /** Server-private durable idempotency marker for one detached subagent turn. */ subagentTask?: { attemptKey: string; + /** Parent response that initiated this exact child task. */ + parentRunId?: string; requestFingerprint?: string; status: 'running' | 'completed' | 'error' | 'cancelled'; - /** Records which polling invocation collected this terminal result. */ resultClaim?: { + kind: 'manual' | 'wakeup'; claimId: string; claimedAt: Date; }; From b91691937e0e39933325a925027192fd29681f47 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 14:21:20 -0400 Subject: [PATCH 09/15] =?UTF-8?q?=F0=9F=99=8B=20fix:=20Free=20the=20Compos?= =?UTF-8?q?er=20When=20a=20Question=20Pause=20Collapses=20(#15011)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🙋 fix: Free the Composer When a Question Pause Collapses Collapsing a live `ask_user_question` left the user with nothing to do. A batch of questions disables the composer, the send button, and the stop button for as long as the pause is active — and `collapse` deliberately keeps it active, while hiding the popover that carried the only dismiss. After the chevron there was no way to type, send, or stop the run short of reloading the page. Split the composer's role out of `active`: `composerAnswers` (a single question, answered IN the composer) and `composerLocked` (a batch, answered in its own card — and only while the popover is up). Collapsing a batch now hands the composer back to the thread; the stop button follows `composerAnswers`, so a paused run stays stoppable. Both collapsed cards also carry the popover's ×, so dismiss survives the handover, and `submitText` declines a batch's composer text instead of claiming it — the old `return true` reported success and dropped whatever was staged when the pause began. Contrast, per feedback that the questions were hard to read: the answer options, the answer textarea, and the digit chips all drew their edge from `border-light`, which measures 1.20:1 against the panel (WCAG 1.4.11 wants 3:1 for a UI component boundary) — a column of choices read as flat text. Adds a `choice` Button variant carrying its own fill and a `border-xheavy` edge (5.49:1 dark / 6.54:1 light), at `font-normal` so the question above stays the heading, and replaces the single-question popover's hardcoded `bg-white`/`dark:bg-gray-700` with the semantic surface role it should have been using. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01UPtmUb6VLBhhxXkS3PfV6r * 🧹 refactor: Render Popover Answers Through the Choice Variant The popover's option rows re-stated the shared `choice` variant's border, fill, weight, and hover on a raw ` + ); })}
diff --git a/client/src/components/Chat/Input/ChatForm.tsx b/client/src/components/Chat/Input/ChatForm.tsx index 3c52b13d0f1..91b56393a2f 100644 --- a/client/src/components/Chat/Input/ChatForm.tsx +++ b/client/src/components/Chat/Input/ChatForm.tsx @@ -179,6 +179,10 @@ const ChatForm = memo(function ChatForm({ const answerPlaceholder = answerMode.batchMode ? localize('com_ui_answer_questions_above') : (answerMode.otherLabel ?? localize('com_ui_something_else')); + /** The composer is not a plain chat composer: it either IS this pause's + * answer box, or is locked behind the batch card that owns the answer. A + * collapsed batch is neither — it hands the composer back to the thread. */ + const composerReserved = answerMode.composerAnswers || answerMode.composerLocked; useAutoSave({ index, @@ -269,7 +273,7 @@ const ChatForm = memo(function ChatForm({ conversationId, conversation, isSubmitting, - answerModeActive: answerMode.active, + answerModeActive: composerReserved, files, setFiles, filesLoading, @@ -287,8 +291,8 @@ const ChatForm = memo(function ChatForm({ const liveFilesRef = useRef(files); liveFilesRef.current = files; /** Same reason: the run can pause on `ask_user_question` mid-reclaim. */ - const liveAnswerModeRef = useRef(answerMode.active); - liveAnswerModeRef.current = answerMode.active; + const liveAnswerModeRef = useRef(composerReserved); + liveAnswerModeRef.current = composerReserved; /** A reclaim can resolve after this form unmounts (left the route, closed the * pane). Its refs still hold the origin chat, so the restore would pass its * checks and write into a dead form — reporting success and making the caller @@ -393,11 +397,11 @@ const ChatForm = memo(function ChatForm({ setIsScrollable, disabled: disableInputs, // The composer IS the free-form answer box while a question pause is live. - placeholder: answerMode.active ? answerPlaceholder : placeholder, + placeholder: composerReserved ? answerPlaceholder : placeholder, // Enter stays live during a run when it can steer/queue instead of send. allowSubmitWhileGenerating: steering.duringRunActive, onDuringRunModifier: steering.duringRunActive ? handleDuringRunModifier : undefined, - answerModeActive: answerMode.active && !answerMode.batchMode, + answerModeActive: answerMode.composerAnswers, }); useQueryParams({ textAreaRef }); @@ -406,7 +410,7 @@ const ChatForm = memo(function ChatForm({ * hands the composer text straight to the paused run, which answers with * values and cannot consume files, so an empty draft must stay unsubmittable * there rather than enabling a button whose submit is silently dropped. */ - const submittableFileCount = answerMode.active ? 0 : files.size; + const submittableFileCount = composerReserved ? 0 : files.size; const { ref, ...registerProps } = methods.register('text', { required: submittableFileCount === 0, @@ -485,7 +489,8 @@ const ChatForm = memo(function ChatForm({ onSubmit={methods.handleSubmit((data) => { // Answer mode: composer text answers the paused run instead of // starting a new turn (submitText resets the composer itself). - // Dismissing the popover restores normal sends. + // Dismissing the popover — or collapsing a batch, which answers in its + // own card — restores normal sends. if (answerMode.active && answerMode.submitText(data.text)) { return; } @@ -608,11 +613,7 @@ const ChatForm = memo(function ChatForm({ textAreaRef as React.MutableRefObject ).current = e; }} - disabled={ - disableInputs || - isNotAppendable || - (answerMode.active && answerMode.batchMode) - } + disabled={disableInputs || isNotAppendable || answerMode.composerLocked} onPaste={handlePaste} onKeyDown={(e) => { // Answer mode consumes option-navigation keys from the @@ -705,7 +706,7 @@ const ChatForm = memo(function ChatForm({
{isSubmitting && (showStopButton || steering.duringRunActive) && - !answerMode.active + !answerMode.composerAnswers ? duringRunSlot : endpoint && ( )} diff --git a/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx b/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx index 8817090a2e7..a4d3ad13832 100644 --- a/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx +++ b/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx @@ -1,5 +1,5 @@ import { useContext, useMemo, useState } from 'react'; -import { ChevronUp, TriangleAlert } from 'lucide-react'; +import { ChevronUp, TriangleAlert, X } from 'lucide-react'; import { Button, TextareaAutosize } from '@librechat/client'; import type { Agents } from 'librechat-data-provider'; import { useApprovalContext, useAskSubmitStatus, useResumeSubmit } from './ApprovalContext'; @@ -40,6 +40,7 @@ export default function AskUserQuestion({ questions={questions} className="my-2 max-h-[70vh] w-full rounded-lg border border-border-light bg-surface-secondary" onExpand={answerMode.collapsed && isLivePause ? answerMode.expand : undefined} + onDismiss={answerMode.collapsed && isLivePause ? answerMode.dismiss : undefined} /> ); } @@ -68,7 +69,7 @@ function AskUserQuestionSingle({ * chevron re-expands it) or dismissed (and in contexts without a * ChatContext, where the popover can't exist). */ - const { popoverVisible, collapsed, expand, liveAsk } = answerMode; + const { popoverVisible, collapsed, expand, dismiss, liveAsk } = answerMode; const isLivePause = liveAsk?.actionId === actionId; /** Same fold as the popover: a model-supplied catch-all "Other" option @@ -144,14 +145,26 @@ function AskUserQuestionSingle({ {question.question}

{collapsed && isLivePause && ( - +
+ + {/** Mirrors the popover's ×: the collapsed card is the only chrome + * left, so the way out of answer mode has to live here too. */} + +
)}
{question.description != null && question.description.length > 0 && ( @@ -166,7 +179,7 @@ function AskUserQuestionSingle({ + {(onExpand != null || onDismiss != null) && ( +
+ {onExpand != null && ( + + )} + {/** The collapsed card is the ONLY surface left for a collapsed + * batch, so it has to carry the popover's dismiss too — without it + * the pause can only be answered or skipped. */} + {onDismiss != null && ( + + )}
)} {stepped && ( @@ -175,7 +192,7 @@ export default function AskUserQuestions({ {question.question}

{question.description != null && question.description.length > 0 && ( -

+

{question.description}

)} @@ -188,7 +205,7 @@ export default function AskUserQuestions({ key={option.value} type="button" size="sm" - variant={isSelected ? 'submit' : 'outline'} + variant={isSelected ? 'submit' : 'choice'} role={question.multiSelect === true ? 'checkbox' : undefined} aria-checked={question.multiSelect === true ? isSelected : undefined} aria-pressed={question.multiSelect === true ? undefined : isSelected} @@ -212,7 +229,7 @@ export default function AskUserQuestions({ minRows={1} maxRows={6} placeholder={otherLabel ?? localize('com_ui_your_answer')} - className="mt-2 w-full resize-none rounded-md border border-border-light bg-surface-primary p-2 text-sm text-text-primary" + className="mt-2 w-full resize-none rounded-md border border-border-xheavy bg-surface-primary p-2 text-sm text-text-primary" aria-label={`${question.question} ${localize('com_ui_your_answer')}`} /> diff --git a/client/src/components/Chat/Messages/Content/__tests__/AskUserQuestionCollapse.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/AskUserQuestionCollapse.test.tsx new file mode 100644 index 00000000000..eedb7287b35 --- /dev/null +++ b/client/src/components/Chat/Messages/Content/__tests__/AskUserQuestionCollapse.test.tsx @@ -0,0 +1,141 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import { fireEvent, render, screen } from '@testing-library/react'; +import type { Agents } from 'librechat-data-provider'; +import useAskAnswerMode from '~/hooks/Input/useAskAnswerMode'; +import { ChatContext } from '~/Providers/ChatContext'; +import AskUserQuestion from '../AskUserQuestion'; + +/** + * Collapsing a live pause used to strand the user: the popover carried the only + * dismiss, and a batch kept the composer disabled for as long as the pause was + * active — so after the chevron there was nothing to type in, nothing to send + * with, no stop button, and no × anywhere. Reloading the page was the only way + * out. These cover both halves of the handover across the REAL recoil atoms the + * popover and the chat card share. + */ + +const mockBatch: Agents.AskUserQuestionBatchItem[] = [ + { id: 'scope', header: 'Dashboard scope', question: 'North star or full funnel?' }, + { id: 'window', question: 'Which time window?' }, +]; + +const mockSingle = { + question: 'Which environment?', + options: [{ label: 'Staging', value: 'staging' }], +} as Agents.AskUserQuestionRequest; + +let mockLiveAsk: { + actionId: string; + question: Agents.AskUserQuestionRequest; + questions?: Agents.AskUserQuestionBatchItem[]; + messageId: string; +} = { + actionId: 'act-1', + question: mockSingle, + questions: mockBatch, + messageId: 'message-1', +}; + +jest.mock('~/data-provider', () => ({ + useGetMessagesByConvoId: () => ({ data: mockLiveAsk }), +})); +jest.mock('~/components/Chat/Messages/Content/ApprovalContext', () => ({ + useApprovalContext: () => ({ getAskAnswerDraft: () => '', setAskAnswerDraft: jest.fn() }), + useAskSubmitStatus: () => ({ getAskStatus: () => 'idle' }), + useResumeSubmit: () => ({ submitAskAnswer: jest.fn() }), +})); +jest.mock('~/Providers', () => ({ useOptionalChatFormContext: () => null })); + +/** Stands in for the composer: publishes the flags ChatForm gates on. */ +function ComposerProbe() { + const ask = useAskAnswerMode('conversation-1'); + return ( +
+ {String(ask.composerLocked)} + {String(ask.composerAnswers)} + {String(ask.popoverVisible)} + {String(ask.active)} +
+ ); +} + +const renderPause = () => + render( + + + + + + , + ); + +describe('collapsing a live ask_user_question', () => { + describe('batched questions', () => { + beforeEach(() => { + mockLiveAsk = { + actionId: 'act-1', + question: mockSingle, + questions: mockBatch, + messageId: 'message-1', + }; + }); + + it('locks the composer only while the popover is up', () => { + renderPause(); + expect(screen.getByTestId('composer-locked').textContent).toBe('true'); + /** The popover owns the question, so the card stays out of the way. */ + expect(screen.queryByText('North star or full funnel?')).toBeNull(); + + fireEvent.click(screen.getByTestId('collapse-from-popover')); + + expect(screen.getByTestId('popover-visible').textContent).toBe('false'); + /** The pause is still live — the card has it now... */ + expect(screen.getByTestId('active').textContent).toBe('true'); + expect(screen.getByText('North star or full funnel?')).toBeInTheDocument(); + /** ...and the composer is a composer again. */ + expect(screen.getByTestId('composer-locked').textContent).toBe('false'); + expect(screen.getByTestId('composer-answers').textContent).toBe('false'); + }); + + it('gives the collapsed card the popover’s dismiss', () => { + renderPause(); + fireEvent.click(screen.getByTestId('collapse-from-popover')); + + expect(screen.getByLabelText('Expand')).toBeInTheDocument(); + fireEvent.click(screen.getByLabelText('Close')); + + /** Dismiss exits answer mode entirely, exactly as the popover’s × did. */ + expect(screen.getByTestId('active').textContent).toBe('false'); + }); + }); + + describe('a single question', () => { + beforeEach(() => { + mockLiveAsk = { actionId: 'act-1', question: mockSingle, messageId: 'message-1' }; + }); + + it('keeps the composer as the answer box, and still offers a dismiss', () => { + renderPause(); + expect(screen.getByTestId('composer-answers').textContent).toBe('true'); + expect(screen.getByTestId('composer-locked').textContent).toBe('false'); + + fireEvent.click(screen.getByTestId('collapse-from-popover')); + + /** A single question is answered IN the composer, so it stays the answer + * box past collapse — the card is only the display handing over. */ + expect(screen.getByTestId('composer-answers').textContent).toBe('true'); + expect(screen.getByText('Which environment?')).toBeInTheDocument(); + + fireEvent.click(screen.getByLabelText('Close')); + + expect(screen.getByTestId('active').textContent).toBe('false'); + expect(screen.getByTestId('composer-answers').textContent).toBe('false'); + }); + }); +}); diff --git a/client/src/hooks/Input/useAskAnswerMode.spec.ts b/client/src/hooks/Input/useAskAnswerMode.spec.ts index d2ffce31f58..691275ea12e 100644 --- a/client/src/hooks/Input/useAskAnswerMode.spec.ts +++ b/client/src/hooks/Input/useAskAnswerMode.spec.ts @@ -75,7 +75,7 @@ describe('useAskAnswerMode', () => { expect(result.current.popoverVisible).toBe(false); }); - it('keeps batch mode active while reserving answers for the bounded form', () => { + it('locks the composer for a batch, and hands it back the moment it collapses', () => { mockUseGetMessages.mockReturnValue({ data: { ...liveAsk, @@ -92,7 +92,12 @@ describe('useAskAnswerMode', () => { expect(result.current.batchMode).toBe(true); expect(result.current.options).toEqual([]); expect(result.current.draftId).toBeNull(); - expect(result.current.submitText('must stay out of the normal send path')).toBe(true); + /** The bounded form owns the answer, so the composer never speaks for it. */ + expect(result.current.composerAnswers).toBe(false); + expect(result.current.composerLocked).toBe(true); + /** Text is declined, not claimed: claiming it dropped whatever was staged + * when the pause began. */ + expect(result.current.submitText('must stay out of the normal send path')).toBe(false); expect(mockSubmitAskAnswer).not.toHaveBeenCalled(); }); diff --git a/client/src/hooks/Input/useAskAnswerMode.ts b/client/src/hooks/Input/useAskAnswerMode.ts index 436b9433aaf..a78b3afb6e9 100644 --- a/client/src/hooks/Input/useAskAnswerMode.ts +++ b/client/src/hooks/Input/useAskAnswerMode.ts @@ -52,9 +52,10 @@ const askAnswerCheckedAtom = atom({ * resolves. * * Two ways out short of answering: `collapse` hides the popover chrome but - * KEEPS answer mode live (the question renders in the chat card; the composer - * still answers), while the × `dismiss` exits answer mode entirely. `Skip` - * resumes the run with a canned decline notice. + * KEEPS the pause live (the question renders in the chat card, which still + * answers it), while the × `dismiss` exits answer mode entirely. `Skip` + * resumes the run with a canned decline notice. Collapsing always returns the + * composer to normal chat — see `composerLocked`. * * `handleComposerKeyDown` only steers selection from the EMPTY composer and * reports whether it consumed the key. @@ -128,10 +129,22 @@ export default function useAskAnswerMode(conversationId?: string | null) { */ const active = liveAsk != null && !dismissed && status !== 'expired' && status !== 'submitted'; const collapsed = active && collapsedIds.includes(liveAsk.actionId); - /** The popover renders only while expanded; collapse keeps `active` (and the - * composer's answer role) but hands the question display to the chat card. */ + /** The popover renders only while expanded; collapse keeps `active` but + * hands the question display to the chat card. */ const popoverVisible = active && !collapsed; const batchMode = (liveAsk?.questions?.length ?? 0) > 0; + /** + * Which role the composer plays for this pause. A single question is + * answered IN the composer, so it stays live for as long as the pause does. + * A batch is answered in its own card, so the composer has nothing to + * contribute and locks — but ONLY while the popover is up. Collapsing has to + * hand the composer back, because every other way out is gone once the + * popover closes: the stop button hides behind `composerAnswers` and Escape + * would reach a disabled textarea. Staying locked past collapse left no way + * to type, send, or stop the run short of reloading the page. + */ + const composerAnswers = active && !batchMode; + const composerLocked = popoverVisible && batchMode; const multiSelect = !batchMode && liveAsk != null && liveAsk.question.multiSelect === true; /** Answer-phase draft key: handed to useAutoSave so the composer drafts * under the question's own key while answer mode is live, leaving the @@ -291,16 +304,20 @@ export default function useAskAnswerMode(conversationId?: string | null) { [multiSelect, checkedValues, selected, submitValues, submitOption], ); - /** Composer text answers the question directly; true when consumed. On a - * multi-select question any checked options ride along with the text. */ + /** + * Composer text answers the question directly; true when consumed. On a + * multi-select question any checked options ride along with the text. + * + * A batch answers in its card, so the composer's text is none of its + * business: report it UNconsumed and let the normal send/steer path have it. + * Claiming it (the old `return true`) silently swallowed whatever was staged + * when the pause began — the submit reported success and dropped the words. + */ const submitText = useCallback( (text: string): boolean => { - if (!active || !liveAsk) { + if (!active || !liveAsk || batchMode) { return false; } - if (batchMode) { - return true; - } const trimmed = text.trim(); if (trimmed.length > 0) { submitValues(multiSelect ? [...checkedValues(), trimmed] : [trimmed], true); @@ -353,8 +370,11 @@ export default function useAskAnswerMode(conversationId?: string | null) { const composerText = e.currentTarget.value; if (composerText.trim().length > 0) { // The composer IS the free-form answer box: Enter submits the typed - // text (before useTextarea's submitting-lock can swallow it). - if (e.key === 'Enter' && !e.shiftKey) { + // text (before useTextarea's submitting-lock can swallow it). Not for + // a batch, which answers in its card — its Enter belongs to the normal + // send path, so leave the event untouched rather than preventDefault + // an event we are about to decline. + if (e.key === 'Enter' && !e.shiftKey && !batchMode) { e.preventDefault(); return submitText(composerText); } @@ -409,6 +429,7 @@ export default function useAskAnswerMode(conversationId?: string | null) { active, options, selected, + batchMode, multiSelect, popoverVisible, canSubmit, @@ -460,6 +481,8 @@ export default function useAskAnswerMode(conversationId?: string | null) { collapse, expand, popoverVisible, + composerAnswers, + composerLocked, multiSelect, locked, selected, diff --git a/packages/client/src/components/Button.tsx b/packages/client/src/components/Button.tsx index 1998ee8e26e..abde7d06892 100644 --- a/packages/client/src/components/Button.tsx +++ b/packages/client/src/components/Button.tsx @@ -11,6 +11,7 @@ type ButtonVariantOptions = | 'link' | 'submit' | 'outline' + | 'choice' | 'subtle' | 'destructive' | 'secondary' @@ -35,6 +36,16 @@ const buttonVariantRecipe = cva( 'bg-surface-destructive text-text-on-status hover:bg-surface-destructive-hover', outline: 'text-text-primary border border-border-light bg-transparent hover:bg-surface-hover hover:text-text-primary', + /** + * A selectable answer inside a question card. `outline` is wrong here: + * its `border-light` edge measures ~1.2:1 against the panel these sit + * on, well under WCAG 1.4.11's 3:1 for a UI component boundary, so a + * column of choices reads as flat text rather than as controls. Carries + * its own fill so the answers are a different colour from the prompt, + * and drops to `font-normal` so the question above stays the heading. + */ + choice: + 'border border-border-xheavy bg-surface-tertiary font-normal text-text-primary hover:bg-surface-hover hover:text-text-primary', subtle: 'border border-border-light bg-transparent text-text-primary hover:bg-surface-secondary focus-visible:ring-text-primary focus-visible:ring-offset-0', secondary: 'bg-surface-secondary text-text-primary hover:bg-surface-hover', From 7c71d6dc1a582801dafe4e40bfc68b37cd081cd1 Mon Sep 17 00:00:00 2001 From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:21:39 -0700 Subject: [PATCH 10/15] =?UTF-8?q?=F0=9F=9B=82=20fix:=20Preserve=20OpenID?= =?UTF-8?q?=20Re-Auth=20Errors=20Through=20MCP=20Tool=20Error=20Classifica?= =?UTF-8?q?tion=20(#15010)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MCP tool-call catch block classifies errors by message substring, so an OpenIDReauthRequiredError raised during header resolution was rewritten into an MCP OAuth configuration prompt and its class identity discarded. The typed error now passes through ahead of the heuristic, so the actionable re-authentication message reaches the caller intact. --- api/server/services/MCP.js | 6 +++++ api/server/services/MCP.spec.js | 48 +++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index d00b4053d3f..0c351996faf 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -35,6 +35,7 @@ const { hasRuntimeUrlPlaceholders, containsGraphTokenPlaceholder, isOAuthServer, + OpenIDReauthRequiredError, } = require('@librechat/api'); const { Time, @@ -1125,6 +1126,11 @@ function createToolInstance({ error, ); + /** Carries the actionable re-auth message; the substring heuristic below would misreport it as an OAuth configuration problem */ + if (error instanceof OpenIDReauthRequiredError) { + throw error; + } + /** OAuth error, provide a helpful message */ const isOAuthError = error.message?.includes('401') || diff --git a/api/server/services/MCP.spec.js b/api/server/services/MCP.spec.js index c03d0e89028..976e6f81e82 100644 --- a/api/server/services/MCP.spec.js +++ b/api/server/services/MCP.spec.js @@ -1642,6 +1642,54 @@ describe('User parameter passing tests', () => { ); }); + it('preserves OpenIDReauthRequiredError through the OAuth error classification', async () => { + const { OpenIDReauthRequiredError } = require('@librechat/api'); + const mockUser = { id: 'reauth-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + const reauthError = new OpenIDReauthRequiredError( + 'OpenID token is expired or unavailable; re-authentication is required to resolve {{LIBRECHAT_OPENID_ACCESS_TOKEN}}.', + ); + mockGetMCPManager.mockReturnValue({ + callTool: jest.fn().mockRejectedValue(reauthError), + }); + + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + config: { requiresOAuth: false }, + toolKey: `test-tool${D}test-server`, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [`test-tool${D}test-server`]: { + function: { + description: 'Cached tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + + await expect( + mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }, + ), + ).rejects.toBe(reauthError); + }); + it('does not label OBO authentication failures as unconfigured MCP OAuth', async () => { const mockUser = { id: 'obo-user', role: 'USER' }; const mockRes = { write: jest.fn(), flush: jest.fn() }; From c8953b8f328563b53d675ab1399bcdbe64e0ec5e Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 14:23:08 -0400 Subject: [PATCH 11/15] =?UTF-8?q?=F0=9F=AA=82=20fix:=20Land=20Navigation?= =?UTF-8?q?=20Auto-Scroll=20on=20the=20Rendered=20Thread=20(#15014)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The "auto scroll to latest message" setting stopped taking readers to the newest message when opening a conversation, most visibly on long threads. `useMessageScrolling` fired its landing on the conversation id alone. That id reaches the hook a commit or more before the tree does, so `scrollIntoView` ran against the OUTGOING conversation's rows: it scrolled that thread to its end, and — having no dependency on the tree — never ran again once the requested thread mounted. The reader was left at whatever offset the old thread's bottom happened to be, which on a long thread is the top. Key the landing on the conversation that owns the RENDERED rows instead, using the same `messagesTree[0].conversationId` fallback `MessagesView` already uses to key the mount window, and land once per conversation so the tree identities a stream mints cannot haul back a reader who scrolled away. This is independent of the progressive row mounting: that window only ever grows upward from the newest row, so the end of the mounted content is already the end of the thread, and the landing needs no full mount to be correct. Measured against the real client (react-scan render tallies over a 10-message to 120-message navigation), render counts are unchanged at ~16k and the thread still mounts progressively; distance from the bottom on arrival goes 841px to 0. With progressive mounting disabled the same navigation landed 15421px from the bottom, confirming the anchoring was masking this rather than causing it. Also moves the `autoScroll` setting from Recoil to Jotai, keeping the same `autoScroll` localStorage key so a stored preference survives, and matching the `showThinking`/`smoothStreaming` atoms already served through `ToggleSwitch`. Claude-Session: https://claude.ai/code/session_01BDQSLdbwvtSqCmQSw7Nz91 Co-authored-by: Claude --- .../components/Chat/Messages/MessagesView.tsx | 3 +- .../src/components/Nav/Settings/registry.tsx | 3 +- .../__tests__/useMessageScrolling.spec.tsx | 110 ++++++++++++++++++ .../src/hooks/Messages/useMessageScrolling.ts | 61 ++++++++-- client/src/store/autoScroll.ts | 10 ++ client/src/store/settings.ts | 1 - 6 files changed, 178 insertions(+), 10 deletions(-) create mode 100644 client/src/store/autoScroll.ts diff --git a/client/src/components/Chat/Messages/MessagesView.tsx b/client/src/components/Chat/Messages/MessagesView.tsx index ed76987ccdc..8166a5d9ce5 100644 --- a/client/src/components/Chat/Messages/MessagesView.tsx +++ b/client/src/components/Chat/Messages/MessagesView.tsx @@ -9,6 +9,7 @@ import { RowMountProvider, useProgressiveRowMount } from '~/hooks/Messages'; import { MessagesViewProvider, useChatContext } from '~/Providers'; import ScrollToBottom from '~/components/Messages/ScrollToBottom'; import { steerOverlayHeightFamily } from '~/store/steer'; +import { autoScrollAtom } from '~/store/autoScroll'; import { fontSizeAtom } from '~/store/fontSize'; import MultiMessage from './MultiMessage'; import MessageNav from './MessageNav'; @@ -117,7 +118,7 @@ function MessagesViewContent({ const { index, latestMessageDepth } = useChatContext(); const isSubmitting = useRecoilValue(store.isSubmittingFamily(index)); - const autoScroll = useRecoilValue(store.autoScroll); + const autoScroll = useAtomValue(autoScrollAtom); /** Re-arm from the conversation that owns the RENDERED tree: the Recoil * conversation id lags the route during warm-cache navigation, and keying * off it would first mount the new tree unwindowed, then narrow it after diff --git a/client/src/components/Nav/Settings/registry.tsx b/client/src/components/Nav/Settings/registry.tsx index 8a81a7f40df..2c08d8fb745 100644 --- a/client/src/components/Nav/Settings/registry.tsx +++ b/client/src/components/Nav/Settings/registry.tsx @@ -43,6 +43,7 @@ import SharedLinks from '../SettingsTabs/Data/SharedLinks'; import ImageResize from '../SettingsTabs/Chat/ImageResize'; import { showThinkingAtom } from '~/store/showThinking'; import ProviderKeys from '../SettingsTabs/ProviderKeys'; +import { autoScrollAtom } from '~/store/autoScroll'; import Avatar from '../SettingsTabs/Account/Avatar'; import About from '../SettingsTabs/About/About'; import ApiKeys from '../SettingsTabs/ApiKeys'; @@ -352,7 +353,7 @@ export const registry: SettingEntry[] = [ section: 'conversations', labelKey: 'com_nav_auto_scroll', Component: toggleControl({ - stateAtom: store.autoScroll, + stateAtom: autoScrollAtom, localizationKey: 'com_nav_auto_scroll', switchId: 'autoScroll', }), diff --git a/client/src/hooks/Messages/__tests__/useMessageScrolling.spec.tsx b/client/src/hooks/Messages/__tests__/useMessageScrolling.spec.tsx index cc0771b60bb..0330e1583c2 100644 --- a/client/src/hooks/Messages/__tests__/useMessageScrolling.spec.tsx +++ b/client/src/hooks/Messages/__tests__/useMessageScrolling.spec.tsx @@ -1,5 +1,6 @@ import React from 'react'; import { RecoilRoot } from 'recoil'; +import { Provider, createStore } from 'jotai'; import { act, fireEvent, render, screen } from '@testing-library/react'; import type { TConversation, TMessage } from 'librechat-data-provider'; import { @@ -35,6 +36,7 @@ jest.mock('../messageLayout', () => ({ import useMessageScrolling from '../useMessageScrolling'; import { reconcileMessageContentLayout } from '../messageLayout'; +import { autoScrollAtom } from '~/store/autoScroll'; const mockReconcileMessageContentLayout = reconcileMessageContentLayout as jest.Mock; @@ -588,3 +590,111 @@ describe('useMessageScrolling resize reconciliation', () => { expect(mockScrollToBottom).not.toHaveBeenCalled(); }); }); + +describe('useMessageScrolling navigation landing', () => { + const treeFor = (conversationId: string): TMessage[] => [ + { ...message, conversationId } as TMessage, + ]; + + const harness = ( + store: ReturnType, + conversationId: string, + messagesTree: TMessage[] | null, + ) => ( + + + + + + + + ); + + function renderLanding( + conversationId: string, + messagesTree: TMessage[] | null, + autoScroll = true, + ) { + const store = createStore(); + store.set(autoScrollAtom, autoScroll); + const view = render(harness(store, conversationId, messagesTree)); + return { + ...view, + rerenderWith: (nextId: string, nextTree: TMessage[] | null) => + view.rerender(harness(store, nextId, nextTree)), + }; + } + + beforeEach(() => { + MockResizeObserver.reset(); + MockIntersectionObserver.reset(); + mockScrollToBottom.mockClear(); + mockScrollToBottom.cancel.mockClear(); + (global as unknown as { ResizeObserver: typeof MockResizeObserver }).ResizeObserver = + MockResizeObserver; + ( + global as unknown as { IntersectionObserver: typeof MockIntersectionObserver } + ).IntersectionObserver = MockIntersectionObserver; + }); + + afterEach(() => { + (global as unknown as { ResizeObserver: typeof ResizeObserver | undefined }).ResizeObserver = + originalResizeObserver; + ( + global as unknown as { IntersectionObserver: typeof IntersectionObserver | undefined } + ).IntersectionObserver = originalIntersectionObserver; + }); + + it('waits for the opened conversation to own the rendered rows', () => { + /** The id reaches the hook commits before the tree does. Landing on the + * outgoing thread spends the one scroll this navigation gets, and leaves + * the reader at the top of the thread they actually asked for. */ + renderLanding('conversation-2', treeFor('conversation-1')); + + expect(mockScrollToBottom).not.toHaveBeenCalled(); + }); + + it('lands once the rendered tree names the opened conversation', () => { + const { rerenderWith } = renderLanding('conversation-2', treeFor('conversation-1')); + expect(mockScrollToBottom).not.toHaveBeenCalled(); + + rerenderWith('conversation-2', treeFor('conversation-2')); + + expect(mockScrollToBottom).toHaveBeenCalledTimes(1); + }); + + it('does not re-land on the tree identities a stream mints', () => { + const { rerenderWith } = renderLanding('conversation-2', treeFor('conversation-2')); + expect(mockScrollToBottom).toHaveBeenCalledTimes(1); + + /** Every delta writes a fresh array; re-landing on those would haul back a + * reader who deliberately scrolled up. */ + rerenderWith('conversation-2', treeFor('conversation-2')); + rerenderWith('conversation-2', treeFor('conversation-2')); + + expect(mockScrollToBottom).toHaveBeenCalledTimes(1); + }); + + it('lands again for the next conversation opened', () => { + const { rerenderWith } = renderLanding('conversation-2', treeFor('conversation-2')); + expect(mockScrollToBottom).toHaveBeenCalledTimes(1); + + rerenderWith('conversation-3', treeFor('conversation-2')); + expect(mockScrollToBottom).toHaveBeenCalledTimes(1); + + rerenderWith('conversation-3', treeFor('conversation-3')); + expect(mockScrollToBottom).toHaveBeenCalledTimes(2); + }); + + it('stays put when the setting is off', () => { + renderLanding('conversation-2', treeFor('conversation-2'), false); + + expect(mockScrollToBottom).not.toHaveBeenCalled(); + }); +}); diff --git a/client/src/hooks/Messages/useMessageScrolling.ts b/client/src/hooks/Messages/useMessageScrolling.ts index 03cce408783..cd5ae9df63f 100644 --- a/client/src/hooks/Messages/useMessageScrolling.ts +++ b/client/src/hooks/Messages/useMessageScrolling.ts @@ -1,11 +1,11 @@ import { useRef, useCallback, useEffect } from 'react'; -import { useRecoilValue } from 'recoil'; +import { useAtomValue } from 'jotai'; import { Constants } from 'librechat-data-provider'; import type { TMessage } from 'librechat-data-provider'; import { useMessagesConversation, useMessagesSubmission } from '~/Providers'; import { reconcileMessageContentLayout } from './messageLayout'; import useScrollToRef from '~/hooks/useScrollToRef'; -import store from '~/store'; +import { autoScrollAtom } from '~/store/autoScroll'; const resizeFollowThreshold = 120; @@ -26,7 +26,7 @@ const prefersReducedMotion = () => window.matchMedia('(prefers-reduced-motion: reduce)').matches; export default function useMessageScrolling(messagesTree?: TMessage[] | null) { - const autoScroll = useRecoilValue(store.autoScroll); + const autoScroll = useAtomValue(autoScrollAtom); const scrollableRef = useRef(null); const contentRef = useRef(null); @@ -48,6 +48,9 @@ export default function useMessageScrolling(messagesTree?: TMessage[] | null) { const lastScrollTopRef = useRef(-1); const wasSubmittingRef = useRef(false); const suppressNextResizeFollowRef = useRef(false); + /** The conversation whose newest message has already been landed on, so the + * stream deltas that follow cannot re-take a reader who scrolled away. */ + const landedConversationRef = useRef(null); const { conversation, conversationId } = useMessagesConversation(); const { setAbortScroll, isSubmitting, abortScroll } = useMessagesSubmission(); @@ -411,15 +414,59 @@ export default function useMessageScrolling(messagesTree?: TMessage[] | null) { }; }, [isSubmitting, messagesTree, scrollToBottom, abortScroll, followBottom]); + /** + * Land on the newest message when a conversation is opened. + * + * Keyed on the conversation whose rows are actually MOUNTED, not on the id + * alone. The id reaches this hook a commit or more before the tree does, and + * firing on it scrolled the outgoing thread to its end and then never ran + * again — the reader was left wherever that put them, which on a long thread + * is the top. Waiting for the rendered tree to name the same conversation + * costs nothing and is the only moment `messages-end` is where the reader + * expects it. + * + * One landing per conversation: the tree's identity changes on every stream + * delta and cache reconcile, and re-running on those would drag a reader who + * has scrolled away back to the bottom. + */ useEffect(() => { - if (!messagesEndRef.current || !scrollableRef.current) { + if (!autoScroll) { + /** Switching the setting off releases the landing, so switching it back + * on while the same conversation is open honours it again. */ + landedConversationRef.current = null; return; } - if (scrollToBottom && autoScroll && conversationId !== Constants.NEW_CONVO) { - scrollToBottom(); + if (conversationId == null || conversationId === Constants.NEW_CONVO) { + return; } - }, [autoScroll, conversationId, scrollToBottom]); + + if (!scrollToBottom || !messagesEndRef.current || !scrollableRef.current) { + return; + } + + /** Rows are gated by the progressive mount window during a first commit, + * but that window only ever grows UPWARD from the newest row, so the end + * of the mounted content is already the end of the thread. */ + if (!messagesTree?.length) { + return; + } + + /** Same fallback `MessagesView` uses to key the mount window: a tree whose + * rows carry no conversation id is taken to be this conversation's, so a + * locally-built thread still lands instead of waiting forever. */ + const renderedConversationId = messagesTree[0]?.conversationId ?? conversationId; + if (renderedConversationId !== conversationId) { + return; + } + + if (landedConversationRef.current === conversationId) { + return; + } + + landedConversationRef.current = conversationId; + scrollToBottom(); + }, [autoScroll, conversationId, messagesTree, scrollToBottom]); return { conversation, diff --git a/client/src/store/autoScroll.ts b/client/src/store/autoScroll.ts new file mode 100644 index 00000000000..e5e1fa8b2ea --- /dev/null +++ b/client/src/store/autoScroll.ts @@ -0,0 +1,10 @@ +import { createStorageAtom } from './jotai-utils'; + +const DEFAULT_AUTO_SCROLL = false; + +/** + * Whether opening a conversation lands the reader on its newest message. + * Persisted under the same `autoScroll` key the Recoil atom used, so a stored + * preference survives the migration. + */ +export const autoScrollAtom = createStorageAtom('autoScroll', DEFAULT_AUTO_SCROLL); diff --git a/client/src/store/settings.ts b/client/src/store/settings.ts index 67e3c39b7e0..109e63e4477 100644 --- a/client/src/store/settings.ts +++ b/client/src/store/settings.ts @@ -41,7 +41,6 @@ function isSmallViewport(): boolean { const localStorageAtoms = { // General settings - autoScroll: atomWithLocalStorage('autoScroll', false), sidebarExpanded: atomWithLocalStorage( 'unifiedSidebarExpanded', !isSmallViewport(), From 2f50e38217a4fc8b0590185816b9c1cb99f4b428 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 14:24:11 -0400 Subject: [PATCH 12/15] =?UTF-8?q?=F0=9F=93=8E=20fix:=20Never=20Let=20a=20S?= =?UTF-8?q?talled=20Attachment=20Disable=20the=20Composer=20(#15013)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🖼️ fix: Keep Composer Send Enabled When an Attachment Stalls The composer's send button is gated on `hasIncompleteFiles(files)`, so any attachment that can never reach `progress: 1` reads as "still uploading" and disables send for the rest of the session — draft text intact, no error, no way out but removing the chip or reloading. Two paths could park an attachment there: - `loadImage` starts the upload from `img.onload` and had no `onerror`, so an image the browser refuses to decode (unsupported codec, truncated bytes, a revoked object URL) never uploaded at all and stranded the file at `progress: 0.2`. Drop the file and surface the error instead. - Upload completion reconciled against `temp_file_id`, the server's echo of the id the request was sent with, while every client-side handle for that upload — file map key, delayed-toast timer, recovery callbacks — is keyed by the id the client owns. A mismatch applied the completion update to a key that does not exist, leaving the attachment at `progress: 0.9`. Covered by unit regressions in the file-handling suite and a composer-level spec that drives a real upload through `ChatForm`, plus a render-bound guard on typing (react-scan measures one ChatForm render per keystroke in a browser; the guard fails on a multiplier). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Cq3tev2rPbc2pWyVJwnh6u * 🧹 fix: Stop the Draft Restore From Clobbering Live Composer Attachments `restoreFiles` runs on every `QueryKeys.files` write — an upload landing, an SSE attachment mid-run — not just on a conversation swap, and it was written as if the draft were always the whole truth: - An empty draft cleared the composer outright. On the swap path that is redundant (the effect already clears explicitly one line earlier); on the cache path an empty draft only means the draft write has not caught up, so clearing there discards an attachment the user just added — and with no text typed, the send button has nothing left to submit. Restoring now only adds. - A match replaced the composer's entry with the persisted record, dropping the local `File`, the blob preview the chip renders from (`FileRow` falls back to refetching `filepath`), and the tool resource the upload was staged under, and stamping `attached: true` so removing a chip the composer still owns leaves the file orphaned server-side. It now layers the record over the live entry and leaves `attached` to files actually adopted from a draft. Confirmed against a real browser run: the entry is at `progress: 0.9` when this restore fires, so it — not the upload's own completion — is what was re-enabling send. react-scan render counts are unchanged (typing 20 keystrokes: 111 renders, ChatForm=20; attaching an image: 1373, FileRow=6). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Cq3tev2rPbc2pWyVJwnh6u * 🔗 fix: Keep an Attachment's Stored Temporary Id Equal to Its Map Key Two follow-ups from review on the upload reconciliation. Completion stored the server's `temp_file_id` echo in the entry's value while keying the map by the id the request was sent with. `useFileDeletion` deletes map entries by the value's own `file_id` and `temp_file_id`, so where the two disagreed — the exact case the reconciliation exists to tolerate — Remove would delete the file server-side and leave the chip behind, and the draft restore could not correlate its saved key with the cached record. Store the request id. A refused image decode also left its `uploadScope.recent` reservation behind: reservations are released by the render that observes the file in the shared state, which a decode failing before that render never reaches, and once the file is deleted no later render can either. The ghost is merged into every later batch's validation, so re-picking the same file reads as a duplicate and its size keeps counting against the composer's limits. Both covered; both new guards fail without their fix. Also sorts the composer spec's imports, which the static-checks import-order gate flagged. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Cq3tev2rPbc2pWyVJwnh6u * 🧷 fix: Normalize an Upload's Temporary Id at the Cache Boundary The composer keys its file map — and the draft it saves — by the `file_id` the upload request was sent with; `temp_file_id` is only the server's echo of that id. The previous commit reconciled the composer's own entry against the request id but left the record the mutation inserts into `QueryKeys.files` carrying the raw echo, and `restoreFiles` can only correlate a saved draft id by matching a cached record's `file_id` or `temp_file_id`. Where the echo disagreed the draft matched neither, so the attachment was silently dropped on the next conversation switch or reload — the same class of loss, one layer further out. Normalize once where the response enters client state, and hand the normalized record to the mutation's callers, so the cache, the composer entry and the draft all agree on one id. An agreeing response is passed through untouched. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Cq3tev2rPbc2pWyVJwnh6u --------- Co-authored-by: Claude --- .../__tests__/ChatForm.attachments.spec.tsx | 226 +++++++++++++++++ .../__tests__/uploadNormalization.spec.tsx | 115 +++++++++ client/src/data-provider/Files/mutations.ts | 18 +- .../Files/__tests__/useFileHandling.test.ts | 231 ++++++++++++++++-- client/src/hooks/Files/useFileHandling.ts | 60 +++-- client/src/hooks/Input/useAutoSave.spec.ts | 128 ++++++++++ client/src/hooks/Input/useAutoSave.ts | 23 +- 7 files changed, 752 insertions(+), 49 deletions(-) create mode 100644 client/src/components/Chat/Input/__tests__/ChatForm.attachments.spec.tsx create mode 100644 client/src/data-provider/Files/__tests__/uploadNormalization.spec.tsx diff --git a/client/src/components/Chat/Input/__tests__/ChatForm.attachments.spec.tsx b/client/src/components/Chat/Input/__tests__/ChatForm.attachments.spec.tsx new file mode 100644 index 00000000000..dc21d840f23 --- /dev/null +++ b/client/src/components/Chat/Input/__tests__/ChatForm.attachments.spec.tsx @@ -0,0 +1,226 @@ +import React, { Profiler, useMemo, useState } from 'react'; +import '@testing-library/jest-dom'; +import { DndProvider } from 'react-dnd'; +import { useForm } from 'react-hook-form'; +import { RecoilRoot, useRecoilState } from 'recoil'; +import userEvent from '@testing-library/user-event'; +import { HTML5Backend } from 'react-dnd-html5-backend'; +import { BrowserRouter as Router } from 'react-router-dom'; +import { render, screen, waitFor } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { QueryKeys, FileSources, EModelEndpoint } from 'librechat-data-provider'; +import type { TFile, TFileUpload, TConversation } from 'librechat-data-provider'; +import type { ChatFormValues } from '~/common'; +import { ChatContext, ChatFormProvider } from '~/Providers'; +import { AuthContextProvider } from '~/hooks/AuthContext'; +import ChatForm from '../ChatForm'; +import store from '~/store'; + +const mockUpload = jest.fn(); + +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + uploadImage: (...args: unknown[]) => mockUpload(...args), + uploadFile: (...args: unknown[]) => mockUpload(...args), + }, + }; +}); + +const conversation = { + conversationId: 'new', + endpoint: EModelEndpoint.openAI, + model: 'gpt-4o', + title: 'New Chat', +} as TConversation; + +const uploadResponse = { + message: 'File uploaded', + file_id: 'server-file-id', + temp_file_id: 'temp-file-id', + filename: 'cat.png', + filepath: '/images/cat.png', + type: 'image/png', + bytes: 2048, + height: 100, + width: 100, + source: FileSources.local, + embedded: false, +} as unknown as TFileUpload; + +/** jsdom never decodes images; `decodes` mirrors a browser that can or cannot. */ +let decodes = true; + +class StubImage { + onload: (() => void) | null = null; + onerror: (() => void) | null = null; + width = 100; + height = 100; + set src(_value: string) { + setTimeout(() => (decodes ? this.onload?.() : this.onerror?.()), 0); + } +} + +let commits = 0; + +function Harness() { + const [files, setFiles] = useRecoilState(store.filesByIndex(0)); + const [isSubmitting] = useRecoilState(store.isSubmittingFamily(0)); + const [, setFilesLoading] = useState(false); + const methods = useForm({ defaultValues: { text: '' } }); + + const chatHelpers = useMemo( + () => + ({ + index: 0, + conversation, + setConversation: () => undefined, + files, + setFiles, + isSubmitting, + setIsSubmitting: () => undefined, + filesLoading: false, + setFilesLoading, + newConversation: () => undefined, + handleStopGenerating: () => undefined, + stopGenerating: () => undefined, + getMessages: () => undefined, + setMessages: () => undefined, + ask: () => undefined, + regenerate: () => undefined, + setSiblingIdx: () => undefined, + showPopover: false, + setShowPopover: () => undefined, + abortScroll: false, + setAbortScroll: () => undefined, + preset: null, + setPreset: () => undefined, + optionSettings: {}, + setOptionSettings: () => undefined, + handleRegenerate: () => undefined, + handleContinue: () => undefined, + }) as unknown as React.ContextType, + [files, setFiles, isSubmitting], + ); + + return ( + + + (commits += 1)}> + + + + + ); +} + +function renderComposer() { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + queryClient.setQueryData([QueryKeys.fileConfig], {}); + queryClient.setQueryData([QueryKeys.files], []); + queryClient.setQueryData([QueryKeys.endpoints], { [EModelEndpoint.openAI]: { order: 0 } }); + + return render( + + + + + + + + + + + , + ); +} + +const sendButton = () => screen.getByTestId('send-button'); +const attach = (container: HTMLElement, file: File) => + userEvent.upload(container.querySelector('input[type="file"]') as HTMLInputElement, file); +const image = () => new File(['image-bytes'], 'cat.png', { type: 'image/png' }); + +describe('ChatForm attachments', () => { + beforeEach(() => { + localStorage.clear(); + decodes = true; + commits = 0; + global.URL.createObjectURL = jest.fn(() => 'blob:preview'); + global.URL.revokeObjectURL = jest.fn(); + (global as unknown as { Image: unknown }).Image = StubImage; + mockUpload.mockReset(); + /** The server echoes the id the client sent back as `temp_file_id`. */ + mockUpload.mockImplementation((body: FormData) => + Promise.resolve({ ...uploadResponse, temp_file_id: body.get('file_id') as string }), + ); + }); + + test('re-enables send once an attachment finishes uploading', async () => { + const { container } = renderComposer(); + + const textarea = await screen.findByTestId('text-input'); + await userEvent.type(textarea, 'hi'); + expect(sendButton()).toBeEnabled(); + + await attach(container, image()); + await waitFor(() => expect(mockUpload).toHaveBeenCalled()); + + await waitFor(() => expect(sendButton()).toBeEnabled()); + expect(textarea).toHaveValue('hi'); + }, 20000); + + test('enables send for an attachment with no composer text', async () => { + const { container } = renderComposer(); + await screen.findByTestId('text-input'); + expect(sendButton()).toBeDisabled(); + + await attach(container, image()); + await waitFor(() => expect(mockUpload).toHaveBeenCalled()); + + await waitFor(() => expect(sendButton()).toBeEnabled()); + }, 20000); + + /** + * The upload only starts once the browser has decoded the image. A decode it + * refuses used to leave the attachment below `progress: 1`, which reads as + * "still uploading" and disabled the send button for the rest of the session. + */ + test('drops an image the browser cannot decode instead of disabling send', async () => { + decodes = false; + const { container } = renderComposer(); + + const textarea = await screen.findByTestId('text-input'); + await userEvent.type(textarea, 'hi'); + + await attach(container, image()); + await waitFor(() => expect(screen.queryByLabelText('Remove file')).not.toBeInTheDocument()); + + expect(mockUpload).not.toHaveBeenCalled(); + expect(sendButton()).toBeEnabled(); + expect(textarea).toHaveValue('hi'); + }, 20000); + + /** + * The composer is the app's busiest surface: every keystroke already re-renders + * it for the row count and the send button's enabled state, so anything that + * multiplies that work per character is a regression worth failing on. The + * measured cost is ~2.5 commits per character (react-scan reports one ChatForm + * render per keystroke in a real browser); the bound leaves headroom for jsdom + * scheduling without tolerating a doubling. + */ + test('keeps typing render-bounded', async () => { + renderComposer(); + const textarea = await screen.findByTestId('text-input'); + await waitFor(() => expect(sendButton()).toBeInTheDocument()); + + commits = 0; + await userEvent.type(textarea, 'hello there'); + + expect(commits).toBeLessThanOrEqual('hello there'.length * 3); + }, 20000); +}); diff --git a/client/src/data-provider/Files/__tests__/uploadNormalization.spec.tsx b/client/src/data-provider/Files/__tests__/uploadNormalization.spec.tsx new file mode 100644 index 00000000000..92326598ab0 --- /dev/null +++ b/client/src/data-provider/Files/__tests__/uploadNormalization.spec.tsx @@ -0,0 +1,115 @@ +import React from 'react'; +import { QueryKeys, FileSources } from 'librechat-data-provider'; +import { renderHook, act, waitFor } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import type { TFile, TFileUpload } from 'librechat-data-provider'; + +const mockUploadFile = jest.fn(); + +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + uploadFile: (...args: unknown[]) => mockUploadFile(...args), + uploadImage: (...args: unknown[]) => mockUploadFile(...args), + }, + }; +}); + +jest.mock('../../Endpoints', () => ({ + useGetStartupConfig: () => ({ data: undefined }), +})); + +import { useUploadFileMutation } from '../mutations'; + +const REQUEST_FILE_ID = 'client-request-id'; + +const uploadResponse = { + file_id: 'server-file-id', + temp_file_id: 'an-id-the-client-never-sent', + filename: 'notes.txt', + filepath: '/files/notes.txt', + type: 'text/plain', + bytes: 12, + object: 'file', + usage: 0, + user: 'user-1', + embedded: false, + source: FileSources.local, +} as unknown as TFileUpload; + +const body = () => { + const formData = new FormData(); + formData.append('file_id', REQUEST_FILE_ID); + formData.append('message_file', 'true'); + return formData; +}; + +/** + * The composer keys its file map, and the draft it saves, by the `file_id` the + * request was sent with; `temp_file_id` is only the server's echo of that id. + * A cached record carrying an echo that disagrees cannot be correlated back to + * the draft, so `useAutoSave` cannot restore the attachment after a conversation + * switch or a reload. + */ +describe('useUploadFileMutation — temporary id normalization', () => { + let queryClient: QueryClient; + + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ); + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + mockUploadFile.mockReset(); + mockUploadFile.mockResolvedValue(uploadResponse); + }); + + test('caches the uploaded record under the id the request was sent with', async () => { + const { result } = renderHook(() => useUploadFileMutation(), { wrapper }); + + act(() => { + result.current.mutate(body()); + }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + const [cached] = queryClient.getQueryData([QueryKeys.files]) ?? []; + expect(cached).toMatchObject({ + file_id: 'server-file-id', + temp_file_id: REQUEST_FILE_ID, + }); + }); + + test('hands the normalized record to the caller', async () => { + const onSuccess = jest.fn(); + const { result } = renderHook(() => useUploadFileMutation({ onSuccess }), { wrapper }); + + act(() => { + result.current.mutate(body()); + }); + + await waitFor(() => expect(onSuccess).toHaveBeenCalled()); + + expect(onSuccess.mock.calls[0][0]).toMatchObject({ temp_file_id: REQUEST_FILE_ID }); + }); + + test('leaves an agreeing response untouched', async () => { + mockUploadFile.mockResolvedValue({ ...uploadResponse, temp_file_id: REQUEST_FILE_ID }); + const onSuccess = jest.fn(); + const { result } = renderHook(() => useUploadFileMutation({ onSuccess }), { wrapper }); + + act(() => { + result.current.mutate(body()); + }); + + await waitFor(() => expect(onSuccess).toHaveBeenCalled()); + + const [cached] = queryClient.getQueryData([QueryKeys.files]) ?? []; + expect(cached).toBe(onSuccess.mock.calls[0][0]); + }); +}); diff --git a/client/src/data-provider/Files/mutations.ts b/client/src/data-provider/Files/mutations.ts index a538eea4b84..097829022ae 100644 --- a/client/src/data-provider/Files/mutations.ts +++ b/client/src/data-provider/Files/mutations.ts @@ -44,8 +44,18 @@ export const useUploadFileMutation = ( }, ...options, onSuccess: (data, formData, context) => { + /** `temp_file_id` is the server's echo of the `file_id` the request was sent + * with, and that request id is what every client-side handle for the upload + * is keyed by — the composer's file map, the draft it saves, the delayed + * toast timer. A record cached under an echo that disagrees cannot be + * correlated back to the draft, so the attachment is silently dropped on the + * next conversation switch or reload. Normalize once, on the way in. */ + const requestFileId = (formData.get('file_id') as string | null) ?? data.temp_file_id; + const file = + data.temp_file_id === requestFileId ? data : { ...data, temp_file_id: requestFileId }; + queryClient.setQueryData([QueryKeys.files], (_files) => [ - data, + file, ...(_files ?? []), ]); @@ -56,7 +66,7 @@ export const useUploadFileMutation = ( const tool_resource = (formData.get('tool_resource') as string | undefined) ?? ''; if (message_file === 'true') { - onSuccess?.(data, formData, context); + onSuccess?.(file, formData, context); return; } @@ -92,7 +102,7 @@ export const useUploadFileMutation = ( } if (!assistant_id) { - onSuccess?.(data, formData, context); + onSuccess?.(file, formData, context); return; } @@ -136,7 +146,7 @@ export const useUploadFileMutation = ( }; }, ); - onSuccess?.(data, formData, context); + onSuccess?.(file, formData, context); }, }); }; diff --git a/client/src/hooks/Files/__tests__/useFileHandling.test.ts b/client/src/hooks/Files/__tests__/useFileHandling.test.ts index c643b4b9f40..0378f58b042 100644 --- a/client/src/hooks/Files/__tests__/useFileHandling.test.ts +++ b/client/src/hooks/Files/__tests__/useFileHandling.test.ts @@ -8,20 +8,26 @@ import { } from 'librechat-data-provider'; type MockUploadMutationOptions = { - onSuccess?: (data: { - temp_file_id: string; - file_id: string; - filepath: string; - type: string; - filename: string; - source: string; - embedded: boolean; - height?: number; - width?: number; - }) => void; + onSuccess?: ( + data: { + temp_file_id: string; + file_id: string; + filepath: string; + type: string; + filename: string; + source: string; + embedded: boolean; + height?: number; + width?: number; + }, + body: FormData, + ) => void; onError?: (error: unknown, body: FormData) => void; }; +/** Mirrors a browser that can (or cannot) decode the bytes it was handed. */ +let mockImageDecodes = true; + beforeAll(() => { global.URL.createObjectURL = jest.fn(() => 'blob:mock-url'); global.URL.revokeObjectURL = jest.fn(); @@ -31,9 +37,10 @@ beforeAll(() => { width = 640; height = 480; onload: (() => void) | null = null; + onerror: (() => void) | null = null; set src(_src: string) { - queueMicrotask(() => this.onload?.()); + queueMicrotask(() => (mockImageDecodes ? this.onload?.() : this.onerror?.())); } }, }); @@ -131,13 +138,18 @@ jest.mock('../useClientResize', () => ({ })), })); +const mockAddFile = jest.fn(); +const mockReplaceFile = jest.fn(); +const mockUpdateFileById = jest.fn(); +const mockDeleteFileById = jest.fn(); + jest.mock('../useUpdateFiles', () => ({ __esModule: true, default: jest.fn(() => ({ - addFile: jest.fn(), - replaceFile: jest.fn(), - updateFileById: jest.fn(), - deleteFileById: jest.fn(), + addFile: mockAddFile, + replaceFile: mockReplaceFile, + updateFileById: mockUpdateFileById, + deleteFileById: mockDeleteFileById, })), })); @@ -173,6 +185,7 @@ describe('useFileHandling', () => { mockFileConfig = null; mockIsConfigPending = false; mockIsTemporary = false; + mockImageDecodes = true; mockUploadOptions = {}; }); @@ -1057,15 +1070,18 @@ describe('useFileHandling', () => { expect(onStart).toHaveBeenCalledWith(fileId); act(() => { - mockUploadOptions.onSuccess?.({ - temp_file_id: fileId, - file_id: 'saved-file-id', - filepath: '/files/notes.txt', - type: 'text/plain', - filename: 'notes.txt', - source: 'local', - embedded: false, - }); + mockUploadOptions.onSuccess?.( + { + temp_file_id: fileId, + file_id: 'saved-file-id', + filepath: '/files/notes.txt', + type: 'text/plain', + filename: 'notes.txt', + source: 'local', + embedded: false, + }, + uploadBody, + ); }); expect(onSuccess).toHaveBeenCalledWith(fileId); @@ -1160,4 +1176,169 @@ describe('useFileHandling', () => { consoleLog.mockRestore(); }); }); + describe('stalled attachments', () => { + /** + * The upload only starts once the browser has decoded the image, so a decode + * it refuses must not leave the attachment parked below `progress: 1` — the + * composer reads that as "still uploading" and disables send for the rest of + * the session. + */ + it('drops an image the browser cannot decode instead of parking it mid-upload', async () => { + mockImageDecodes = false; + const useFileHandling = await loadHook(); + const { result } = renderHook(() => useFileHandling()); + + await act(async () => { + await result.current.handleFiles([ + new File(['broken'], 'photo.png', { type: 'image/png' }), + ]); + }); + await act(async () => { + await Promise.resolve(); + }); + + const addedFile = mockAddFile.mock.calls[0][0] as { file_id: string }; + expect(mockMutate).not.toHaveBeenCalled(); + expect(mockDeleteFileById).toHaveBeenCalledWith(addedFile.file_id); + }); + + it('reports the failure through the upload lifecycle when a decode fails', async () => { + mockImageDecodes = false; + const onError = jest.fn(); + const useFileHandling = await loadHook(); + const { result } = renderHook(() => useFileHandling()); + + await act(async () => { + await result.current.handleFiles( + [new File(['broken'], 'photo.png', { type: 'image/png' })], + undefined, + { fileId: 'pending-paste-id', onError }, + ); + }); + await act(async () => { + await Promise.resolve(); + }); + + expect(onError).toHaveBeenCalledWith('pending-paste-id'); + }); + + /** + * Every client-side handle for an upload is keyed by the id the request was + * sent with; `temp_file_id` is only the server's echo of it. Reconciling + * against the echo strands the attachment when the two disagree. + */ + it('keys the completed attachment temporary id to the id the request was sent with', async () => { + jest.useFakeTimers(); + const consoleLog = jest.spyOn(console, 'log').mockImplementation(() => undefined); + const useFileHandling = await loadHook(); + const { result } = renderHook(() => useFileHandling()); + + await act(async () => { + await result.current.handleFiles([ + new File(['hello'], 'notes.txt', { type: 'text/plain' }), + ]); + }); + + const uploadBody = mockMutate.mock.calls[0][0] as FormData; + const fileId = uploadBody.get('file_id') as string; + + act(() => { + mockUploadOptions.onSuccess?.( + { + temp_file_id: 'an-id-the-client-never-sent', + file_id: 'saved-file-id', + filepath: '/files/notes.txt', + type: 'text/plain', + filename: 'notes.txt', + source: 'local', + embedded: false, + }, + uploadBody, + ); + }); + act(() => { + jest.runAllTimers(); + }); + jest.useRealTimers(); + + /** Removal deletes by the value's own ids, so a temporary id that is not the + * map key leaves a chip the user cannot clear. */ + const [, completion] = mockUpdateFileById.mock.calls.at(-1) as [string, { progress: number }]; + expect(completion).toMatchObject({ progress: 1, temp_file_id: fileId }); + consoleLog.mockRestore(); + }); + + it('releases the upload reservation when a decode fails', async () => { + /** A stable setter so both batches share one upload scope, and a file map that + * never observes the file — the render that would release the reservation + * cannot happen once the failed decode has removed it. */ + const sharedState = { + files: new Map(), + setFiles: jest.fn(), + setFilesLoading: mockSetFilesLoading, + }; + const { useFileHandlingNoChatContext } = await import('../useFileHandling'); + const { result, rerender } = renderHook(() => + useFileHandlingNoChatContext(undefined, sharedState), + ); + const pick = () => new File(['broken'], 'photo.png', { type: 'image/png' }); + + mockImageDecodes = false; + await act(async () => { + await result.current.handleFiles([pick()]); + }); + await act(async () => { + await Promise.resolve(); + }); + + rerender(); + mockValidateFileDuplicates.mockClear(); + mockImageDecodes = true; + await act(async () => { + await result.current.handleFiles([pick()]); + }); + + /** A reservation the failed decode left behind is merged into the next batch's + * validation, so re-picking the same file reads as a duplicate and its size + * keeps counting against the composer's limits. */ + const [{ files: validatedAgainst }] = mockValidateFileDuplicates.mock.calls[0]; + expect(validatedAgainst.size).toBe(0); + expect(mockMutate).toHaveBeenCalledTimes(1); + }); + + it('completes the attachment against the id the request was sent with', async () => { + const consoleLog = jest.spyOn(console, 'log').mockImplementation(() => undefined); + const useFileHandling = await loadHook(); + const { result } = renderHook(() => useFileHandling()); + + await act(async () => { + await result.current.handleFiles([ + new File(['hello'], 'notes.txt', { type: 'text/plain' }), + ]); + }); + + const uploadBody = mockMutate.mock.calls[0][0] as FormData; + const fileId = uploadBody.get('file_id') as string; + + act(() => { + mockUploadOptions.onSuccess?.( + { + temp_file_id: 'an-id-the-client-never-sent', + file_id: 'saved-file-id', + filepath: '/files/notes.txt', + type: 'text/plain', + filename: 'notes.txt', + source: 'local', + embedded: false, + }, + uploadBody, + ); + }); + + const updatedIds = mockUpdateFileById.mock.calls.map(([id]) => id); + expect(updatedIds).toContain(fileId); + expect(updatedIds).not.toContain('an-id-the-client-never-sent'); + consoleLog.mockRestore(); + }); + }); }); diff --git a/client/src/hooks/Files/useFileHandling.ts b/client/src/hooks/Files/useFileHandling.ts index 8ea2ef3f79d..1eac7e627cd 100644 --- a/client/src/hooks/Files/useFileHandling.ts +++ b/client/src/hooks/Files/useFileHandling.ts @@ -214,16 +214,23 @@ const useFileHandlingCore = (params: UseFileHandling | undefined, fileState: Fil const uploadFile = useUploadFileMutation( { - onSuccess: (data) => { - takeUploadRecovery(data.temp_file_id)?.onSuccess?.(data.temp_file_id); - clearUploadTimer(data.temp_file_id); + onSuccess: (data, body) => { + /** Every client-side handle for this upload — the file map key, the delayed + * toast timer, the recovery callbacks — is the id the request was sent with. + * `temp_file_id` is the server's echo of it, so trusting the echo turns any + * mismatch into a completion update applied to a key that does not exist: + * the attachment stays below `progress: 1` and the send button never + * re-enables. Reconcile against the id we own. */ + const fileId = (body.get('file_id') as string | null) ?? data.temp_file_id; + takeUploadRecovery(fileId)?.onSuccess?.(fileId); + clearUploadTimer(fileId); console.log('upload success', data); if (agent_id) { queryClient.refetchQueries([QueryKeys.agent, agent_id]); return; } updateFileById( - data.temp_file_id, + fileId, { progress: 0.9, filepath: data.filepath, @@ -232,17 +239,22 @@ const useFileHandlingCore = (params: UseFileHandling | undefined, fileState: Fil ); setTimeout(() => { - const cachedBlob = getCachedPreview(data.temp_file_id); - if (cachedBlob && data.file_id !== data.temp_file_id) { + const cachedBlob = getCachedPreview(fileId); + if (cachedBlob && data.file_id !== fileId) { cachePreview(data.file_id, cachedBlob); - removePreviewEntry(data.temp_file_id); + removePreviewEntry(fileId); } updateFileById( - data.temp_file_id, + fileId, { progress: 1, file_id: data.file_id, - temp_file_id: data.temp_file_id, + /** The stored temporary id has to stay the one this entry is keyed + * by: removal reads `file_id` and `temp_file_id` off the value and + * deletes those keys, and the draft restore correlates the cached + * record by the key it saved. Keeping the server's echo here would + * leave a chip that Remove deletes server-side but cannot clear. */ + temp_file_id: fileId, filepath: data.filepath, type: data.type, height: data.height, @@ -389,15 +401,35 @@ const useFileHandlingCore = (params: UseFileHandling | undefined, fileState: Fil ) => { const img = new Image(); img.onload = async () => { - extendedFile.width = img.width; - extendedFile.height = img.height; - extendedFile = { + const measuredFile: ExtendedFile = { ...extendedFile, + width: img.width, + height: img.height, progress: 0.6, }; - replaceFile(extendedFile); + replaceFile(measuredFile); - await startUpload(extendedFile, uploadLifecycle); + await startUpload(measuredFile, uploadLifecycle); + }; + /** The upload only starts once the browser has decoded the image, so a decode + * it refuses (unsupported codec, truncated bytes, a revoked object URL) would + * otherwise strand the attachment below `progress: 1` — which reads as "still + * uploading" and keeps the composer's send button disabled for the rest of the + * session, with nothing to click and no error to explain it. Drop the file and + * say so instead. */ + img.onerror = () => { + clearUploadTimer(extendedFile.file_id); + takeUploadRecovery(extendedFile.file_id)?.onError?.(extendedFile.file_id); + deleteFileById(extendedFile.file_id); + /** Reservations are released by the render that observes the file in the + * shared state, which a decode failing before that render never reaches — + * and once the file is gone no later render can either. A leaked one is + * merged into every subsequent batch's validation, so re-picking the same + * file reads as a duplicate and its size keeps counting against the limit. */ + uploadScope.recent.delete(extendedFile.file_id); + removePreviewEntry(extendedFile.file_id); + URL.revokeObjectURL(preview); + setError('com_error_files_process'); }; img.src = preview; }; diff --git a/client/src/hooks/Input/useAutoSave.spec.ts b/client/src/hooks/Input/useAutoSave.spec.ts index 52932078481..8331f9b54f9 100644 --- a/client/src/hooks/Input/useAutoSave.spec.ts +++ b/client/src/hooks/Input/useAutoSave.spec.ts @@ -529,3 +529,131 @@ describe('useAutoSave — side-by-side pending drafts', () => { expect(getFilesDraft(`${Constants.NEW_CONVO}:1`)).toEqual({ fileIds: [], pendingPastes: {} }); }); }); + +describe('useAutoSave — file cache updates', () => { + const liveAttachment = { + file_id: 'client-temp-id', + type: 'image/png', + size: 2048, + progress: 0.9, + preview: 'blob:local-preview', + tool_resource: 'file_search', + file: new File(['bytes'], 'cat.png', { type: 'image/png' }), + }; + const persistedRecord = { + file_id: 'server-file-id', + temp_file_id: 'client-temp-id', + filename: 'cat.png', + filepath: '/images/cat.png', + type: 'image/png', + bytes: 2048, + object: 'file', + usage: 0, + user: 'user-1', + embedded: false, + }; + + const applySetFiles = (setFiles: jest.Mock, current: Map) => + setFiles.mock.calls.reduce( + (files, [update]) => (typeof update === 'function' ? update(files) : update), + current, + ) as Map>; + + /** + * The file cache is rewritten on every upload and on every attachment an agent + * emits mid-run, and this hook restores from it. An empty draft there means the + * draft write has not caught up — not that the composer is empty — so clearing + * would drop an attachment the user just added (and, with no text typed, leave + * them with nothing submittable). + */ + it('leaves live attachments alone when the file cache changes with no saved draft', () => { + const setFiles = jest.fn(); + const files = new Map([['client-temp-id', liveAttachment]]); + + const { rerender } = renderHook( + ({ fileList }: { fileList: unknown[] }) => { + (useGetFiles as jest.Mock).mockReturnValue({ data: fileList }); + return useAutoSave({ + conversationId: 'convo-1', + textAreaRef: makeTextAreaRef(), + files, + setFiles, + }); + }, + { initialProps: { fileList: [] as unknown[] } }, + ); + + setFiles.mockClear(); + /** The draft is gone the moment storage refuses or evicts the write — another + * tab clearing it, a quota failure, private browsing. The attachment the user + * just added is still in the composer either way. */ + localStorage.clear(); + act(() => { + rerender({ fileList: [persistedRecord] }); + }); + + expect(applySetFiles(setFiles, files).size).toBe(1); + }); + + /** + * The restore also lands on entries the composer still owns, so it has to layer + * the persisted record over them rather than replace them: the blob preview the + * chip renders from and the tool resource the upload was staged under exist only + * locally, and `attached` decides whether removing the chip deletes the file. + */ + it('layers the persisted record over a live attachment instead of replacing it', () => { + const setFiles = jest.fn(); + const files = new Map([['client-temp-id', liveAttachment]]); + setFilesDraft('convo-1', { fileIds: ['client-temp-id'], pendingPastes: {} }); + + const { rerender } = renderHook( + ({ fileList }: { fileList: unknown[] }) => { + (useGetFiles as jest.Mock).mockReturnValue({ data: fileList }); + return useAutoSave({ + conversationId: 'convo-1', + textAreaRef: makeTextAreaRef(), + files, + setFiles, + }); + }, + { initialProps: { fileList: [] as unknown[] } }, + ); + + /** Past the mount swap, which clears the composer itself before restoring. */ + setFiles.mockClear(); + act(() => { + rerender({ fileList: [persistedRecord] }); + }); + + const restored = applySetFiles(setFiles, files).get('client-temp-id'); + expect(restored).toMatchObject({ + file_id: 'server-file-id', + filepath: '/images/cat.png', + progress: 1, + preview: 'blob:local-preview', + tool_resource: 'file_search', + attached: false, + }); + expect(restored?.file).toBeInstanceOf(File); + }); + + it('marks a file restored from a draft alone as attached', () => { + const setFiles = jest.fn(); + setFilesDraft('convo-1', { fileIds: ['client-temp-id'], pendingPastes: {} }); + + renderHook(() => { + (useGetFiles as jest.Mock).mockReturnValue({ data: [persistedRecord] }); + return useAutoSave({ + conversationId: 'convo-1', + textAreaRef: makeTextAreaRef(), + files: new Map(), + setFiles, + }); + }); + + expect(applySetFiles(setFiles, new Map()).get('client-temp-id')).toMatchObject({ + attached: true, + progress: 1, + }); + }); +}); diff --git a/client/src/hooks/Input/useAutoSave.ts b/client/src/hooks/Input/useAutoSave.ts index 15687b686ad..ee9ce33ca2d 100644 --- a/client/src/hooks/Input/useAutoSave.ts +++ b/client/src/hooks/Input/useAutoSave.ts @@ -65,11 +65,12 @@ export const useAutoSave = ({ (id: string): PendingTextAttachmentDraft[] => { const filesDraft = getFilesDraft(id); - if (filesDraft.fileIds.length === 0) { - setFiles(new Map()); - return []; - } - if (fileList == null) { + /** Restoring adds what the draft holds; it never clears. The conversation + * swap below owns that, and does it explicitly before calling here — while + * the file-cache path runs on every `QueryKeys.files` write (an upload + * landing, an SSE attachment during a run), where an empty draft means the + * write has not caught up yet, not that the user has no attachments. */ + if (filesDraft.fileIds.length === 0 || fileList == null) { return []; } @@ -95,10 +96,20 @@ export const useAutoSave = ({ delete pendingPastes[fileId]; setFiles((currentFiles) => { const updatedFiles = new Map(currentFiles); + /** The same entry may still be live in the composer — this path also + * runs when the upload that created it writes the file cache. Layer + * the persisted record over it instead of replacing it: the local + * `File`, the blob preview the chip renders from, and the tool + * resource the upload was staged under exist nowhere on the server + * record, and `attached` decides whether removing the chip also + * deletes the file, which a composer-owned upload still needs. */ + const live = updatedFiles.get(fileIdToRecover); updatedFiles.set(fileIdToRecover, { + ...live, ...fileToRecover, + preview: live?.preview ?? fileToRecover.preview, progress: 1, - attached: true, + attached: live ? (live.attached ?? false) : true, size: fileToRecover.bytes, }); return updatedFiles; From ec5283452e41ca58ffbca070ee6d303b8850d3c4 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 15:40:27 -0400 Subject: [PATCH 13/15] =?UTF-8?q?=F0=9F=AA=AA=20fix:=20Preserve=20Detached?= =?UTF-8?q?=20Subagent=20Owner=20Context=20(#15015)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: preserve detached subagent owner context * style: sort detached subagent imports --- .../api/src/agents/subagentThreads.spec.ts | 53 ++- packages/api/src/agents/subagentThreads.ts | 303 ++++++++++-------- 2 files changed, 225 insertions(+), 131 deletions(-) diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index 649bb8a0105..e513f12e2ae 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -2,8 +2,15 @@ import mongoose from 'mongoose'; import { randomUUID } from 'node:crypto'; import { MongoMemoryServer } from 'mongodb-memory-server'; import { Constants, EModelEndpoint } from 'librechat-data-provider'; -import { createMethods, createModels, logger } from '@librechat/data-schemas'; import { AIMessage, HumanMessage } from '@librechat/agents/langchain/messages'; +import { + createMethods, + createModels, + getTenantId, + getUserId, + logger, + tenantStorage, +} from '@librechat/data-schemas'; import type { SubagentTaskClaim, SubagentTaskControlCommand, @@ -376,6 +383,50 @@ describe('SubagentThreadTaskStore', () => { expect(await methods.getConvo(userId, requireThreadId(started))).not.toBeNull(); }); + it('reconstructs trusted owner context after the admitting request has ended', async () => { + const userId = 'detached-context-user'; + const tenantId = 'detached-context-tenant'; + const parentConversationId = randomUUID(); + await tenantStorage.run({ tenantId, userId }, async () => + saveParent(userId, parentConversationId, { tenantId }), + ); + + const observedContexts: Array<{ tenantId?: string; userId?: string }> = []; + const observeContext = () => { + observedContexts.push({ tenantId: getTenantId(), userId: getUserId() }); + }; + const store = new SubagentThreadTaskStore(methods, { + isOwnerActive: async () => { + observeContext(); + return true; + }, + }); + const config = buildSubagentThreadTaskConfig(store, { + userId, + tenantId, + parentConversationId, + }); + const defaultRun = taskRequest(config.scopeId).run; + const run = jest.fn(async (...args: Parameters) => { + observeContext(); + return defaultRun(...args); + }); + + /** `start` deliberately runs outside `tenantStorage.run`: the detached task + * owns only its serialized host scope once the HTTP request has returned. */ + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).toHaveBeenCalledTimes(1); + expect(observedContexts.length).toBeGreaterThan(0); + expect(observedContexts).toEqual(observedContexts.map(() => ({ tenantId, userId }))); + const messages = await tenantStorage.run({ tenantId, userId }, async () => + methods.getMessages({ user: userId, conversationId: requireThreadId(started) }), + ); + expect(messages).toHaveLength(2); + expect(messages.every((message) => message.tenantId === tenantId)).toBe(true); + }); + it('fails without leaving an orphan when parent persistence rejects', async () => { const userId = 'parent-gate-failure-user'; const parentConversationId = randomUUID(); diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index becac411041..eb3f21fcac7 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -1,6 +1,6 @@ import { randomUUID } from 'node:crypto'; -import { logger } from '@librechat/data-schemas'; import { InMemorySubagentTaskStore } from '@librechat/agents'; +import { logger, tenantStorage } from '@librechat/data-schemas'; import { EModelEndpoint, Constants } from 'librechat-data-provider'; import { mapChatMessagesToStoredMessages, @@ -63,6 +63,7 @@ export interface SubagentCancellationPlan { } /** Three missed 10-second transport heartbeats retire a crashed owner. */ const DEFAULT_TASK_ROUTING_TTL_MS = 30_000; +const SLOW_PREPARATION_WARN_MS = 5_000; const MAX_TRANSCRIPT_BYTES = 12 * 1024 * 1024; const TRANSCRIPT_SELECT = 'messageId parentMessageId text createdAt +subagentTranscript +subagentTask'; @@ -380,6 +381,21 @@ function safeErrorMessage(error: unknown): string { return `Subagent task failed: ${publicFailureDetail(error).slice(0, 2_000)}`; } +async function observeSlowPreparation( + operation: Promise, + context: { stage: string; taskId: string; threadId: string }, +): Promise { + const warning = setTimeout(() => { + logger.warn('[subagentThreads] Child-thread preparation is still waiting', context); + }, SLOW_PREPARATION_WARN_MS); + warning.unref?.(); + try { + return await operation; + } finally { + clearTimeout(warning); + } +} + /** Persists view-only logical child threads with owner-routed controls and a shared execution fence. */ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { readonly supportsThreadContinuation = true; @@ -519,128 +535,129 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { started = super.start({ ...request, threadId, - run: async (runtime: SubagentTaskRuntime) => { - lease.taskId = runtime.taskId; - lease.running = true; - const detachedUsage: UsageMetadata[] = []; - let prepared: PreparedThread | undefined; - try { - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - /** Publish the owner address before any provider work: a child running - * while unaddressable cannot be polled, controlled, or cancelled, and its - * side effects would already have happened by the time a heartbeat - * republished it. A failed registration fails the task closed instead. */ - await this.taskControlTransport?.registerTask( - request.scopeId, - runtime.taskId, - this.taskRoutingTtlMs, - ); - await parentReady; - prepared = await this.prepareThread( - request.scopeId, - scope, - threadId, - isContinuation, - request, - runtime.taskId, - lease, - ); - await this.registerTaskWakeup(scope, prepared.conversation.conversationId, request, { - taskId: prepared.replay?.taskId ?? runtime.taskId, - parentRunId: prepared.replay?.parentRunId ?? request.parentRunId, - createdAt: prepared.taskCreatedAt, - }); - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - if (prepared.replay != null) { - if (prepared.replay.status === 'completed') { - return { content: prepared.replay.content }; + run: (runtime: SubagentTaskRuntime) => + this.runWithOwnerContext(scope, async () => { + lease.taskId = runtime.taskId; + lease.running = true; + const detachedUsage: UsageMetadata[] = []; + let prepared: PreparedThread | undefined; + try { + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); } - throw new SubagentThreadPublicError(prepared.replay.content); - } - if (!(await this.renewSharedLease(scope, threadId, lease))) { - throw new SubagentThreadPublicError( - 'This child thread is already being continued by another run.', + /** Publish the owner address before any provider work: a child running + * while unaddressable cannot be polled, controlled, or cancelled, and its + * side effects would already have happened by the time a heartbeat + * republished it. A failed registration fails the task closed instead. */ + await this.taskControlTransport?.registerTask( + request.scopeId, + runtime.taskId, + this.taskRoutingTtlMs, ); - } - const preparedThread = prepared; - const result = await runWithDetachedSubagentUsage(detachedUsage, () => - request.run(runtime, preparedThread.initialMessages), - ); - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - if (!(await this.renewSharedLease(scope, threadId, lease))) { - throw new SubagentThreadPublicError( - 'This child thread is already being continued by another run.', + await parentReady; + prepared = await this.prepareThread( + request.scopeId, + scope, + threadId, + isContinuation, + request, + runtime.taskId, + lease, + ); + await this.registerTaskWakeup(scope, prepared.conversation.conversationId, request, { + taskId: prepared.replay?.taskId ?? runtime.taskId, + parentRunId: prepared.replay?.parentRunId ?? request.parentRunId, + createdAt: prepared.taskCreatedAt, + }); + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + if (prepared.replay != null) { + if (prepared.replay.status === 'completed') { + return { content: prepared.replay.content }; + } + throw new SubagentThreadPublicError(prepared.replay.content); + } + if (!(await this.renewSharedLease(scope, threadId, lease))) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + const preparedThread = prepared; + const result = await runWithDetachedSubagentUsage(detachedUsage, () => + request.run(runtime, preparedThread.initialMessages), + ); + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + if (!(await this.renewSharedLease(scope, threadId, lease))) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + lease.settling = true; + await this.persistResult( + scope, + request, + runtime.taskId, + prepared, + result, + detachedUsage, + ); + return result; + } catch (error) { + /** A replay is already terminal in Mongo. A temporary wakeup-queue + * outage must not overwrite that canonical result with a new error. */ + if (prepared?.replay != null) { + throw error; + } + const mayPersist = + lease.shared == null || (await this.renewSharedLease(scope, threadId, lease)); + const terminalTask = this.get(request.scopeId, runtime.taskId); + if (runtime.signal.aborted && terminalTask?.status === 'cancelled') { + if (mayPersist) { + await this.persistCancellation( + scope, + threadId, + request, + runtime.taskId, + detachedUsage, + ).catch((persistError) => { + logger.error( + '[subagentThreads] Failed to persist child-thread cancellation', + persistError, + ); + }); + } + throw error; + } + logger.error( + '[subagentThreads] Child-thread execution failed', + publicFailureDetail(error), ); - } - lease.settling = true; - await this.persistResult( - scope, - request, - runtime.taskId, - prepared, - result, - detachedUsage, - ); - return result; - } catch (error) { - /** A replay is already terminal in Mongo. A temporary wakeup-queue - * outage must not overwrite that canonical result with a new error. */ - if (prepared?.replay != null) { - throw error; - } - const mayPersist = - lease.shared == null || (await this.renewSharedLease(scope, threadId, lease)); - const terminalTask = this.get(request.scopeId, runtime.taskId); - if (runtime.signal.aborted && terminalTask?.status === 'cancelled') { if (mayPersist) { - await this.persistCancellation( + await this.persistFailure( scope, threadId, request, runtime.taskId, + error, detachedUsage, ).catch((persistError) => { logger.error( - '[subagentThreads] Failed to persist child-thread cancellation', + '[subagentThreads] Failed to persist child-thread failure', persistError, ); }); } - throw error; - } - logger.error( - '[subagentThreads] Child-thread execution failed', - publicFailureDetail(error), - ); - if (mayPersist) { - await this.persistFailure( - scope, - threadId, - request, - runtime.taskId, - error, - detachedUsage, - ).catch((persistError) => { - logger.error( - '[subagentThreads] Failed to persist child-thread failure', - persistError, - ); - }); - } - throw new Error(publicFailureDetail(error)); - } finally { - await this.stopAndReleaseSharedLease(scope, threadId, lease); - if (this.activeThreads.get(lockKey) === lease) { - this.activeThreads.delete(lockKey); + throw new Error(publicFailureDetail(error)); + } finally { + await this.stopAndReleaseSharedLease(scope, threadId, lease); + if (this.activeThreads.get(lockKey) === lease) { + this.activeThreads.delete(lockKey); + } } - } - }, + }), }); } catch (error) { if (ownsLease && this.activeThreads.get(lockKey) === lease) { @@ -657,6 +674,20 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { return started; } + /** Detached tasks intentionally outlive the HTTP request that admitted them. + * Reconstruct only the trusted owner identity carried by the opaque task scope + * so tenant-isolated database reads and lazy child initialization do not depend + * on request AsyncLocalStorage remaining alive after the parent turn returns. */ + private runWithOwnerContext(scope: SubagentThreadScope, run: () => Promise): Promise { + return tenantStorage.run( + { + userId: scope.userId, + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + }, + run, + ); + } + /** * Claims locally when possible, otherwise asks the registered owning replica. * @@ -1566,13 +1597,22 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { /** Account deletion can fence the owner after the optimistic probe but before * this lease exists. Once the lease is visible, revalidate so deletion either * observes and drains us or wins before any provider work can begin. */ - if (!(await this.isOwnerActive(scope.userId))) { + if ( + !(await observeSlowPreparation(this.isOwnerActive(scope.userId), { + stage: 'owner_recheck', + taskId, + threadId, + })) + ) { throw new SubagentThreadDeletedError('The thread owner is unavailable.'); } - const allMessages = (await this.methods.getMessages( - { conversationId: threadId, user: scope.userId }, - TRANSCRIPT_SELECT, - { sort: { createdAt: 1, _id: 1 } }, + const allMessages = (await observeSlowPreparation( + this.methods.getMessages( + { conversationId: threadId, user: scope.userId }, + TRANSCRIPT_SELECT, + { sort: { createdAt: 1, _id: 1 } }, + ), + { stage: 'transcript_read', taskId, threadId }, )) as ThreadMessage[]; const attemptKey = createSubagentAttemptKey(scopeId, request.idempotencyKey); const requestFingerprint = normalizedRequestFingerprint(request); @@ -1685,25 +1725,28 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { break; } } - const savedUserMessage = await this.methods.saveMessage( - { userId: scope.userId }, - { - messageId: userMessageId, - conversationId: threadId, - parentMessageId, - sender: 'User', - text: request.input, - endpoint: EModelEndpoint.agents, - isCreatedByUser: true, - subagentTask: { - attemptKey, - parentRunId: request.parentRunId, - ...(requestFingerprint == null ? {} : { requestFingerprint }), - status: 'running', + const savedUserMessage = await observeSlowPreparation( + this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: userMessageId, + conversationId: threadId, + parentMessageId, + sender: 'User', + text: request.input, + endpoint: EModelEndpoint.agents, + isCreatedByUser: true, + subagentTask: { + attemptKey, + parentRunId: request.parentRunId, + ...(requestFingerprint == null ? {} : { requestFingerprint }), + status: 'running', + }, + ...retentionFields(conversation), }, - ...retentionFields(conversation), - }, - { context: 'SubagentThreadTaskStore.prepareThread' }, + { context: 'SubagentThreadTaskStore.prepareThread' }, + ), + { stage: 'seed_write', taskId, threadId }, ); if (savedUserMessage == null) { throw new Error('Unable to persist the child-thread input.'); From f431b01d1f14a630f8fcf959cba440a3d72c2041 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 19 Aug 2026 15:48:14 -0400 Subject: [PATCH 14/15] =?UTF-8?q?=F0=9F=9A=9A=20chore:=20Bump=20`@librecha?= =?UTF-8?q?t/agents`=20to=20v3.6.8=20For=20Token,=20Trace,=20and=20Stream?= =?UTF-8?q?=20Fixes=20(#15012)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: bump `librechat/agents` to v3.6.7 * v3.6.8 --- api/package.json | 2 +- package-lock.json | 50 +++++++++++++++++++-------------------- packages/api/package.json | 2 +- 3 files changed, 27 insertions(+), 27 deletions(-) diff --git a/api/package.json b/api/package.json index 4ffb1a52375..c625b46fca2 100644 --- a/api/package.json +++ b/api/package.json @@ -46,7 +46,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.6", + "@librechat/agents": "^3.6.8", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", diff --git a/package-lock.json b/package-lock.json index 15d6c5225a4..7f4c6826a7d 100644 --- a/package-lock.json +++ b/package-lock.json @@ -63,7 +63,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.6", + "@librechat/agents": "^3.6.8", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", @@ -10521,22 +10521,22 @@ } }, "node_modules/@langfuse/core": { - "version": "5.9.1", - "resolved": "https://registry.npmjs.org/@langfuse/core/-/core-5.9.1.tgz", - "integrity": "sha512-KvyAskAO+2ixJwr9wy148ttR/Zn3oapvOJ2Br4Xcp6zhDpjSIIGB4jW/jkp53/KU9MyEHj0RSFPK/yKoxLC4KA==", + "version": "5.10.1", + "resolved": "https://registry.npmjs.org/@langfuse/core/-/core-5.10.1.tgz", + "integrity": "sha512-W8UArizWSy1DdeLGTsTwJwl7bkA7OQQcGZW8RtoopXyJZ93O0rwG7wzzeiZjhjpj5OtWOUTEaJuNkwOrF31UDw==", "license": "MIT", "peerDependencies": { "@opentelemetry/api": "^1.9.0" } }, "node_modules/@langfuse/langchain": { - "version": "5.9.1", - "resolved": "https://registry.npmjs.org/@langfuse/langchain/-/langchain-5.9.1.tgz", - "integrity": "sha512-Qv5Wn8EO2cRiVTBlb/zQwRhjD+93rVjD02in9L0+hyg9OBvc2rzqTfr8zKVDv9pgpsDfWD7Wm/sXyo7vGZmsTA==", + "version": "5.10.1", + "resolved": "https://registry.npmjs.org/@langfuse/langchain/-/langchain-5.10.1.tgz", + "integrity": "sha512-roKCdlyTmBVw1mT91yz3TUy+7xnvuBD1FaQqb6eR4H7/U8l40UGThP3c1wPKUOfIO57EGa9A/YwjGoc7YC2AIw==", "license": "MIT", "dependencies": { - "@langfuse/core": "^5.9.1", - "@langfuse/tracing": "^5.9.1" + "@langfuse/core": "^5.10.1", + "@langfuse/tracing": "^5.10.1" }, "peerDependencies": { "@langchain/core": ">=0.3.8", @@ -10544,12 +10544,12 @@ } }, "node_modules/@langfuse/otel": { - "version": "5.9.1", - "resolved": "https://registry.npmjs.org/@langfuse/otel/-/otel-5.9.1.tgz", - "integrity": "sha512-viM5Qq/AIPZPXfO7YdSmDHxEDSrNU/MyGzuE9zKA6hGBII1iLowU+qL99O+TjnXqxoADqlNibhY2pg1/WTIPcw==", + "version": "5.10.1", + "resolved": "https://registry.npmjs.org/@langfuse/otel/-/otel-5.10.1.tgz", + "integrity": "sha512-F2153e4PoJ1cN+5tM/xnsS44aQCQwK3p0nPk4NEpITV5pMTqiQVyvpkAvly8GKQ5Qjjr7heJ1dFtghW43ysyPQ==", "license": "MIT", "dependencies": { - "@langfuse/core": "^5.9.1" + "@langfuse/core": "^5.10.1" }, "engines": { "node": ">=20" @@ -10562,12 +10562,12 @@ } }, "node_modules/@langfuse/tracing": { - "version": "5.9.1", - "resolved": "https://registry.npmjs.org/@langfuse/tracing/-/tracing-5.9.1.tgz", - "integrity": "sha512-tJRyVAv1JkuOPh4Uz5eWUNH8U4jcJVtK2F5QNy5cZUzXCSrXobCSHusPbxY6VFZcLcFgtpDtSaxL7ev1tV2JNQ==", + "version": "5.10.1", + "resolved": "https://registry.npmjs.org/@langfuse/tracing/-/tracing-5.10.1.tgz", + "integrity": "sha512-m2kK4D0MsH8g4Og6KpnlYk8NLdQTYe0JR5M4KKpfNj99XXLlbdpXE/g3uJSqkcrWFhpiIb+3cyS9+uV6wQ6WtA==", "license": "MIT", "dependencies": { - "@langfuse/core": "^5.9.1" + "@langfuse/core": "^5.10.1" }, "engines": { "node": ">=20" @@ -10628,9 +10628,9 @@ } }, "node_modules/@librechat/agents": { - "version": "3.6.6", - "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.6.tgz", - "integrity": "sha512-YJqZ4Dsw/+dIu/Kbp3oMAc36zFXIrecHCcxYHdlFZA7WNnBLlBk2RwDSvf2WZUP0KfPACdx5K5LZOhTG4s+7ZQ==", + "version": "3.6.8", + "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.8.tgz", + "integrity": "sha512-FiPf8ggOoKJp7hpGKnk6dCtMITN1mzTqld3W9mOmGKgRAKp6PqCKiryseTdL0+5B6Mj1zBBx98GGLSXyxt2Ybg==", "license": "MIT", "dependencies": { "@anthropic-ai/sdk": "^0.115.0", @@ -10648,10 +10648,10 @@ "@langchain/openai": "1.5.8", "@langchain/textsplitters": "^1.0.1", "@langchain/xai": "^1.4.3", - "@langfuse/core": "^5.4.1", - "@langfuse/langchain": "^5.4.1", - "@langfuse/otel": "^5.4.1", - "@langfuse/tracing": "^5.4.1", + "@langfuse/core": "^5.10.1", + "@langfuse/langchain": "^5.10.1", + "@langfuse/otel": "^5.10.1", + "@langfuse/tracing": "^5.10.1", "@opentelemetry/context-async-hooks": "^2.9.0", "@opentelemetry/sdk-node": "^0.220.0", "@types/diff": "^7.0.2", @@ -42875,7 +42875,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.6", + "@librechat/agents": "^3.6.8", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/package.json b/packages/api/package.json index 4f9e0319fa7..210f8572ad7 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -113,7 +113,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.6", + "@librechat/agents": "^3.6.8", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", From 16e4d1419107a31c00fa23442ef07688e64d6366 Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:49:47 +0200 Subject: [PATCH 15/15] =?UTF-8?q?=E2=9C=A8=20refactor:=20Presets,=20Skills?= =?UTF-8?q?=20Motion=20and=20Model=20Selector=20Polish=20(#14953)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor: presets, skills motion and model selector polish Four surfaces that had drifted from the rest of the app, plus the CI fragility that surfaced while getting them green. Two were functional bugs rather than styling: Keyboard focus was invisible in the model selector. The highlight rule existed and the background was painted, but it used surface-secondary and the menu sits on bg-presentation, which resolve to the same value in dark and to within 3/255 in light, so only the thin indicator bar ever showed. Keyboard focus now uses the same surface a pointer gets. Importing a malformed preset raised com_ui_upload_invalid, which talks about image size limits, and FileUpload's JSON.parse had nothing catching it at that call site. The overflow menu owns the input and reports the existing preset import error instead. The rest is polish: preset surfaces use the theme radius roles rather than raw values; the edit dialog stops nesting a fixed 350px scroll box inside an already scrolling dialog and pins its title and actions, with the endpoint picker moved to ControlCombobox and kept out of any clipping ancestor; Clear all and Import move into a three-dots menu matching the conversation row; the Skills sections and pinned chats adopt the Collapse that Projects already used; the rendered/source toggle slides between states, is extracted rather than duplicated, and gains the accessible name and RTL mirroring it lacked; the header toggle loses its fill and the mobile new chat button hides when you are already in a new chat. The CI changes are unrelated to the UI but blocked it: the MCP and Redis cache jobs installed Redis with a bare apt-get and lost a race against the runner's own apt-daily work, failing four times and once hanging for 30 minutes. They now stop that background work and wait for the lock. DPkg::Lock::Timeout alone does not help, since it covers the dpkg frontend lock and not the lists lock. * refactor: move the section label appearance into the Label primitive The preset dialog reached into the agent panel's private `Advanced/ui` for its field eyebrow, so an agent-only refactor could change the dialog. Give the shared `Label` a `section` variant and export the recipe for the agent id row, which heads its value on a span and must not inherit the label's block layout. Each variant carries its own size, leading and color: the recipe output reaches that span unmerged, and a font size declared after `leading-none` drops it. * fix: derive the mobile new chat action from the route The context conversation still holds the previous chat for a render after a history or link navigation, a lag ChatView already guards against, so the action could show on /c/new or hide while an existing chat loaded. * style: sort imports in the touched files * fix: return focus to the menu item after the clear dialog The dialog is controlled and has no trigger, so Radix restored focus to whatever held it when the content mounted, the menu's own focus trap, and a keyboard user was left on the document. The menu stays open behind the dialog, so the invoking item is still there to take focus back. * fix: fall back to the trigger when clearing removes the invoking item Confirming empties the presets optimistically, so React commits the removed menu item together with the dialog close and the saved invoker is already disconnected when focus is handed back. --------- Co-authored-by: Danny Avila --- .github/workflows/cache-integration-tests.yml | 24 ++- client/src/components/Chat/Header.tsx | 17 +- .../Chat/Menus/Endpoints/CustomMenu.tsx | 13 +- .../Chat/Menus/Presets/EditPresetDialog.tsx | 121 ++++++------ .../Chat/Menus/Presets/PresetItems.tsx | 180 +++++++++++++----- .../Presets/__tests__/PresetItems.spec.tsx | 118 ++++++++++++ .../src/components/Chat/Menus/PresetsMenu.tsx | 4 +- .../Conversations/PinnedSection.tsx | 5 +- .../__tests__/PinnedSection.spec.tsx | 4 +- .../Endpoints/SaveAsPresetDialog.tsx | 2 +- .../Agents/Advanced/AdvancedPanel.tsx | 8 +- .../SidePanel/Agents/Advanced/ui.tsx | 3 - .../components/Skills/display/SkillDetail.tsx | 52 +---- .../Skills/display/SkillFileViewer.tsx | 38 +--- .../components/Skills/display/ViewToggle.tsx | 63 ++++++ .../src/components/Skills/lists/SkillList.tsx | 5 +- .../components/Skills/lists/SkillListItem.tsx | 14 +- .../__tests__/SkillsSidePanel.spec.tsx | 7 +- .../client/src/components/Button.spec.tsx | 4 +- packages/client/src/components/Button.tsx | 2 +- packages/client/src/components/Label.spec.tsx | 49 +++++ packages/client/src/components/Label.tsx | 38 +++- 22 files changed, 544 insertions(+), 227 deletions(-) create mode 100644 client/src/components/Chat/Menus/Presets/__tests__/PresetItems.spec.tsx create mode 100644 client/src/components/Skills/display/ViewToggle.tsx create mode 100644 packages/client/src/components/Label.spec.tsx diff --git a/.github/workflows/cache-integration-tests.yml b/.github/workflows/cache-integration-tests.yml index 7fda84602ac..9634569cf5b 100644 --- a/.github/workflows/cache-integration-tests.yml +++ b/.github/workflows/cache-integration-tests.yml @@ -45,9 +45,29 @@ jobs: node-version: '24.16.0' - name: Install Redis tools + timeout-minutes: 10 run: | - sudo apt-get update - sudo apt-get install -y redis-server redis-tools + # Same runner apt contention that broke the MCP job in + # playwright-mock.yml: apt-daily/unattended-upgrades hold + # /var/lib/apt/lists/lock at boot. Without a step timeout this hung + # until the job-level one fired, taking the whole leg with it. + sudo systemctl stop apt-daily.service apt-daily-upgrade.service \ + unattended-upgrades.service 2>/dev/null || true + sudo systemctl kill --kill-who=all apt-daily.service \ + apt-daily-upgrade.service 2>/dev/null || true + + apt_with_lock_wait() { + for attempt in $(seq 1 30); do + if sudo apt-get -o DPkg::Lock::Timeout=60 "$@"; then + return 0 + fi + echo "apt-get $1 could not take the lock (attempt ${attempt}/30), retrying" + sleep 10 + done + return 1 + } + apt_with_lock_wait update + apt_with_lock_wait install -y redis-server redis-tools - name: Start Single Redis Instance run: | diff --git a/client/src/components/Chat/Header.tsx b/client/src/components/Chat/Header.tsx index 454b872e0f7..16a73ce6000 100644 --- a/client/src/components/Chat/Header.tsx +++ b/client/src/components/Chat/Header.tsx @@ -1,6 +1,12 @@ import { memo, useMemo } from 'react'; import { useRecoilValue } from 'recoil'; -import { getConfigDefaults, PermissionTypes, Permissions } from 'librechat-data-provider'; +import { useParams } from 'react-router-dom'; +import { + getConfigDefaults, + Constants, + PermissionTypes, + Permissions, +} from 'librechat-data-provider'; import { OpenSidebar, PresetsMenu, NewChat, HeaderMenu } from './Menus'; import ModelSelector from './Menus/Endpoints/ModelSelector'; import { useGetStartupConfig } from '~/data-provider'; @@ -31,6 +37,13 @@ function Header({ const { data: startupConfig } = useGetStartupConfig(); const navVisible = useRecoilValue(store.sidebarExpanded); + /** The mobile row only offers a new chat when there is one to leave. Read + * from the route rather than the context conversation, which still holds the + * previous chat for a render after a history or link navigation. An unsaved + * conversation has no id in the route yet, so absence counts as new too. */ + const { conversationId: routeConversationId } = useParams(); + const isNewChat = routeConversationId == null || routeConversationId === Constants.NEW_CONVO; + const interfaceConfig = useMemo( () => startupConfig?.interface ?? defaultInterface, [startupConfig], @@ -90,7 +103,7 @@ function Header({
- + {!isNewChat && }
diff --git a/client/src/components/Chat/Menus/Endpoints/CustomMenu.tsx b/client/src/components/Chat/Menus/Endpoints/CustomMenu.tsx index 06db64662c9..1190a113deb 100644 --- a/client/src/components/Chat/Menus/Endpoints/CustomMenu.tsx +++ b/client/src/components/Chat/Menus/Endpoints/CustomMenu.tsx @@ -42,9 +42,12 @@ export const CustomMenu = React.forwardRef(func const rootMenuStateClass = isOpen ? 'bg-surface-active-alt hover:bg-surface-active-alt' : 'bg-presentation hover:bg-surface-active-alt'; + /** Nested triggers sit on the popover, whose bg-presentation resolves to the + * same value as surface-secondary in dark and within 3/255 of it in light, + * so highlighting with it leaves keyboard focus invisible. */ const nestedMenuStateClass = isOpen - ? 'bg-surface-secondary hover:bg-surface-hover data-[active-item]:bg-surface-secondary data-[active-item]:hover:bg-surface-hover' - : 'hover:bg-surface-hover data-[active-item]:bg-surface-secondary data-[active-item]:hover:bg-surface-hover'; + ? 'bg-surface-hover' + : 'hover:bg-surface-hover data-[active-item]:bg-surface-hover'; const element = ( @@ -172,7 +175,11 @@ export const CustomMenuItem = React.forwardRef !isAgentsEndpoint(endpoint)); }, [_endpoints]); + const endpointItems = useMemo( + () => + availableEndpoints.map((value) => ({ + value, + label: alternateName[value] ?? value, + })), + [availableEndpoints], + ); + useEffect(() => { if (!preset) { return; @@ -133,45 +136,51 @@ const EditPresetDialog = ({ return ( - - + + {localize('com_ui_edit_preset_title', { title: preset?.title })} -
- {/* Header section with preset name and endpoint */} -
-
- - -
-
- - -
+ {/* Pinned above the scroller, and the dialog itself is overflow-visible: + ControlCombobox renders its popover in place (portal={false} for the + dialog's focus trap), so no ancestor may clip it. The flex column + still bounds the dialog because the settings region below owns the + only scroll. */} +
+
+ +
+
+ + +
+
+ {/* Only this region scrolls, so the title, the fields above and the actions stay put */} +
{/* PopoverButtons section */}
- {/* Settings section */} -
+ {/* Settings section. The shared component ships a fixed-height scroll + box; overriding it to auto lets the dialog own the single scroll + rather than nesting one inside another. */} +
+
- {/* Action buttons */} -
- - -
+ {/* Action buttons */} +
+ +
diff --git a/client/src/components/Chat/Menus/Presets/PresetItems.tsx b/client/src/components/Chat/Menus/Presets/PresetItems.tsx index 98f14f90e23..b7f426e1a0e 100644 --- a/client/src/components/Chat/Menus/Presets/PresetItems.tsx +++ b/client/src/components/Chat/Menus/Presets/PresetItems.tsx @@ -1,14 +1,18 @@ +import { useRef, useState } from 'react'; import { useRecoilValue } from 'recoil'; +import * as Ariakit from '@ariakit/react'; import { Close } from '@radix-ui/react-popover'; -import { BookCopy, FileX2 } from 'lucide-react'; import { Flipper, Flipped } from 'react-flip-toolkit'; import { getEndpointField } from 'librechat-data-provider'; +import { BookCopy, FileUp, FileX2, Ellipsis } from 'lucide-react'; import { Button, PinIcon, EditIcon, TrashIcon, + DropdownPopup, TooltipAnchor, + useToastContext, AlertDialog, AlertDialogAction, AlertDialogCancel, @@ -17,11 +21,10 @@ import { AlertDialogFooter, AlertDialogHeader, AlertDialogTitle, - AlertDialogTrigger, } from '@librechat/client'; +import type { MenuItemProps } from '@librechat/client'; import type { TPreset } from 'librechat-data-provider'; -import type { FC } from 'react'; -import FileUpload from '~/components/Chat/Input/Files/FileUpload'; +import type { ChangeEvent, FC } from 'react'; import { useGetEndpointsQuery } from '~/data-provider'; import { getPresetTitle, getIconKey } from '~/utils'; import { icons } from '~/hooks/Endpoint/Icons'; @@ -30,6 +33,9 @@ import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; import store from '~/store'; +/** Shared by the trigger and the clear dialog's focus fallback. */ +const PRESET_MENU_ID = 'preset-options-button'; + const PresetItems: FC<{ presets?: Array; onSetDefaultPreset: (preset: TPreset, remove?: boolean) => void; @@ -52,7 +58,54 @@ const PresetItems: FC<{ const { data: endpointsConfig } = useGetEndpointsQuery(); const defaultPreset = useRecoilValue(store.defaultPreset); const localize = useLocalize(); + const { showToast } = useToastContext(); const hasPresets = (presets?.length ?? 0) > 0; + const [isMenuOpen, setIsMenuOpen] = useState(false); + const [isClearDialogOpen, setIsClearDialogOpen] = useState(false); + const importInputRef = useRef(null); + /** Radix restores focus to whatever held it when the dialog mounted, which by + * then is the menu's own focus trap rather than the item that opened it. */ + const clearInvokerRef = useRef(null); + + const handleImportChange = (event: ChangeEvent) => { + const file = event.target.files?.[0]; + /** Cleared so re-picking the same file still fires a change event */ + event.target.value = ''; + if (!file) { + return; + } + + const reader = new FileReader(); + reader.onload = (e) => { + try { + onFileSelected(JSON.parse(e.target?.result as string)); + } catch { + showToast({ message: localize('com_endpoint_preset_import_error'), status: 'error' }); + } + }; + reader.readAsText(file); + }; + + const menuItems: MenuItemProps[] = [ + { + label: localize('com_ui_import'), + onClick: () => importInputRef.current?.click(), + icon:
-
- {hasPresets && ( - - - - - - - {localize('com_ui_clear_presets')} - - {localize('com_endpoint_presets_clear_warning')} - - - - {localize('com_ui_cancel')} - - {localize('com_ui_clear')} - - - - - )} - -
+ +
+ + + + + { + const saved = clearInvokerRef.current; + clearInvokerRef.current = null; + /** Confirming removes the item itself, since it only shows while + * presets exist, so fall back to the trigger that opened the menu. */ + const invoker = + saved?.isConnected === true ? saved : document.getElementById(PRESET_MENU_ID); + if (invoker == null) { + return; + } + event.preventDefault(); + invoker.focus(); + }} + className="w-11/12 max-w-md rounded-theme-surface sm:rounded-theme-surface" + > + + {localize('com_ui_clear_presets')} + + {localize('com_endpoint_presets_clear_warning')} + + + + {localize('com_ui_cancel')} + + {localize('com_ui_clear')} + + + + {presets && presets.length === 0 && (
@@ -152,11 +234,11 @@ const PresetItems: FC<{
-
+
@@ -88,7 +88,7 @@ const PresetsMenu: FC = () => { sideOffset={8} collisionPadding={16} aria-label={localize('com_endpoint_examples')} - className="z-50 max-h-[495px] overflow-x-hidden rounded-lg border border-border-light bg-presentation text-text-primary shadow-lg md:min-w-[400px]" + className="z-50 max-h-[495px] overflow-x-hidden rounded-theme-surface border border-border-light bg-presentation text-text-primary shadow-lg md:min-w-[400px]" > {
- {isExpanded && ( +
    {conversations.map((convo) => ( @@ -65,7 +66,7 @@ const PinnedSection = ({ conversations, toggleNav }: PinnedSectionProps) => { ))}
- )} +
); }; diff --git a/client/src/components/Conversations/__tests__/PinnedSection.spec.tsx b/client/src/components/Conversations/__tests__/PinnedSection.spec.tsx index b7460df8b36..304bd183774 100644 --- a/client/src/components/Conversations/__tests__/PinnedSection.spec.tsx +++ b/client/src/components/Conversations/__tests__/PinnedSection.spec.tsx @@ -74,7 +74,9 @@ describe('PinnedSection', () => { 'aria-expanded', 'false', ); - expect(screen.queryByText('Pinned Chat')).not.toBeInTheDocument(); + /** Collapse keeps children mounted so the height can tween, and hides them + * from assistive tech instead, the same as ProjectsSection above it. */ + expect(screen.getByText('Pinned Chat').closest('[aria-hidden="true"]')).not.toBeNull(); }); it('toggles the section when the header is clicked', () => { diff --git a/client/src/components/Endpoints/SaveAsPresetDialog.tsx b/client/src/components/Endpoints/SaveAsPresetDialog.tsx index 61d246a4fa1..fa9acd65038 100644 --- a/client/src/components/Endpoints/SaveAsPresetDialog.tsx +++ b/client/src/components/Endpoints/SaveAsPresetDialog.tsx @@ -77,7 +77,7 @@ const SaveAsPresetDialog = ({ open, onOpenChange, preset }: TEditPresetProps) => value={title || ''} onChange={(e) => setTitle(e.target.value || '')} placeholder={localize('com_endpoint_preset_custom_name_placeholder')} - className="flex h-10 max-h-10 w-full resize-none border-border-medium px-3 py-2" + className="flex h-10 max-h-10 w-full resize-none rounded-theme-control border-border-medium px-3 py-2" />
diff --git a/client/src/components/SidePanel/Agents/Advanced/AdvancedPanel.tsx b/client/src/components/SidePanel/Agents/Advanced/AdvancedPanel.tsx index 78041c1ce3e..85588b76f61 100644 --- a/client/src/components/SidePanel/Agents/Advanced/AdvancedPanel.tsx +++ b/client/src/components/SidePanel/Agents/Advanced/AdvancedPanel.tsx @@ -2,13 +2,13 @@ import { useMemo, useState } from 'react'; import { useFormContext } from 'react-hook-form'; import { ChevronLeft, Check, Copy } from 'lucide-react'; import { AgentCapabilities } from 'librechat-data-provider'; -import { Button, TooltipAnchor, useToastContext } from '@librechat/client'; +import { Button, TooltipAnchor, labelVariants, useToastContext } from '@librechat/client'; import type { AgentForm } from '~/common'; -import { sectionLabelClass, groupHeadingClass } from './ui'; import { useAgentPanelContext } from '~/Providers'; import StatefulSessions from './StatefulSessions'; import OrchestrationHub from './OrchestrationHub'; import MaxAgentSteps from './MaxAgentSteps'; +import { groupHeadingClass } from './ui'; import { useLocalize } from '~/hooks'; import { Panel } from '~/common'; @@ -66,7 +66,9 @@ export default function AdvancedPanel() { {currentAgentId && (
- {localize('com_ui_agent_id')} + + {localize('com_ui_agent_id')} + void; - localize: ReturnType; -}) { - return ( -
- - -
- ); -} - export default function SkillDetail({ skill, onEdit, onDelete }: SkillDetailProps) { const localize = useLocalize(); const { user } = useAuthContext(); @@ -159,7 +113,7 @@ export default function SkillDetail({ skill, onEdit, onDelete }: SkillDetailProp {/* Divider with view toggle */}

- +
{/* Frontmatter metadata */} diff --git a/client/src/components/Skills/display/SkillFileViewer.tsx b/client/src/components/Skills/display/SkillFileViewer.tsx index fd195eb9207..281c14fb252 100644 --- a/client/src/components/Skills/display/SkillFileViewer.tsx +++ b/client/src/components/Skills/display/SkillFileViewer.tsx @@ -2,12 +2,12 @@ import React, { memo, useMemo, useState, useCallback, useRef } from 'react'; import { useNavigate } from 'react-router-dom'; import { apiBaseUrl } from 'librechat-data-provider'; import { Spinner, TooltipAnchor, useToastContext } from '@librechat/client'; -import { ArrowLeft, Eye, Code, Copy, Check, FileText, FileQuestion } from 'lucide-react'; +import { ArrowLeft, Copy, Check, FileText, FileQuestion } from 'lucide-react'; import { useGetSkillFileContentQuery } from '~/data-provider'; import SkillMarkdownRenderer from './SkillMarkdownRenderer'; import { parseFrontmatter } from '../utils'; +import ViewToggle from './ViewToggle'; import { useLocalize } from '~/hooks'; -import { cn } from '~/utils'; interface SkillFileViewerProps { skillId: string; @@ -97,39 +97,7 @@ function SkillFileViewer({ skillId, relativePath }: SkillFileViewerProps) { )} {/* View toggle (markdown only) */} - {isMarkdown && isText && ( -
- - -
- )} + {isMarkdown && isText && }
diff --git a/client/src/components/Skills/display/ViewToggle.tsx b/client/src/components/Skills/display/ViewToggle.tsx new file mode 100644 index 00000000000..f635400609e --- /dev/null +++ b/client/src/components/Skills/display/ViewToggle.tsx @@ -0,0 +1,63 @@ +import { Eye, Code } from 'lucide-react'; +import type { LucideIcon } from 'lucide-react'; +import type { TranslationKeys } from '~/hooks'; +import { useLocalize } from '~/hooks'; +import { cn } from '~/utils'; + +export type SkillViewMode = 'rendered' | 'source'; + +interface ViewToggleProps { + viewMode: SkillViewMode; + setViewMode: (mode: SkillViewMode) => void; +} + +const MODES: ReadonlyArray<{ mode: SkillViewMode; Icon: LucideIcon; labelKey: TranslationKeys }> = [ + { mode: 'rendered', Icon: Eye, labelKey: 'com_ui_skill_view_rendered' }, + { mode: 'source', Icon: Code, labelKey: 'com_ui_skill_view_source' }, +]; + +/** + * Segmented control for the rendered/source swap. + * + * The active state is one thumb that slides between the options rather than a + * background appearing on one button as it disappears from the other, so the + * change reads as a single movement. Option widths are fixed so the thumb can + * travel by exactly one option without measuring. + */ +export default function ViewToggle({ viewMode, setViewMode }: ViewToggleProps) { + const localize = useLocalize(); + + return ( +
+
+ ); +} diff --git a/client/src/components/Skills/lists/SkillList.tsx b/client/src/components/Skills/lists/SkillList.tsx index 745589e4bc7..a7523b766c5 100644 --- a/client/src/components/Skills/lists/SkillList.tsx +++ b/client/src/components/Skills/lists/SkillList.tsx @@ -3,6 +3,7 @@ import { ChevronRight } from 'lucide-react'; import { useSearchParams } from 'react-router-dom'; import type { TSkillSummary } from 'librechat-data-provider'; import SkillListItem from './SkillListItem'; +import { Collapse } from '~/components/ui'; import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; @@ -47,7 +48,7 @@ export default function SkillList({
{/* Skill items */} - {sectionOpen && ( +
{skills.length === 0 ? (

@@ -66,7 +67,7 @@ export default function SkillList({ )) )}

- )} +
); } diff --git a/client/src/components/Skills/lists/SkillListItem.tsx b/client/src/components/Skills/lists/SkillListItem.tsx index 22eaa0f2050..bf97bfc7a91 100644 --- a/client/src/components/Skills/lists/SkillListItem.tsx +++ b/client/src/components/Skills/lists/SkillListItem.tsx @@ -5,6 +5,7 @@ import { ScrollText, ChevronDown, ChevronRight, Folder, Pin } from 'lucide-react import type { FixedSizeNodeData, TreeWalkerValue, TreeWalker } from 'react-vtree'; import type { TSkillSummary, TSkillFile } from 'librechat-data-provider'; import { useListSkillFilesQuery } from '~/data-provider'; +import { Collapse } from '~/components/ui'; import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; @@ -323,7 +324,7 @@ function SkillListItem({ - {skill.name} + {skill.name} {skill.alwaysApply === true && ( {/* Inline file tree */} -
+ -
+
); } diff --git a/client/src/components/Skills/sidebar/__tests__/SkillsSidePanel.spec.tsx b/client/src/components/Skills/sidebar/__tests__/SkillsSidePanel.spec.tsx index d9bb2c390b2..7d8f4788f6a 100644 --- a/client/src/components/Skills/sidebar/__tests__/SkillsSidePanel.spec.tsx +++ b/client/src/components/Skills/sidebar/__tests__/SkillsSidePanel.spec.tsx @@ -41,7 +41,12 @@ jest.mock('~/components/ui', () => { const PanelContent = ReactModule.forwardRef( ({ children }, ref) =>
{children}
, ); - return { PanelContent }; + /** SkillList renders its body through Collapse, which keeps children mounted + * and marks them hidden when closed rather than unmounting them. */ + const Collapse = ({ open, children }: { open: boolean; children?: React.ReactNode }) => ( +
{children}
+ ); + return { PanelContent, Collapse }; }); jest.mock('../FilterSkills', () => ({ diff --git a/packages/client/src/components/Button.spec.tsx b/packages/client/src/components/Button.spec.tsx index 74cd664fa86..67cfe5b0c45 100644 --- a/packages/client/src/components/Button.spec.tsx +++ b/packages/client/src/components/Button.spec.tsx @@ -21,8 +21,10 @@ describe('Button', () => { it('renders the header-action toggle from semantic tokens', () => { render(); + /** Transparent so the toggle reads as an icon on the header rather than a + * raised control; the border and hover still mark it as hit-able. */ expect(screen.getByRole('button', { name: 'Toggle' })).toHaveClass( - 'bg-presentation', + 'bg-transparent', 'border-border-light', 'rounded-xl', 'duration-0', diff --git a/packages/client/src/components/Button.tsx b/packages/client/src/components/Button.tsx index abde7d06892..49a8e082e9e 100644 --- a/packages/client/src/components/Button.tsx +++ b/packages/client/src/components/Button.tsx @@ -70,7 +70,7 @@ const buttonVariantRecipe = cva( * lag rather than polish. */ 'header-action': - 'rounded-xl border border-border-light bg-presentation text-text-primary duration-0 hover:bg-surface-active-alt hover:text-text-primary', + 'rounded-xl border border-border-light bg-transparent text-text-primary duration-0 hover:bg-surface-active-alt hover:text-text-primary', }, size: { default: 'h-10 px-4 py-2', diff --git a/packages/client/src/components/Label.spec.tsx b/packages/client/src/components/Label.spec.tsx new file mode 100644 index 00000000000..d70a971e449 --- /dev/null +++ b/packages/client/src/components/Label.spec.tsx @@ -0,0 +1,49 @@ +import '@testing-library/jest-dom'; +import { render, screen } from '@testing-library/react'; +import { Label, labelVariants } from './Label'; + +describe('Label', () => { + it('keeps the default appearance when no variant is selected', () => { + render(); + const label = screen.getByText('Name'); + + expect(label).toHaveClass('block', 'w-full', 'break-all', 'leading-none', 'text-sm'); + expect(label).toHaveClass('text-text-primary', 'peer-disabled:opacity-70'); + }); + + it('renders the section eyebrow from the shared variant', () => { + render( + , + ); + const label = screen.getByText('Endpoint'); + + expect(label).toHaveClass( + 'text-[11px]', + 'font-medium', + 'uppercase', + 'tracking-wide', + 'text-text-secondary', + ); + /** The variant owns size, leading and color outright: an arbitrary font size + * also clears `leading-none`, which is what the label read before. */ + expect(label).not.toHaveClass('text-sm', 'text-text-primary', 'leading-none'); + }); + + /** + * A settings row heads its value with this appearance on a non-label element, + * so the recipe has to stay free of the label's block layout: `block w-full` + * would break the row's `justify-between`. + */ + it('exposes the eyebrow to non-label elements without layout', () => { + const section = labelVariants({ variant: 'section' }); + + expect(section).toContain('text-[11px]'); + expect(section).toContain('text-text-secondary'); + expect(section).not.toContain('block'); + expect(section).not.toContain('w-full'); + /** Unmerged recipe output, so a conflicting base color would survive it. */ + expect(section).not.toContain('text-text-primary'); + }); +}); diff --git a/packages/client/src/components/Label.tsx b/packages/client/src/components/Label.tsx index 26c350b1098..04bb94e52be 100644 --- a/packages/client/src/components/Label.tsx +++ b/packages/client/src/components/Label.tsx @@ -1,23 +1,51 @@ import * as React from 'react'; import * as LabelPrimitive from '@radix-ui/react-label'; +import { ClassProp } from 'class-variance-authority/types'; +import { cva, type VariantProps } from 'class-variance-authority'; import { cn } from '~/utils'; +type LabelVariantOptions = + | ({ variant?: 'default' | 'section' | null | undefined } & ClassProp) + | undefined; + +/** + * Typography only, so a non-label element that heads a settings row can reuse a + * variant without inheriting the label's block layout. Each variant carries its + * own size, leading and color rather than overriding a shared base: the raw + * recipe output is not merged for those consumers, and a font size declared + * after `leading-none` would drop it. + */ +const labelVariants: (props?: LabelVariantOptions) => string = cva('', { + variants: { + variant: { + default: 'text-sm leading-none text-text-primary', + /** Eyebrow above a field or settings group. */ + section: 'text-[11px] font-medium uppercase tracking-wide text-text-secondary', + }, + }, + defaultVariants: { + variant: 'default', + }, +}); + const Label: React.ForwardRefExoticComponent< Omit, 'ref'> & { className?: string; - } & React.RefAttributes + } & VariantProps & + React.RefAttributes > = React.forwardRef< React.ElementRef, React.ComponentPropsWithoutRef & { className?: string; - } ->(({ className = '', ...props }, ref) => ( + } & VariantProps +>(({ className = '', variant, ...props }, ref) => (