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/.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/.github/workflows/playwright-mock.yml b/.github/workflows/playwright-mock.yml index db3a275b255..5f984eb6bb3 100644 --- a/.github/workflows/playwright-mock.yml +++ b/.github/workflows/playwright-mock.yml @@ -251,18 +251,17 @@ 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 + # 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: | diff --git a/CONTEXT.md b/CONTEXT.md index 80f2c15602b..276b1f8ecab 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -2,4 +2,6 @@ - **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/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/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/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/__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/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/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/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/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/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/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/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..5f7ff966f6c 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -1,16 +1,34 @@ -const { createSubagentThreadTaskStore } = require('@librechat/api'); +const { + cacheConfig, + ioredisClient, + isEnabled, + registerShutdownTask, + duplicateIoRedisClient, + createSubagentThreadTaskStore, + createSubagentCompletionWakeupHandler, + RedisSubagentTaskControlTransport, +} = require('@librechat/api'); const db = require('~/models'); +const { enqueueAgentTrigger } = require('../../Agents/triggers'); -/** Durable logical threads use normal LibreChat conversations/messages. Live - * controls stay process-local; Mongo fences continuation across API replicas. */ +/** 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. */ const subagentThreadTaskStore = createSubagentThreadTaskStore( { acquireSubagentThreadLease: db.acquireSubagentThreadLease, + claimSubagentTaskResult: db.claimSubagentTaskResult, + releaseSubagentTaskResultClaim: db.releaseSubagentTaskResultClaim, 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 +36,52 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( saveMessage: db.saveMessage, }, { - isOwnerActive: db.isAgentTriggerPrincipalActive, + isOwnerActive: db.isSubagentOwnerAdmissible, + fenceOwnerAdmission: db.fenceSubagentAdmission, + renewOwnerAdmission: db.renewSubagentAdmission, + releaseOwnerAdmission: db.releaseSubagentAdmission, + ...(completionWakeupsEnabled && { + onTaskPrepared: createSubagentCompletionWakeupHandler(enqueueAgentTrigger), + }), }, ); +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/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() }; 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/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/Input/AskUserQuestionPopover.tsx b/client/src/components/Chat/Input/AskUserQuestionPopover.tsx index a855e572eee..d6568da1af3 100644 --- a/client/src/components/Chat/Input/AskUserQuestionPopover.tsx +++ b/client/src/components/Chat/Input/AskUserQuestionPopover.tsx @@ -156,7 +156,7 @@ function AskUserQuestionPopoverPanel({ scroll region: the panel is absolutely positioned, so anything that overflows it is unreachable by page scroll. */}
@@ -165,7 +165,7 @@ function AskUserQuestionPopoverPanel({ {liveAsk.question.question}

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

+

{liveAsk.question.description}

)} @@ -193,19 +193,23 @@ function AskUserQuestionPopoverPanel({ {options.map((option, index) => { const isChecked = multiSelect && checked.includes(index); return ( - + ); })}
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/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/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]" > ); } @@ -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/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__/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/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/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/Conversations/PinnedSection.tsx b/client/src/components/Conversations/PinnedSection.tsx index 717bf43ddf8..0bdcb95f697 100644 --- a/client/src/components/Conversations/PinnedSection.tsx +++ b/client/src/components/Conversations/PinnedSection.tsx @@ -3,6 +3,7 @@ import { ChevronDown } from 'lucide-react'; import type { TConversation } from 'librechat-data-provider'; import { useLocalize, useLocalStorage } from '~/hooks'; import { useActiveJobs } from '~/data-provider'; +import { Collapse } from '~/components/ui'; import { cn } from '~/utils'; import Convo from './Convo'; @@ -50,7 +51,7 @@ const PinnedSection = ({ conversations, toggleNav }: PinnedSectionProps) => {
- {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/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/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')} +
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/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/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/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; 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/locales/en/translation.json b/client/src/locales/en/translation.json index a9499bb3bc7..3f2405e8a53 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", @@ -1594,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", @@ -1657,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", @@ -2293,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/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(), 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") 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", 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..cd77218f36b 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -32,6 +32,8 @@ export * from './skills'; export * from './phases'; export * from './startup'; export * from './subagentThreads'; +export * from './subagentCompletionWakeup'; +export * from './subagentTaskRouting'; export * from './skillConfigurable'; export * from './skillFiles'; export * from './codeFilesSession'; 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.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..d4488776669 --- /dev/null +++ b/packages/api/src/agents/subagentTaskRouting.ts @@ -0,0 +1,1384 @@ +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 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') { + return claim; + } + const task = boundedSnapshot(claim.task); + if (claim.status === 'completed') { + 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) }; + } + 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..e513f12e2ae 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -2,17 +2,38 @@ 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, + 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 { SubagentTaskWakeupRegistration } from './subagentThreads'; 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 +41,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 +141,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 +183,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 { @@ -167,6 +293,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(); @@ -195,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(); @@ -334,22 +566,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( @@ -365,6 +610,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.', @@ -583,8 +840,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 +869,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 +887,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 +907,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 +1629,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..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, @@ -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,30 @@ 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 SLOW_PREPARATION_WARN_MS = 5_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 +76,13 @@ class SubagentThreadDeletedError extends SubagentThreadPublicError {} type SubagentThreadMethods = Pick< AllMethods, | 'acquireSubagentThreadLease' + | 'claimSubagentTaskResult' | 'countActiveSubagentThreadLeases' | 'deleteConvos' | 'deleteMessages' | 'getConvo' | 'getMessages' + | 'listActiveSubagentThreadLeases' | 'reserveSubagentThread' | 'releaseSubagentThreadLease' | 'renewSubagentThreadLease' @@ -68,10 +102,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; }; } @@ -88,6 +126,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 +139,26 @@ 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; + 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 { @@ -142,10 +201,22 @@ 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; } +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, @@ -284,21 +355,81 @@ 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. */ +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; 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 readonly onTaskPrepared?: SubagentThreadTaskStoreOptions['onTaskPrepared']; + private taskControlTransport?: SubagentTaskControlTransport; constructor( private readonly methods: SubagentThreadMethods, @@ -319,7 +450,38 @@ 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; + this.onTaskPrepared = options.onTaskPrepared; + } + + /** 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. */ @@ -373,107 +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[] = []; - try { - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - await parentReady; - const prepared = await this.prepareThread( - request.scopeId, - scope, - threadId, - isContinuation, - request, - runtime.taskId, - lease, - ); - 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 result = await runWithDetachedSubagentUsage(detachedUsage, () => - request.run(runtime, prepared.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) { - 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) { @@ -490,6 +674,300 @@ 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. + * + * 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, + kind: 'manual', + 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 +1022,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 +1171,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 +1441,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; + } + 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; } - return renewed; + shared.expiresAt = renewedUntil; + return true; } catch (error) { shared.lost = true; logger.warn('[subagentThreads] Lost the shared child-thread lease', error); @@ -779,18 +1588,31 @@ 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 * 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); @@ -811,13 +1633,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' @@ -846,6 +1684,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { error: true, subagentTask: { attemptKey, + parentRunId: request.parentRunId, ...(requestFingerprint == null ? {} : { requestFingerprint }), status: 'error', }, @@ -862,7 +1701,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); @@ -877,24 +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, - ...(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.'); @@ -908,6 +1760,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) { @@ -990,6 +1846,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ...(subagentTranscript == null ? {} : { subagentTranscript }), subagentTask: { attemptKey: prepared.attemptKey, + parentRunId: request.parentRunId, ...(normalizedRequestFingerprint(request) == null ? {} : { requestFingerprint: normalizedRequestFingerprint(request) }), @@ -1033,6 +1890,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { error: true, subagentTask: { attemptKey: createSubagentAttemptKey(request.scopeId, request.idempotencyKey), + parentRunId: request.parentRunId, ...(normalizedRequestFingerprint(request) == null ? {} : { requestFingerprint: normalizedRequestFingerprint(request) }), @@ -1049,6 +1907,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, @@ -1074,6 +1954,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { unfinished: false, subagentTask: { attemptKey: createSubagentAttemptKey(request.scopeId, request.idempotencyKey), + parentRunId: request.parentRunId, ...(normalizedRequestFingerprint(request) == null ? {} : { requestFingerprint: normalizedRequestFingerprint(request) }), @@ -1226,6 +2107,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 +2130,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 +2160,6 @@ export function buildSubagentThreadTaskConfig( ): SubagentTaskConfig { return { store, - scopeId: JSON.stringify({ version: SCOPE_VERSION, ...scope }), + scopeId: serializeScope(scope), }; } 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/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); } 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/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/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/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; } 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 1998ee8e26e..49a8e082e9e 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', @@ -59,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) => ( { + 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'; 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; 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..b94dd885286 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -28,6 +28,10 @@ let bulkSaveMessages: ReturnType['bulkSaveMessages' let updateMessageText: ReturnType['updateMessageText']; let deleteMessagesSince: ReturnType['deleteMessagesSince']; let recordMessage: ReturnType['recordMessage']; +let claimSubagentTaskResult: ReturnType['claimSubagentTaskResult']; +let releaseSubagentTaskResultClaim: ReturnType< + typeof createMessageMethods +>['releaseSubagentTaskResultClaim']; beforeAll(async () => { mongoServer = await MongoMemoryServer.create(); @@ -47,6 +51,8 @@ beforeAll(async () => { updateMessageText = methods.updateMessageText; deleteMessagesSince = methods.deleteMessagesSince; recordMessage = methods.recordMessage; + claimSubagentTaskResult = methods.claimSubagentTaskResult; + releaseSubagentTaskResultClaim = methods.releaseSubagentTaskResultClaim; await mongoose.connect(mongoUri); }); @@ -1549,4 +1555,207 @@ 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, + kind: 'manual', + 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, + 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, + 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 () => { + const runningTaskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(runningTaskId, conversationId, 'running'); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: runningTaskId, + kind: 'manual', + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: uuidv4(), + kind: 'manual', + 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, + 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 a864acc5766..500fb05ea60 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'; message: IMessage } + | { status: 'acquired'; message: IMessage }; + export interface MessageMethods { saveMessage( ctx: { userId: string; isTemporary?: boolean; interfaceConfig?: AppConfig['interfaceConfig'] }, @@ -82,6 +87,20 @@ export interface MessageMethods { message: Partial & { newMessageId?: string }, metadata?: { context?: string }, ): Promise>; + claimSubagentTaskResult(params: { + 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 }, @@ -518,6 +537,134 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } } + /** 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 ( + 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'); + } + const Message = mongoose.models.Message as Model; + 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( + { + 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.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; + } + /** * Deletes messages in a conversation since a specific message. */ @@ -655,6 +802,8 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateMessageText, updateToolCallResult, updateMessage, + claimSubagentTaskResult, + releaseSubagentTaskResultClaim, 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..ad1f59db9e6 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -138,12 +138,22 @@ const messageSchema: Schema = new Schema( subagentTask: { type: { attemptKey: { type: String, required: true }, + parentRunId: { type: String }, requestFingerprint: { type: String }, status: { type: String, enum: ['running', 'completed', 'error', 'cancelled'], required: true, }, + resultClaim: { + type: { + kind: { type: String, enum: ['manual', 'wakeup'], required: true }, + 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..bd3da444ee5 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -51,8 +51,15 @@ 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'; + resultClaim?: { + kind: 'manual' | 'wakeup'; + 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;