diff --git a/.github/workflows/playwright-mock.yml b/.github/workflows/playwright-mock.yml index db3a275b25..6c57b4fe10 100644 --- a/.github/workflows/playwright-mock.yml +++ b/.github/workflows/playwright-mock.yml @@ -251,12 +251,6 @@ jobs: continue-on-error: true run: timeout -k 10 90 npx playwright install ffmpeg - # Optional fonts only — see the note in the e2e_shards job. - - name: Install optional Playwright font dependencies (best effort) - timeout-minutes: 4 - continue-on-error: true - run: .github/scripts/install-playwright-fonts.sh - # Redis is a hard requirement for this job, so this step stays fatal. - name: Install Redis runtime dependencies timeout-minutes: 5 diff --git a/CONTEXT.md b/CONTEXT.md index 80f2c15602..6e411689c0 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -2,4 +2,5 @@ - **Agent run envelope**: the versioned, JSON-safe request contract created after ingress authentication and protocol validation but before agent, provider, tool, or MCP initialization. It carries only the validated protocol payload and the minimum trusted principal identifiers. The execution host rehydrates all runtime state from those identifiers. - **Subagent thread**: a durable, view-only child conversation owned by one parent conversation and subagent identity. A parent agent may continue it by stable `threadId`; each continuation uses a fresh execution lease restored from the canonical child transcript. It is not an ordinary human-writable chat. +- **Live subagent task owner**: the one API process holding a detached child execution, its abort controller, and its bounded control queue. Redis may route trusted poll/control envelopes to that owner, but it does not migrate or persist the executor; Mongo persists only the logical child thread and its continuation fence. - **Theme definition**: a versioned, data-only description of LibreChat semantic colors and shared appearance roles, optionally specialized by light or dark mode. The theme module validates and resolves partial definitions against bundled defaults before adapters apply them. A theme definition does not contain arbitrary CSS, application behavior, or alternate feature layouts. diff --git a/api/server/experimental.js b/api/server/experimental.js index 2d2ff5ec6f..698cb52ff2 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 c4364fb6b8..4ff9d6b924 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 f3cf7c67ec..0768a27c50 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 6461a148e3..013ece204e 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 d6ac3ca0da..5233d27ef0 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 ec5b23f263..c09a8ac075 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 87a9d0a719..fcba4f9a83 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -118,6 +118,26 @@ router.get('/gen_title/:conversationId', async (req, res) => { } }); +const POST_DELETE_CANCEL_ATTEMPTS = 3; +const POST_DELETE_CANCEL_BACKOFF_MS = 250; + +/** Replays a cancellation plan after deletion, retrying a transiently unreachable + * owner rather than losing the only pass that can stop a late-admitted child. */ +async function retryPostDeleteCancellation(cancellationPlan, deletedConversationIds) { + for (let attempt = 1; attempt <= POST_DELETE_CANCEL_ATTEMPTS; attempt += 1) { + try { + await subagentThreadTaskStore.cancelPlan(cancellationPlan, deletedConversationIds); + return; + } catch (error) { + if (attempt === POST_DELETE_CANCEL_ATTEMPTS) { + logger.warn('Post-delete subagent cancellation failed', error); + return; + } + await new Promise((resolve) => setTimeout(resolve, POST_DELETE_CANCEL_BACKOFF_MS * attempt)); + } + } +} + router.delete('/', configMiddleware, async (req, res) => { let filter = {}; const { conversationId, source, thread_id, endpoint } = req.body?.arg ?? {}; @@ -154,17 +174,36 @@ router.delete('/', configMiddleware, async (req, res) => { typeof req.user.tenantId === 'string' && req.user.tenantId !== '' ? req.user.tenantId : undefined; + let cancellationPlan; + let dbResponse; if (filter.conversationId) { - subagentThreadTaskStore.cancelForConversations( + /** Resolve the targets while the conversations still exist: the second pass + * runs after their rows are gone and can only reach registered owners. */ + cancellationPlan = await subagentThreadTaskStore.planCancellationForConversations( req.user.id, [filter.conversationId], tenantId, ); + await subagentThreadTaskStore.cancelPlan(cancellationPlan); + dbResponse = await db.deleteConvos(req.user.id, filter); + } else { + /** An empty filter deletes every conversation this owner has, so it runs behind + * the same admission fence as `DELETE /all` rather than a bare drain. */ + dbResponse = await subagentThreadTaskStore.withOwnerDeletionFence(req.user.id, tenantId, () => + db.deleteConvos(req.user.id, filter), + ); } - const dbResponse = await db.deleteConvos(req.user.id, filter); const deletedConversationIds = dbResponse.conversationIds ?? (filter.conversationId ? [filter.conversationId] : []); - subagentThreadTaskStore.cancelForConversations(req.user.id, deletedConversationIds, tenantId); + /** Root deletion closes new child admission. Replay the plan to catch a task + * admitted after the first pass but before that fence, extended with the cascade + * this deletion reported. */ + if (cancellationPlan != null && deletedConversationIds.length > 0) { + /** The conversations are gone, so this pass is the only thing that can still + * stop a child admitted after the first one. It cannot fail the request — the + * deletion already committed — so it retries briefly before giving up. */ + await retryPostDeleteCancellation(cancellationPlan, deletedConversationIds); + } // HITL: prune the deleted conversations' durable checkpoints — a paused run's // checkpoint would otherwise persist until the Mongo TTL. Never throws. await deleteAgentCheckpoints( @@ -186,13 +225,18 @@ router.delete('/', configMiddleware, async (req, res) => { router.delete('/all', configMiddleware, async (req, res) => { try { - subagentThreadTaskStore.cancelForOwner( - req.user.id, + const tenantId = typeof req.user.tenantId === 'string' && req.user.tenantId !== '' ? req.user.tenantId - : undefined, + : undefined; + /** Fences new child admission for this owner, drains the live ones, and deletes + * inside that fence: a child admitted on another replica mid-deletion would + * otherwise keep running against conversations that no longer exist. */ + const dbResponse = await subagentThreadTaskStore.withOwnerDeletionFence( + req.user.id, + tenantId, + () => db.deleteConvos(req.user.id, {}), ); - const dbResponse = await db.deleteConvos(req.user.id, {}); // HITL: prune ALL the deleted conversations' durable checkpoints in one bulk pass. await deleteAgentCheckpoints( dbResponse.conversationIds, diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 46be766f7c..a3256f54bd 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 a5be3f9852..59910aa2a0 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 ff1dd6c60b..f11c6a2267 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -1,16 +1,25 @@ -const { createSubagentThreadTaskStore } = require('@librechat/api'); +const { + cacheConfig, + ioredisClient, + registerShutdownTask, + duplicateIoRedisClient, + createSubagentThreadTaskStore, + RedisSubagentTaskControlTransport, +} = require('@librechat/api'); const db = require('~/models'); -/** Durable logical threads use normal LibreChat conversations/messages. Live - * controls stay process-local; Mongo fences continuation across API replicas. */ +/** Durable logical threads use normal LibreChat conversations/messages. Mongo + * fences continuation; optional Redis routing reaches the live owning process. */ const subagentThreadTaskStore = createSubagentThreadTaskStore( { acquireSubagentThreadLease: db.acquireSubagentThreadLease, + claimSubagentTaskResult: db.claimSubagentTaskResult, countActiveSubagentThreadLeases: db.countActiveSubagentThreadLeases, deleteConvos: db.deleteConvos, deleteMessages: db.deleteMessages, getConvo: db.getConvo, getMessages: db.getMessages, + listActiveSubagentThreadLeases: db.listActiveSubagentThreadLeases, releaseSubagentThreadLease: db.releaseSubagentThreadLease, reserveSubagentThread: db.reserveSubagentThread, renewSubagentThreadLease: db.renewSubagentThreadLease, @@ -18,8 +27,49 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( saveMessage: db.saveMessage, }, { - isOwnerActive: db.isAgentTriggerPrincipalActive, + isOwnerActive: db.isSubagentOwnerAdmissible, + fenceOwnerAdmission: db.fenceSubagentAdmission, + renewOwnerAdmission: db.renewSubagentAdmission, + releaseOwnerAdmission: db.releaseSubagentAdmission, }, ); +let taskRoutingConfigured = false; + +/** Starts the optional Redis owner directory before HTTP admission opens. */ +async function configureSubagentTaskRouting() { + if (taskRoutingConfigured || !cacheConfig.USE_REDIS) { + return; + } + if (ioredisClient == null || typeof ioredisClient.duplicate !== 'function') { + throw new Error('Redis subagent task routing requires a dedicated subscriber connection.'); + } + const subscriber = ioredisClient.duplicate(); + /** A dedicated publisher without the offline queue: the shared client would hold a + * command issued during a disconnect and deliver it after the caller gave up, so a + * steer the caller was told had failed could still reach the child. Failing fast + * turns that into the honest `unavailable` the caller already handles. */ + const publisher = duplicateIoRedisClient(ioredisClient, { enableOfflineQueue: false }); + const transport = new RedisSubagentTaskControlTransport(publisher, subscriber, { + namespace: cacheConfig.REDIS_KEY_PREFIX, + }); + try { + await subagentThreadTaskStore.configureTaskControlTransport(transport); + } catch (error) { + subscriber.disconnect(); + publisher.disconnect(); + throw error; + } + taskRoutingConfigured = true; + registerShutdownTask( + 'subagent task control transport', + async () => { + await subagentThreadTaskStore.destroyTaskControlTransport(); + publisher.disconnect(); + }, + { priority: 90 }, + ); +} + module.exports = subagentThreadTaskStore; +module.exports.configureSubagentTaskRouting = configureSubagentTaskRouting; diff --git a/packages/api/src/agents/background.spec.ts b/packages/api/src/agents/background.spec.ts index 4c604c8b98..97fbc2eed6 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 0f06c2ad6d..113bd2c7fb 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 0673ac1071..ee2a69c9b2 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 c99e4aea24..4e0442b4ef 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 95d1b34b53..f79b560f11 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -32,6 +32,7 @@ export * from './skills'; export * from './phases'; export * from './startup'; export * from './subagentThreads'; +export * from './subagentTaskRouting'; export * from './skillConfigurable'; export * from './skillFiles'; export * from './codeFilesSession'; diff --git a/packages/api/src/agents/subagentTaskRouting.spec.ts b/packages/api/src/agents/subagentTaskRouting.spec.ts new file mode 100644 index 0000000000..044dec0e23 --- /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 0000000000..7496bbdafd --- /dev/null +++ b/packages/api/src/agents/subagentTaskRouting.ts @@ -0,0 +1,1379 @@ +import { logger } from '@librechat/data-schemas'; +import { createHash, randomUUID } from 'node:crypto'; +import type { + SubagentTaskClaim, + SubagentTaskControlCommand, + SubagentTaskControlResult, + SubagentTaskSnapshot, +} from '@librechat/agents'; +import type { Cluster, Redis } from 'ioredis'; +import { createConcurrencyLimiter } from '~/utils/promise'; + +const PROTOCOL_VERSION = 1; +const DEFAULT_REQUEST_TIMEOUT_MS = 2_000; +const DEFAULT_RETRY_DELAY_MS = 500; +const DEFAULT_READY_TIMEOUT_MS = 10_000; +const DEFAULT_REGISTRATION_HEARTBEAT_MS = 10_000; +const MAX_PENDING_REQUESTS = 1_000; +/** A consumed claim is retained apart from control replays so unrelated command + * traffic cannot displace it while its caller retries. Retention is a fast path, not + * the guarantee: the terminal result is recoverable from its durable child message. */ +const MAX_CLAIM_REPLAY_ENTRIES = 2_000; +const MAX_CLAIM_REPLAY_BYTES = 16 * 1024 * 1024; +const MAX_CONTROL_REPLAY_ENTRIES = 2_000; +const MAX_CONTROL_REPLAY_BYTES = 4 * 1024 * 1024; +const RESPONSE_CACHE_TTL_MS = 5 * 60_000; +/** Absorbs ordinary clock drift between replicas when honouring a request deadline. */ +const REQUEST_CLOCK_SKEW_MS = 30_000; +const MAX_SCOPE_ID_CHARS = 4_096; +const MAX_TASK_ID_CHARS = 256; +const MAX_CONTROL_MESSAGE_CHARS = 64 * 1_024; +const MAX_RESULT_CHARS = 100_000; +const MAX_ERROR_CHARS = 4 * 1_024; +const MAX_THREAD_ID_CHARS = 256; +const MAX_SUBAGENT_TYPE_CHARS = 256; +const MAX_PROGRESS_LABEL_CHARS = 1_024; +/** Bounds the model-facing task list, per owner reply and across the merged result. */ +export const MAX_TASK_SNAPSHOTS = 200; +const MAX_CANCEL_THREAD_IDS = 200; +/** Matches the deletion drain so bounded fan-out stays well inside the lease TTL. */ +const ROUTING_FANOUT_CONCURRENCY = 32; +/** Contains every bounded response even when JSON escapes each retained character. */ +const MAX_ROUTED_MESSAGE_CHARS = 8 * 1_024 * 1_024; + +const REGISTER_TASK_SCRIPT = + "local now = redis.call('TIME'); " + + 'local ttl = tonumber(ARGV[3]); ' + + 'local expiresAt = (tonumber(now[1]) * 1000) + math.floor(tonumber(now[2]) / 1000) + ttl; ' + + "redis.call('HSET', KEYS[1], ARGV[1], tostring(expiresAt) .. '|' .. ARGV[2]); " + + "local directoryTtl = redis.call('PTTL', KEYS[1]); " + + "if directoryTtl < ttl then redis.call('PEXPIRE', KEYS[1], ttl); end; " + + 'return 1'; + +const READ_ACTIVE_REGISTRATIONS_SCRIPT = + "local now = redis.call('TIME'); " + + 'local nowMs = (tonumber(now[1]) * 1000) + math.floor(tonumber(now[2]) / 1000); ' + + "local entries = redis.call('HGETALL', KEYS[1]); " + + 'local active = {}; ' + + 'for i = 1, #entries, 2 do ' + + 'local value = entries[i + 1]; ' + + "local separator = string.find(value, '|', 1, true); " + + 'local expiresAt = separator and tonumber(string.sub(value, 1, separator - 1)); ' + + 'if expiresAt and expiresAt > nowMs then ' + + 'table.insert(active, entries[i]); ' + + 'table.insert(active, string.sub(value, separator + 1)); ' + + "else redis.call('HDEL', KEYS[1], entries[i]); end; " + + 'end; ' + + 'return active'; + +const READ_TASK_OWNER_SCRIPT = + "local value = redis.call('HGET', KEYS[1], ARGV[1]); " + + 'if not value then return nil; end; ' + + "local separator = string.find(value, '|', 1, true); " + + 'local expiresAt = separator and tonumber(string.sub(value, 1, separator - 1)); ' + + "local now = redis.call('TIME'); " + + 'local nowMs = (tonumber(now[1]) * 1000) + math.floor(tonumber(now[2]) / 1000); ' + + "if not expiresAt or expiresAt <= nowMs then redis.call('HDEL', KEYS[1], ARGV[1]); return nil; end; " + + 'return string.sub(value, separator + 1)'; + +type RedisClient = Redis | Cluster; +interface RoutedRequestBase { + version: typeof PROTOCOL_VERSION; + kind: 'request'; + requestId: string; + requesterId: string; + scopeId: string; + /** Epoch milliseconds after which the requester has stopped waiting. */ + expiresAt: number; +} + +type RoutedRequest = RoutedRequestBase & + ( + | { operation: 'claim'; taskId: string } + | { + operation: 'control'; + taskId: string; + command: SubagentTaskControlCommand; + invocationId: string; + } + | { operation: 'list' } + | { operation: 'cancel'; threadIds: string[] | null } + ); + +type RoutedRequestPayload = + | { operation: 'claim'; scopeId: string; taskId: string } + | { + operation: 'control'; + scopeId: string; + taskId: string; + command: SubagentTaskControlCommand; + invocationId: string; + } + | { operation: 'list'; scopeId: string } + | { operation: 'cancel'; scopeId: string; threadIds: string[] | null }; + +interface RoutedResponse { + version: typeof PROTOCOL_VERSION; + kind: 'response'; + requestId: string; + ok: boolean; + result?: unknown; +} + +/** Tells the owner a consumed result reached a caller and no longer needs retaining. */ +interface RoutedAck { + version: typeof PROTOCOL_VERSION; + kind: 'ack'; + scopeId: string; + taskId: string; +} + +interface PendingRequest { + resolve: (value: unknown) => void; + reject: (error: Error) => void; + retry: ReturnType; + timeout: ReturnType; +} + +interface CachedResponse { + value: string; + bytes: number; + expiresAt: number; + /** Content this response answered, so one id cannot replay a different command. */ + fingerprint?: string; +} + +interface ReplayCache { + entries: Map; + bytes: number; + maxEntries: number; + maxBytes: number; +} + +interface RoutedTaskList { + snapshots: SubagentTaskSnapshot[]; + truncated: boolean; +} + +interface RoutedCancelResult { + cancelled: number; +} + +interface OwnedTaskRegistration { + scopeId: string; + taskId: string; + ttlMs: number; +} + +export interface SubagentTaskControlHandler { + claim(scopeId: string, taskId: string): SubagentTaskClaim; + control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): SubagentTaskControlResult; + list(scopeId: string): SubagentTaskSnapshot[]; + cancelScope(scopeId: string, threadIds: string[] | null): number; +} + +/** Optional host transport for reaching the process that owns a live child task. */ +export interface SubagentTaskControlTransport { + bind(handler: SubagentTaskControlHandler): Promise; + registerTask(scopeId: string, taskId: string, ttlMs: number): Promise; + hasTasks(scopeId: string): Promise; + claim(scopeId: string, taskId: string): Promise; + control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise; + list(scopeId: string): Promise; + cancelScope(scopeId: string, threadIds: string[] | null): Promise; + destroy(): Promise; +} + +export class SubagentTaskOwnerUnavailableError extends Error { + constructor() { + super('The process running this subagent task is temporarily unavailable.'); + } +} + +export interface RedisSubagentTaskControlTransportOptions { + /** Separates pub/sub channels for deployments sharing one Redis service. */ + namespace?: string; + instanceId?: string; + requestTimeoutMs?: number; + retryDelayMs?: number; + registrationHeartbeatMs?: number; +} + +function positiveInteger(value: number | undefined, fallback: number): number { + return Number.isSafeInteger(value) && value != null && value > 0 ? value : fallback; +} + +function shortHash(value: string): string { + return createHash('sha256').update(value).digest('base64url').slice(0, 24); +} + +function isBoundedString(value: unknown, maxChars: number): value is string { + return typeof value === 'string' && value.length > 0 && value.length <= maxChars; +} + +function isStringWithin(value: unknown, maxChars: number): value is string { + return typeof value === 'string' && value.length <= maxChars; +} + +function truncateMiddle(value: string, maxChars: number): string { + if (value.length <= maxChars) { + return value; + } + const marker = '\n…[truncated]…\n'; + const available = Math.max(0, maxChars - marker.length); + const head = Math.ceil(available / 2); + return `${value.slice(0, head)}${marker}${value.slice(value.length - (available - head))}`; +} + +function boundedSnapshot(snapshot: SubagentTaskSnapshot): SubagentTaskSnapshot { + return { + taskId: truncateMiddle(snapshot.taskId, MAX_TASK_ID_CHARS), + ...(snapshot.threadId == null + ? {} + : { threadId: truncateMiddle(snapshot.threadId, MAX_THREAD_ID_CHARS) }), + subagentType: truncateMiddle(snapshot.subagentType, MAX_SUBAGENT_TYPE_CHARS), + status: snapshot.status, + createdAt: snapshot.createdAt, + updatedAt: snapshot.updatedAt, + resultAvailable: snapshot.resultAvailable, + resultClaimed: snapshot.resultClaimed, + pendingControls: snapshot.pendingControls, + ...(snapshot.progress == null + ? {} + : { + progress: { + ...snapshot.progress, + ...(snapshot.progress.label == null + ? {} + : { label: truncateMiddle(snapshot.progress.label, MAX_PROGRESS_LABEL_CHARS) }), + }, + }), + ...(snapshot.error == null ? {} : { error: truncateMiddle(snapshot.error, MAX_ERROR_CHARS) }), + }; +} + +/** + * Bounds a model-facing task list, keeping what a caller can still act on: running + * children first, then the most recent settled results. A plain oldest-first slice + * would drop the newest tasks, hiding a child that just started from the only tool + * able to poll it. + */ +export function boundedTaskList(tasks: SubagentTaskSnapshot[]): SubagentTaskSnapshot[] { + const byCreatedAt = (left: SubagentTaskSnapshot, right: SubagentTaskSnapshot): number => + left.createdAt - right.createdAt; + if (tasks.length <= MAX_TASK_SNAPSHOTS) { + return tasks.sort(byCreatedAt); + } + const running: SubagentTaskSnapshot[] = []; + const settled: SubagentTaskSnapshot[] = []; + for (const task of tasks) { + (task.status === 'running' ? running : settled).push(task); + } + running.sort(byCreatedAt); + const keptRunning = running.slice(-MAX_TASK_SNAPSHOTS); + const remaining = MAX_TASK_SNAPSHOTS - keptRunning.length; + if (remaining <= 0) { + return keptRunning; + } + settled.sort(byCreatedAt); + return [...keptRunning, ...settled.slice(-remaining)].sort(byCreatedAt); +} + +/** Applies the routed result and snapshot bounds to a claim from any source. */ +export function boundedClaim(claim: SubagentTaskClaim): SubagentTaskClaim { + if (claim.status === 'not_found') { + return claim; + } + const task = boundedSnapshot(claim.task); + if (claim.status === 'completed') { + return { status: 'completed', task, result: truncateMiddle(claim.result, MAX_RESULT_CHARS) }; + } + if (claim.status === 'error' || claim.status === 'cancelled') { + return { status: claim.status, task, error: truncateMiddle(claim.error, MAX_ERROR_CHARS) }; + } + return { status: claim.status, task }; +} + +function boundedControlResult(result: SubagentTaskControlResult): SubagentTaskControlResult { + if (result.status === 'not_found') { + return result; + } + if (result.status === 'invalid') { + return { status: 'invalid', message: truncateMiddle(result.message, MAX_ERROR_CHARS) }; + } + return { + status: result.status, + task: boundedSnapshot(result.task), + ...(result.status === 'accepted' && result.controlId != null + ? { controlId: truncateMiddle(result.controlId, MAX_TASK_ID_CHARS) } + : {}), + }; +} + +async function waitForRedisConnectionReady(client: RedisClient): Promise { + if (client.status == null || client.status === 'ready') { + return; + } + if (client.status === 'end') { + throw new SubagentTaskOwnerUnavailableError(); + } + await new Promise((resolve, reject) => { + const cleanup = () => { + clearTimeout(timeout); + client.off('ready', onReady); + client.off('end', onEnd); + }; + const onReady = () => { + cleanup(); + resolve(); + }; + const onEnd = () => { + cleanup(); + reject(new SubagentTaskOwnerUnavailableError()); + }; + const timeout = setTimeout(() => { + cleanup(); + reject(new SubagentTaskOwnerUnavailableError()); + }, DEFAULT_READY_TIMEOUT_MS); + timeout.unref?.(); + client.once('ready', onReady); + client.once('end', onEnd); + /** Close the status-check/listener-registration race: ioredis may become ready + * synchronously between the check above and installing these listeners. */ + if (client.status === 'ready') { + onReady(); + } else if (client.status === 'end') { + onEnd(); + } else if (client.status === 'wait') { + client.connect().catch(onEnd); + } + }); +} + +async function waitForRedisReady( + client: RedisClient, + options: { eagerClusterMasters?: boolean } = {}, +): Promise { + await waitForRedisConnectionReady(client); + if (options.eagerClusterMasters !== true || !client.isCluster) { + return; + } + /** A ready Cluster has a slot map but its per-master connections are lazy. The + * fail-fast publisher cannot admit requests until every possible write target is + * connected; otherwise the first command to a cold shard would be rejected. */ + await Promise.all( + (client as Cluster).nodes('master').map((node) => waitForRedisConnectionReady(node)), + ); +} + +function isSnapshot(value: unknown): value is SubagentTaskSnapshot { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + return ( + isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS) && + typeof candidate.subagentType === 'string' && + ['running', 'completed', 'error', 'cancelled'].includes(candidate.status ?? '') && + typeof candidate.createdAt === 'number' && + typeof candidate.updatedAt === 'number' && + typeof candidate.resultAvailable === 'boolean' && + typeof candidate.resultClaimed === 'boolean' && + typeof candidate.pendingControls === 'number' + ); +} + +function isClaim(value: unknown): value is SubagentTaskClaim { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + if (candidate.status === 'not_found') { + return true; + } + if (!('task' in candidate) || !isSnapshot(candidate.task)) { + return false; + } + if (candidate.status === 'completed') { + return 'result' in candidate && typeof candidate.result === 'string'; + } + if (candidate.status === 'error' || candidate.status === 'cancelled') { + return 'error' in candidate && typeof candidate.error === 'string'; + } + return candidate.status === 'running' || candidate.status === 'claimed'; +} + +function isControlResult(value: unknown): value is SubagentTaskControlResult { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + if (candidate.status === 'not_found') { + return true; + } + if (candidate.status === 'invalid') { + return typeof candidate.message === 'string'; + } + return ( + ['accepted', 'cancelled', 'not_running', 'control_not_found'].includes( + candidate.status ?? '', + ) && + 'task' in candidate && + isSnapshot(candidate.task) + ); +} + +function isCancelThreadIds(value: unknown): value is string[] { + return ( + Array.isArray(value) && + value.length > 0 && + value.length <= MAX_CANCEL_THREAD_IDS && + value.every((threadId) => isBoundedString(threadId, MAX_THREAD_ID_CHARS)) + ); +} + +function isCancelResult(value: unknown): value is RoutedCancelResult { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const { cancelled } = value as Partial; + return Number.isSafeInteger(cancelled) && (cancelled as number) >= 0; +} + +function controlContent(command: SubagentTaskControlCommand): string { + if (command.action === 'cancel') { + return 'cancel'; + } + if (command.action === 'cancel_message') { + return `cancel_message\u0000${command.controlId}`; + } + return `${command.action}\u0000${command.message}`; +} + +/** + * Canonical identity of one control's content. Property order cannot vary it, so the + * transport and the owning task store agree on when two commands are the same, and it + * is hashed so retaining one costs a fixed few bytes rather than a whole message. + */ +export function controlFingerprint(command: SubagentTaskControlCommand): string { + return createHash('sha256').update(controlContent(command)).digest('base64url'); +} + +/** True once a claim has consumed the task's one-shot terminal result. */ +function consumesResult(result: SubagentTaskClaim): boolean { + return ( + result.status === 'completed' || result.status === 'error' || result.status === 'cancelled' + ); +} + +function isRoutedTaskList(value: unknown): value is RoutedTaskList { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const candidate = value as Partial; + return ( + Array.isArray(candidate.snapshots) && + candidate.snapshots.every(isSnapshot) && + typeof candidate.truncated === 'boolean' + ); +} + +function parseControlCommand(value: unknown): SubagentTaskControlCommand | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as { action?: unknown; message?: unknown; controlId?: unknown }; + if (candidate.action === 'cancel') { + return { action: 'cancel' }; + } + if (candidate.action === 'cancel_message') { + return isStringWithin(candidate.controlId, MAX_TASK_ID_CHARS) + ? { action: 'cancel_message', controlId: candidate.controlId } + : undefined; + } + if ( + (candidate.action === 'steer' || + candidate.action === 'queue' || + candidate.action === 'interrupt') && + isStringWithin(candidate.message, MAX_CONTROL_MESSAGE_CHARS) + ) { + return { action: candidate.action, message: candidate.message }; + } + return undefined; +} + +function failureResponse(requestId: string): string { + const response: RoutedResponse = { + version: PROTOCOL_VERSION, + kind: 'response', + requestId, + ok: false, + }; + return JSON.stringify(response); +} + +function successResponse(requestId: string, result: string): string { + return `{"version":${PROTOCOL_VERSION},"kind":"response","requestId":${JSON.stringify( + requestId, + )},"ok":true,"result":${result}}`; +} + +function parseRequest(value: unknown): RoutedRequest | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as { + version?: unknown; + kind?: unknown; + requestId?: unknown; + requesterId?: unknown; + operation?: unknown; + scopeId?: unknown; + taskId?: unknown; + command?: unknown; + threadIds?: unknown; + invocationId?: unknown; + expiresAt?: unknown; + }; + if ( + candidate.version !== PROTOCOL_VERSION || + candidate.kind !== 'request' || + !isBoundedString(candidate.requestId, 128) || + !isBoundedString(candidate.requesterId, 128) || + !['claim', 'control', 'list', 'cancel'].includes( + typeof candidate.operation === 'string' ? candidate.operation : '', + ) || + !isBoundedString(candidate.scopeId, MAX_SCOPE_ID_CHARS) || + !Number.isSafeInteger(candidate.expiresAt) + ) { + return undefined; + } + const expiresAt = candidate.expiresAt as number; + if (candidate.operation === 'list') { + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'list', + scopeId: candidate.scopeId, + }; + } + if (candidate.operation === 'cancel') { + if (candidate.threadIds !== null && !isCancelThreadIds(candidate.threadIds)) { + return undefined; + } + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'cancel', + scopeId: candidate.scopeId, + threadIds: candidate.threadIds, + }; + } + if (!isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS)) { + return undefined; + } + if (candidate.operation === 'claim') { + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'claim', + scopeId: candidate.scopeId, + taskId: candidate.taskId, + }; + } + const command = parseControlCommand(candidate.command); + if (command == null || !isBoundedString(candidate.invocationId, 128)) { + return undefined; + } + return { + version: PROTOCOL_VERSION, + kind: 'request', + requestId: candidate.requestId, + requesterId: candidate.requesterId, + expiresAt, + operation: 'control', + scopeId: candidate.scopeId, + taskId: candidate.taskId, + command, + invocationId: candidate.invocationId, + }; +} + +function createReplayCache(maxEntries: number, maxBytes: number): ReplayCache { + return { entries: new Map(), bytes: 0, maxEntries, maxBytes }; +} + +function parseAck(value: unknown): RoutedAck | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as Partial; + if ( + candidate.version !== PROTOCOL_VERSION || + candidate.kind !== 'ack' || + !isBoundedString(candidate.scopeId, MAX_SCOPE_ID_CHARS) || + !isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS) + ) { + return undefined; + } + return { + version: PROTOCOL_VERSION, + kind: 'ack', + scopeId: candidate.scopeId, + taskId: candidate.taskId, + }; +} + +function parseResponse(value: unknown): RoutedResponse | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const candidate = value as Partial; + if ( + candidate.version !== PROTOCOL_VERSION || + candidate.kind !== 'response' || + !isBoundedString(candidate.requestId, 128) || + typeof candidate.ok !== 'boolean' + ) { + return undefined; + } + return candidate as RoutedResponse; +} + +/** + * Routes bounded live-task operations to their owning API replica. Redis keeps + * only an expiring owner directory and request/reply envelopes; the executor, + * transcript, and checkpoint never move between processes. + */ +export class RedisSubagentTaskControlTransport implements SubagentTaskControlTransport { + private readonly instanceId: string; + private readonly namespaceHash: string; + private readonly requestTimeoutMs: number; + private readonly retryDelayMs: number; + private readonly registrationHeartbeatMs: number; + private readonly pending = new Map(); + private readonly claimReplays = createReplayCache( + MAX_CLAIM_REPLAY_ENTRIES, + MAX_CLAIM_REPLAY_BYTES, + ); + + private readonly controlReplays = createReplayCache( + MAX_CONTROL_REPLAY_ENTRIES, + MAX_CONTROL_REPLAY_BYTES, + ); + + private readonly ownedTasks = new Map(); + private handler?: SubagentTaskControlHandler; + private ready?: Promise; + private registrationHeartbeat?: ReturnType; + private registrationRefresh?: Promise; + private destroyed = false; + + constructor( + private readonly publisher: RedisClient, + private readonly subscriber: RedisClient, + options: RedisSubagentTaskControlTransportOptions = {}, + ) { + this.instanceId = options.instanceId?.trim() || randomUUID(); + this.namespaceHash = shortHash(options.namespace?.trim() || 'default'); + this.requestTimeoutMs = positiveInteger(options.requestTimeoutMs, DEFAULT_REQUEST_TIMEOUT_MS); + this.retryDelayMs = Math.min( + positiveInteger(options.retryDelayMs, DEFAULT_RETRY_DELAY_MS), + Math.max(1, Math.floor(this.requestTimeoutMs / 2)), + ); + this.registrationHeartbeatMs = positiveInteger( + options.registrationHeartbeatMs, + DEFAULT_REGISTRATION_HEARTBEAT_MS, + ); + } + + async bind(handler: SubagentTaskControlHandler): Promise { + if (this.destroyed) { + throw new Error('Subagent task control transport is closed.'); + } + if (this.handler != null) { + throw new Error('Subagent task control transport is already bound.'); + } + this.handler = handler; + /** The publisher fails fast instead of queueing commands, so opening HTTP + * admission before it is ready would turn healthy startup lag into false + * `unavailable` results. Both dedicated connections are part of readiness. */ + await Promise.all([ + waitForRedisReady(this.publisher, { eagerClusterMasters: true }), + waitForRedisReady(this.subscriber), + ]); + this.subscriber.on('message', this.onMessage); + this.ready = this.subscriber.subscribe(this.channel(this.instanceId)).then(() => undefined); + await this.ready; + } + + async registerTask(scopeId: string, taskId: string, ttlMs: number): Promise { + this.assertTaskAddress(scopeId, taskId); + const registration = { + scopeId, + taskId, + ttlMs: positiveInteger(ttlMs, 1), + }; + this.ownedTasks.set(this.registrationKey(scopeId, taskId), registration); + this.ensureRegistrationHeartbeat(); + await this.publishRegistration(registration); + } + + async hasTasks(scopeId: string): Promise { + this.assertScope(scopeId); + await this.requireReady(); + try { + return Object.keys(await this.readActiveRegistrations(scopeId)).length > 0; + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to inspect the task owner directory', error); + throw new SubagentTaskOwnerUnavailableError(); + } + } + + async claim(scopeId: string, taskId: string): Promise { + const routed = await this.requestTaskOwner(scopeId, taskId, 'claim'); + if (routed == null) { + return undefined; + } + const { ownerId, result } = routed; + if (!isClaim(result)) { + throw new SubagentTaskOwnerUnavailableError(); + } + if (result.status === 'not_found') { + await this.removeRegistrations(scopeId, [taskId]); + return undefined; + } + if (consumesResult(result)) { + /** Frees the owner's retained copy immediately. An acknowledgement that cannot + * be confirmed only leaves that copy to expire, so the caller still keeps the + * result it is holding rather than trading it for a retry. */ + await this.acknowledgeClaim(ownerId, scopeId, taskId); + } + return result; + } + + async control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise { + const routed = await this.requestTaskOwner(scopeId, taskId, 'control', command, invocationId); + if (routed == null) { + return undefined; + } + const { result } = routed; + if (!isControlResult(result)) { + throw new SubagentTaskOwnerUnavailableError(); + } + if (result.status === 'not_found') { + await this.removeRegistrations(scopeId, [taskId]); + return undefined; + } + return result; + } + + async list(scopeId: string): Promise { + this.assertScope(scopeId); + await this.requireReady(); + let ownersByTask: Record; + try { + ownersByTask = await this.readActiveRegistrations(scopeId); + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to read the task owner directory', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const owners = new Set(Object.values(ownersByTask)); + owners.delete(this.instanceId); + if (owners.size === 0) { + return []; + } + const ownerIds = [...owners]; + const results = await Promise.all( + ownerIds.map((ownerId) => this.sendRequest(ownerId, { operation: 'list', scopeId })), + ); + const snapshots: SubagentTaskSnapshot[] = []; + const staleTaskIds: string[] = []; + for (const [index, value] of results.entries()) { + if (!isRoutedTaskList(value)) { + throw new SubagentTaskOwnerUnavailableError(); + } + const ownerId = ownerIds[index]; + const reportedTaskIds = new Set(value.snapshots.map((snapshot) => snapshot.taskId)); + for (const snapshot of value.snapshots) { + if (ownersByTask[snapshot.taskId] === ownerId) { + snapshots.push(snapshot); + } + } + if (!value.truncated) { + for (const [taskId, registeredOwnerId] of Object.entries(ownersByTask)) { + if (registeredOwnerId === ownerId && !reportedTaskIds.has(taskId)) { + staleTaskIds.push(taskId); + } + } + } + } + if (staleTaskIds.length > 0) { + await this.removeRegistrations(scopeId, staleTaskIds); + } + /** Each owner bounds its own reply, so without an aggregate cap this grows with the + * number of replicas holding the scope. Bounding after the loop rather than during + * it keeps the sweep above reading every owner's reply, and lets the cap choose by + * status instead of by whichever owner answered first. */ + return boundedTaskList(snapshots); + } + + /** + * Cancels live children on every other owner of this scope. The owner applies the + * predicate to its complete local task set, so deletion never depends on the + * bounded model-facing list and cannot miss a task beyond that cap. + */ + async cancelScope(scopeId: string, threadIds: string[] | null): Promise { + this.assertScope(scopeId); + if (threadIds != null && threadIds.length === 0) { + return 0; + } + await this.requireReady(); + let ownersByTask: Record; + try { + ownersByTask = await this.readActiveRegistrations(scopeId); + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to read the task owner directory', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const owners = new Set(Object.values(ownersByTask)); + owners.delete(this.instanceId); + if (owners.size === 0) { + return 0; + } + const batches: Array = []; + if (threadIds == null) { + batches.push(null); + } else { + for (let index = 0; index < threadIds.length; index += MAX_CANCEL_THREAD_IDS) { + batches.push(threadIds.slice(index, index + MAX_CANCEL_THREAD_IDS)); + } + } + const cancelSlot = createConcurrencyLimiter(ROUTING_FANOUT_CONCURRENCY); + const requests: Array> = []; + for (const ownerId of owners) { + for (const batch of batches) { + requests.push( + cancelSlot(() => + this.sendRequest(ownerId, { operation: 'cancel', scopeId, threadIds: batch }), + ), + ); + } + } + let cancelled = 0; + for (const value of await Promise.all(requests)) { + if (!isCancelResult(value)) { + throw new SubagentTaskOwnerUnavailableError(); + } + cancelled += value.cancelled; + } + return cancelled; + } + + async destroy(): Promise { + if (this.destroyed) { + return; + } + this.destroyed = true; + for (const pending of this.pending.values()) { + clearTimeout(pending.retry); + clearTimeout(pending.timeout); + pending.reject(new SubagentTaskOwnerUnavailableError()); + } + this.pending.clear(); + for (const cache of [this.claimReplays, this.controlReplays]) { + cache.entries.clear(); + cache.bytes = 0; + } + this.ownedTasks.clear(); + if (this.registrationHeartbeat != null) { + clearInterval(this.registrationHeartbeat); + this.registrationHeartbeat = undefined; + } + this.subscriber.off('message', this.onMessage); + await this.subscriber.unsubscribe(this.channel(this.instanceId)).catch(() => undefined); + this.subscriber.disconnect(); + } + + private readonly onMessage = (channel: string, message: string): void => { + if (channel !== this.channel(this.instanceId) || message.length > MAX_ROUTED_MESSAGE_CHARS) { + return; + } + let parsed: unknown; + try { + parsed = JSON.parse(message) as unknown; + } catch { + return; + } + const response = parseResponse(parsed); + if (response != null) { + this.handleResponse(response); + return; + } + const ack = parseAck(parsed); + if (ack != null) { + this.releaseClaimReplay(ack.scopeId, ack.taskId); + return; + } + const request = parseRequest(parsed); + if (request != null) { + void this.handleRequest(request).catch((error) => { + logger.warn('[subagentTaskRouting] Failed to answer a routed command', error); + }); + } + }; + + private handleResponse(response: RoutedResponse): void { + const pending = this.pending.get(response.requestId); + if (pending == null) { + return; + } + this.pending.delete(response.requestId); + clearTimeout(pending.retry); + clearTimeout(pending.timeout); + if (!response.ok) { + pending.reject(new SubagentTaskOwnerUnavailableError()); + return; + } + pending.resolve(response.result); + } + + private async handleRequest(request: RoutedRequest): Promise { + if (Date.now() > request.expiresAt + REQUEST_CLOCK_SKEW_MS) { + /** The caller stopped waiting for this long ago and has been told it was + * unavailable, so applying it now would steer a child it believes untouched. */ + logger.warn('[subagentTaskRouting] Dropped a routed command past its deadline'); + return; + } + const replay = this.replayFor(request); + const cached = replay?.cache.entries.get(replay.key); + /** A retransmission replays; the same id carrying different content is a caller + * error, so it reaches the owner, which refuses it, rather than being answered + * from the earlier command's response. */ + if ( + cached != null && + cached.expiresAt > Date.now() && + cached.fingerprint === replay?.fingerprint + ) { + await this.publish( + this.channel(request.requesterId), + successResponse(request.requestId, cached.value), + ); + return; + } + const handler = this.handler; + if (handler == null) { + return; + } + let serialized: string; + try { + let result: + | SubagentTaskClaim + | SubagentTaskControlResult + | RoutedTaskList + | RoutedCancelResult; + /** A claim that consumed nothing stays uncached so a later poll still observes + * the task's live status. */ + let replayable = replay != null; + if (request.operation === 'list') { + const tasks = handler.list(request.scopeId); + /** Bounded the same way the requester bounds the merge: a positional slice here + * would drop this owner's running children before they ever reached it. */ + const bounded = boundedTaskList(tasks); + result = { + snapshots: bounded.map(boundedSnapshot), + truncated: tasks.length > bounded.length, + }; + } else if (request.operation === 'cancel') { + result = { cancelled: handler.cancelScope(request.scopeId, request.threadIds) }; + } else if (request.operation === 'claim') { + const claim = boundedClaim(handler.claim(request.scopeId, request.taskId)); + replayable = consumesResult(claim); + result = claim; + } else { + result = boundedControlResult( + handler.control(request.scopeId, request.taskId, request.command, request.invocationId), + ); + } + const serializedResult = JSON.stringify(result); + /** Retaining the result rather than the envelope lets a later caller retry, + * which carries its own correlation id, recover a response it never received. */ + if (replay != null && replayable) { + this.retainReplay(replay.cache, replay.key, serializedResult, replay.fingerprint); + } + serialized = successResponse(request.requestId, serializedResult); + } catch (error) { + logger.error('[subagentTaskRouting] Owner failed to process a routed command', error); + serialized = failureResponse(request.requestId); + } + await this.publish(this.channel(request.requesterId), serialized); + } + + /** + * Locates a destructive operation's replay slot. A claim consumes the one-shot + * result, so it is keyed by operation—stable across callers and replicas, so a + * later poll still resolves to the response the owner produced. Lists are + * idempotent and recomputable, so their large bodies are never retained. + */ + private replayFor( + request: RoutedRequest, + ): { cache: ReplayCache; key: string; fingerprint?: string } | undefined { + if (request.operation === 'list') { + return undefined; + } + if (request.operation === 'claim') { + return { + cache: this.claimReplays, + key: this.claimReplayKey(request.scopeId, request.taskId), + }; + } + if (request.operation === 'cancel') { + return { cache: this.controlReplays, key: `cancel\u0000${request.requestId}` }; + } + /** Task-scoped: a provider tool-call id such as `call_0` repeats across runs and + * agents, so keying on it alone would answer one task from another's snapshot. */ + return { + cache: this.controlReplays, + key: `control\u0000${shortHash(request.scopeId)}\u0000${request.taskId}\u0000${request.invocationId}`, + fingerprint: controlFingerprint(request.command), + }; + } + + private claimReplayKey(scopeId: string, taskId: string): string { + return `claim\u0000${shortHash(scopeId)}\u0000${taskId}`; + } + + /** + * Releases a retained result once a caller confirms holding it, so a delivered + * result frees its slot immediately instead of waiting out the replay window. + */ + private releaseClaimReplay(scopeId: string, taskId: string): void { + const key = this.claimReplayKey(scopeId, taskId); + const cached = this.claimReplays.entries.get(key); + if (cached == null) { + return; + } + this.claimReplays.entries.delete(key); + this.claimReplays.bytes -= cached.bytes; + } + + /** + * Tells the owner it may release a delivered result. Delivery to zero subscribers is + * not an acknowledgement, so this retries inside the ordinary request window. + */ + private async acknowledgeClaim(ownerId: string, scopeId: string, taskId: string): Promise { + const ack: RoutedAck = { version: PROTOCOL_VERSION, kind: 'ack', scopeId, taskId }; + const serialized = JSON.stringify(ack); + const destination = this.channel(ownerId); + const deadline = Date.now() + this.requestTimeoutMs; + for (;;) { + try { + if ((await this.publish(destination, serialized)) > 0) { + return; + } + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to acknowledge a claimed result', error); + } + if (Date.now() + this.retryDelayMs >= deadline) { + return; + } + await new Promise((resolve) => { + const timer = setTimeout(resolve, this.retryDelayMs); + timer.unref?.(); + }); + } + } + + private async requestTaskOwner( + scopeId: string, + taskId: string, + operation: 'claim' | 'control', + command?: SubagentTaskControlCommand, + invocationId?: string, + ): Promise<{ ownerId: string; result: unknown } | undefined> { + this.assertTaskAddress(scopeId, taskId); + await this.requireReady(); + let ownerId: string | null; + try { + ownerId = (await this.publisher.eval( + READ_TASK_OWNER_SCRIPT, + 1, + this.registryKey(scopeId), + taskId, + )) as string | null; + } catch (error) { + logger.warn('[subagentTaskRouting] Failed to resolve the task owner', error); + throw new SubagentTaskOwnerUnavailableError(); + } + if (!isBoundedString(ownerId, 128)) { + return undefined; + } + if (operation === 'claim') { + return { ownerId, result: await this.sendRequest(ownerId, { operation, scopeId, taskId }) }; + } + if (command == null || invocationId == null) { + throw new Error('A routed subagent control command and invocation id are required.'); + } + return { + ownerId, + result: await this.sendRequest(ownerId, { + operation, + scopeId, + taskId, + command, + invocationId, + }), + }; + } + + private async sendRequest(ownerId: string, request: RoutedRequestPayload): Promise { + await this.requireReady(); + if (this.pending.size >= MAX_PENDING_REQUESTS) { + throw new SubagentTaskOwnerUnavailableError(); + } + const requestId = randomUUID(); + /** Carried so a request the caller has stopped waiting for cannot be applied + * later: a disconnected publisher queues the envelope offline and delivers it + * after this deadline, by which time the caller has been told `unavailable`. */ + const envelope: RoutedRequest = { + version: PROTOCOL_VERSION, + kind: 'request', + requestId, + requesterId: this.instanceId, + expiresAt: Date.now() + this.requestTimeoutMs, + ...request, + }; + const serialized = JSON.stringify(envelope); + const destination = this.channel(ownerId); + return new Promise((resolve, reject) => { + const retry = setTimeout(() => { + void this.publish(destination, serialized).catch((error) => { + logger.warn('[subagentTaskRouting] Routed command retry failed', error); + }); + }, this.retryDelayMs); + retry.unref?.(); + const timeout = setTimeout(() => { + this.pending.delete(requestId); + clearTimeout(retry); + reject(new SubagentTaskOwnerUnavailableError()); + }, this.requestTimeoutMs); + timeout.unref?.(); + this.pending.set(requestId, { resolve, reject, retry, timeout }); + void this.publish(destination, serialized).catch((error) => { + logger.warn('[subagentTaskRouting] Routed command publish failed', error); + }); + }); + } + + private pruneExpiredReplays(cache: ReplayCache): void { + const now = Date.now(); + for (const [id, cached] of cache.entries) { + if (cached.expiresAt != null && cached.expiresAt <= now) { + cache.entries.delete(id); + cache.bytes -= cached.bytes; + } + } + } + + private retainReplay(cache: ReplayCache, key: string, value: string, fingerprint?: string): void { + const bytes = Buffer.byteLength(value, 'utf8'); + if (bytes > cache.maxBytes) { + return; + } + this.pruneExpiredReplays(cache); + /** Replacing a key is not an additional entry: leaving the old one counted would + * inflate the cache's byte total permanently and evict unrelated responses. */ + const replaced = cache.entries.get(key); + if (replaced != null) { + cache.entries.delete(key); + cache.bytes -= replaced.bytes; + } + while (cache.entries.size >= cache.maxEntries || cache.bytes + bytes > cache.maxBytes) { + const oldest = cache.entries.keys().next().value as string | undefined; + if (oldest == null) { + break; + } + const evicted = cache.entries.get(oldest); + cache.entries.delete(oldest); + cache.bytes -= evicted?.bytes ?? 0; + } + cache.entries.set(key, { + value, + bytes, + expiresAt: Date.now() + RESPONSE_CACHE_TTL_MS, + ...(fingerprint == null ? {} : { fingerprint }), + }); + cache.bytes += bytes; + } + + private async publish(channel: string, value: string): Promise { + const delivered = await this.publisher.publish(channel, value); + return typeof delivered === 'number' ? delivered : 0; + } + + private ensureRegistrationHeartbeat(): void { + if (this.registrationHeartbeat != null || this.destroyed) { + return; + } + this.registrationHeartbeat = setInterval(() => { + if (this.registrationRefresh != null) { + return; + } + const refresh = this.refreshRegistrations() + .catch((error) => { + logger.warn('[subagentTaskRouting] Failed to refresh child-task owners', error); + }) + .finally(() => { + if (this.registrationRefresh === refresh) { + this.registrationRefresh = undefined; + } + }); + this.registrationRefresh = refresh; + }, this.registrationHeartbeatMs); + this.registrationHeartbeat.unref?.(); + } + + private async refreshRegistrations(): Promise { + const handler = this.handler; + if (handler == null || this.destroyed || this.ownedTasks.size === 0) { + return; + } + const localTaskIdsByScope = new Map>(); + const staleTaskIdsByScope = new Map(); + const retained: OwnedTaskRegistration[] = []; + for (const registration of this.ownedTasks.values()) { + const { scopeId, taskId } = registration; + let localTaskIds = localTaskIdsByScope.get(scopeId); + if (localTaskIds == null) { + localTaskIds = new Set(handler.list(scopeId).map((task) => task.taskId)); + localTaskIdsByScope.set(scopeId, localTaskIds); + } + /** A retained result is only reachable while its owner stays registered, so the + * address outlives the task itself until the result is acknowledged. */ + if ( + localTaskIds.has(taskId) || + this.claimReplays.entries.has(this.claimReplayKey(scopeId, taskId)) + ) { + retained.push(registration); + continue; + } + this.ownedTasks.delete(this.registrationKey(scopeId, taskId)); + const staleTaskIds = staleTaskIdsByScope.get(scopeId) ?? []; + staleTaskIds.push(taskId); + staleTaskIdsByScope.set(scopeId, staleTaskIds); + } + /** Serializing one EVAL per registration can outlast the lease TTL, so a pass + * refreshes in bounded parallel batches and one failure cannot cancel the rest. */ + const refreshSlot = createConcurrencyLimiter(ROUTING_FANOUT_CONCURRENCY); + await Promise.all([ + ...[...staleTaskIdsByScope].map(([scopeId, taskIds]) => + refreshSlot(() => this.removeRegistrations(scopeId, taskIds)), + ), + ...retained.map((registration) => + refreshSlot(() => + this.publishRegistration(registration).catch((error) => { + logger.warn('[subagentTaskRouting] Failed to refresh a child-task owner', error); + }), + ), + ), + ]); + } + + private async publishRegistration(registration: OwnedTaskRegistration): Promise { + await this.requireReady(); + await this.publisher.eval( + REGISTER_TASK_SCRIPT, + 1, + this.registryKey(registration.scopeId), + registration.taskId, + this.instanceId, + registration.ttlMs.toString(), + ); + } + + private async readActiveRegistrations(scopeId: string): Promise> { + const value = (await this.publisher.eval( + READ_ACTIVE_REGISTRATIONS_SCRIPT, + 1, + this.registryKey(scopeId), + )) as unknown; + if (!Array.isArray(value) || value.length % 2 !== 0) { + throw new SubagentTaskOwnerUnavailableError(); + } + const registrations: Record = {}; + for (let index = 0; index < value.length; index += 2) { + const taskId = value[index]; + const ownerId = value[index + 1]; + if (!isBoundedString(taskId, MAX_TASK_ID_CHARS) || !isBoundedString(ownerId, 128)) { + throw new SubagentTaskOwnerUnavailableError(); + } + registrations[taskId] = ownerId; + } + return registrations; + } + + private async removeRegistrations(scopeId: string, taskIds: string[]): Promise { + if (taskIds.length === 0) { + return; + } + await this.publisher.hdel(this.registryKey(scopeId), ...taskIds).catch((error) => { + logger.warn('[subagentTaskRouting] Failed to prune stale task owners', error); + }); + } + + private registryKey(scopeId: string): string { + return `subagent-task:{${shortHash(scopeId)}}:owners`; + } + + private registrationKey(scopeId: string, taskId: string): string { + return `${scopeId}\u0000${taskId}`; + } + + private channel(instanceId: string): string { + return `subagent-task-control:${this.namespaceHash}:${instanceId}`; + } + + private assertScope(scopeId: string): void { + if (!isBoundedString(scopeId, MAX_SCOPE_ID_CHARS)) { + throw new Error('Invalid subagent task routing scope.'); + } + } + + private assertTaskAddress(scopeId: string, taskId: string): void { + this.assertScope(scopeId); + if (!isBoundedString(taskId, MAX_TASK_ID_CHARS)) { + throw new Error('Invalid subagent task routing identity.'); + } + } + + private async requireReady(): Promise { + if (this.destroyed || this.ready == null) { + throw new SubagentTaskOwnerUnavailableError(); + } + await this.ready; + } +} diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index b75ec3b7ee..a68b3bb054 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -5,14 +5,27 @@ import { Constants, EModelEndpoint } from 'librechat-data-provider'; import { createMethods, createModels, logger } from '@librechat/data-schemas'; import { AIMessage, HumanMessage } from '@librechat/agents/langchain/messages'; import type { + SubagentTaskClaim, + SubagentTaskControlCommand, + SubagentTaskControlResult, SubagentTaskRuntime, + SubagentTaskSnapshot, SubagentTaskStartRequest, SubagentTaskStartResult, } from '@librechat/agents'; import type { AllMethods, IConversation, IMessage } from '@librechat/data-schemas'; import type { BaseMessage } from '@librechat/agents/langchain/messages'; +import type { + SubagentTaskControlHandler, + SubagentTaskControlTransport, +} from './subagentTaskRouting'; import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; -import { buildSubagentThreadTaskConfig, SubagentThreadTaskStore } from './subagentThreads'; +import { + buildSubagentThreadTaskConfig, + createSubagentThreadTaskStore, + SubagentThreadTaskStore, +} from './subagentThreads'; +import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting'; import { createSubagentAttemptKey } from './subagentThreadIds'; import { createSubagentUsageSink } from './usage'; @@ -20,6 +33,75 @@ let mongod: MongoMemoryServer; let methods: AllMethods; let loggerErrorSpy: jest.SpyInstance; +class TestTaskRoutingHub { + readonly owners = new Map(); + + key(scopeId: string, taskId: string): string { + return `${scopeId}\u0000${taskId}`; + } +} + +class TestTaskControlTransport implements SubagentTaskControlTransport { + private handler?: SubagentTaskControlHandler; + readonly registrations: Array<{ scopeId: string; taskId: string; ttlMs: number }> = []; + + constructor(private readonly hub: TestTaskRoutingHub) {} + + async bind(handler: SubagentTaskControlHandler): Promise { + this.handler = handler; + } + + async registerTask(scopeId: string, taskId: string, ttlMs: number): Promise { + this.registrations.push({ scopeId, taskId, ttlMs }); + this.hub.owners.set(this.hub.key(scopeId, taskId), this); + } + + async hasTasks(scopeId: string): Promise { + const prefix = `${scopeId}\u0000`; + return [...this.hub.owners.keys()].some((key) => key.startsWith(prefix)); + } + + async claim(scopeId: string, taskId: string): Promise { + return this.hub.owners.get(this.hub.key(scopeId, taskId))?.handler?.claim(scopeId, taskId); + } + + async control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + _invocationId: string, + ): Promise { + return this.hub.owners + .get(this.hub.key(scopeId, taskId)) + ?.handler?.control(scopeId, taskId, command, _invocationId); + } + + async list(scopeId: string): Promise { + return [...this.remoteOwners(scopeId)].flatMap((owner) => owner.handler?.list(scopeId) ?? []); + } + + async cancelScope(scopeId: string, threadIds: string[] | null): Promise { + let cancelled = 0; + for (const owner of this.remoteOwners(scopeId)) { + cancelled += owner.handler?.cancelScope(scopeId, threadIds) ?? 0; + } + return cancelled; + } + + async destroy(): Promise {} + + private remoteOwners(scopeId: string): Set { + const owners = new Set(); + const prefix = `${scopeId}\u0000`; + for (const [key, owner] of this.hub.owners) { + if (key.startsWith(prefix) && owner !== this) { + owners.add(owner); + } + } + return owners; + } +} + function taskRequest( scopeId: string, overrides: Partial = {}, @@ -51,6 +133,32 @@ function taskRequest( }; } +function replayTransport(claim: SubagentTaskClaim): SubagentTaskControlTransport { + return { + bind: async () => undefined, + registerTask: async () => undefined, + hasTasks: async () => true, + claim: async () => claim, + control: async () => undefined, + list: async () => [], + cancelScope: async () => 0, + destroy: async () => undefined, + }; +} + +function threadSnapshot(taskId: string): SubagentTaskSnapshot { + return { + taskId, + subagentType: 'researcher', + status: 'cancelled', + createdAt: 1, + updatedAt: 2, + resultAvailable: false, + resultClaimed: true, + pendingControls: 0, + }; +} + async function waitForSettled( store: SubagentThreadTaskStore, scopeId: string, @@ -67,6 +175,16 @@ async function waitForSettled( throw new Error('Timed out waiting for the subagent task.'); } +async function waitUntil(condition: () => boolean, description: string): Promise { + for (let attempt = 0; attempt < 400; attempt += 1) { + if (condition()) { + return; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error(`Timed out waiting for ${description}.`); +} + function requireAccepted( started: SubagentTaskStartResult, ): Extract { @@ -583,8 +701,26 @@ describe('SubagentThreadTaskStore', () => { const preparationRelease = new Promise((resolve) => { releasePreparation = resolve; }); + let leaseDeadline = new Date(0); + let renewedPastDeadline = false; const slowMethods = { ...methods, + acquireSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + const acquired = await methods.acquireSubagentThreadLease(...args); + if (acquired) { + leaseDeadline = args[0].expiresAt; + } + return acquired; + }, + ), + renewSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + const renewed = await methods.renewSubagentThreadLease(...args); + renewedPastDeadline ||= renewed && args[0].now > leaseDeadline; + return renewed; + }, + ), getMessages: jest.fn(async (...args: Parameters) => { if (blockNextRead && args[0].conversationId === slowThreadId) { blockNextRead = false; @@ -594,7 +730,7 @@ describe('SubagentThreadTaskStore', () => { return methods.getMessages(...args); }), }; - const options = { leaseTtlMs: 60, leaseHeartbeatMs: 10 }; + const options = { leaseTtlMs: 500, leaseHeartbeatMs: 50 }; const firstWorker = new SubagentThreadTaskStore(slowMethods, options); const secondWorker = new SubagentThreadTaskStore(methods, options); const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); @@ -612,7 +748,9 @@ describe('SubagentThreadTaskStore', () => { }), ); await preparing; - await new Promise((resolve) => setTimeout(resolve, 100)); + /** Wait for evidence rather than a fixed delay: a renewal that succeeds after the + * acquired lease's own deadline proves the heartbeat carried it past expiry. */ + await waitUntil(() => renewedPastDeadline, 'the shared lease to outlive its original deadline'); const overlappingRun = jest.fn(taskRequest(config.scopeId).run); const overlapping = secondWorker.start( @@ -630,6 +768,61 @@ describe('SubagentThreadTaskStore', () => { expect(firstRun).toHaveBeenCalledTimes(1); }); + it('cancels a child when its lease renewal only commits after expiry', async () => { + const userId = 'late-lease-renewal-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let providerEntered = false; + let previousExpiry = 0; + let markLateRenewal = (): void => undefined; + const lateRenewal = new Promise((resolve) => { + markLateRenewal = resolve; + }); + const slowMethods = { + ...methods, + renewSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + if (providerEntered) { + markLateRenewal(); + /** Wait on the last confirmed lease deadline rather than a tiny fixed TTL: + * the renewal definitely commits after the gap, without assuming how fast + * a loaded runner completes preparation and its first Mongo write. */ + await new Promise((resolve) => + setTimeout(resolve, Math.max(0, previousExpiry - Date.now() + 10)), + ); + } + const renewed = await methods.renewSubagentThreadLease(...args); + if (renewed) { + previousExpiry = args[0].expiresAt.getTime(); + } + return renewed; + }, + ), + }; + const store = new SubagentThreadTaskStore(slowMethods, { + leaseTtlMs: 1_000, + leaseHeartbeatMs: 20, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(async (runtime: SubagentTaskRuntime) => { + providerEntered = true; + return new Promise<{ content: string }>((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }); + }); + const started = store.start(taskRequest(config.scopeId, { run })); + + await lateRenewal; + await waitForSettled(store, config.scopeId, started); + + expect(run).toHaveBeenCalledTimes(1); + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'cancelled', + }); + }); + it('rechecks account deletion after acquiring the shared lease', async () => { const userId = 'lease-fence-gap-user'; const parentConversationId = randomUUID(); @@ -1297,6 +1490,999 @@ describe('SubagentThreadTaskStore', () => { expect(JSON.stringify(loggerErrorSpy.mock.calls)).not.toContain('provider-secret'); }); + it('routes live task polling and controls to the replica that owns the execution', async () => { + const userId = 'routed-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const requesterStore = new SubagentThreadTaskStore(methods); + const ownerTransport = new TestTaskControlTransport(hub); + await ownerStore.configureTaskControlTransport(ownerTransport); + await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => result, + }), + ); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + expect(ownerTransport.registrations).toContainEqual({ + scopeId: config.scopeId, + taskId, + ttlMs: 30_000, + }); + + await expect(requesterStore.hasTasks(config.scopeId)).resolves.toBe(true); + await expect(requesterStore.listTasks(config.scopeId)).resolves.toEqual([ + expect.objectContaining({ taskId, status: 'running' }), + ]); + await expect( + requesterStore.controlTask(config.scopeId, taskId, { + action: 'queue', + message: 'Verify the primary source too.', + }), + ).resolves.toMatchObject({ status: 'accepted' }); + + finish({ content: 'Cross-replica result.' }); + await waitForSettled(ownerStore, config.scopeId, started); + await expect(requesterStore.claimTask(config.scopeId, taskId)).resolves.toMatchObject({ + status: 'completed', + result: 'Cross-replica result.', + }); + await expect(requesterStore.claimTask(config.scopeId, taskId)).resolves.toMatchObject({ + status: 'claimed', + }); + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + requesterStore.destroyTaskControlTransport(), + ]); + }); + + it('applies one control invocation once whether it arrives locally or through routing', async () => { + const userId = 'invocation-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const requesterStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => result, + }), + ); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + + const steer = { action: 'queue' as const, message: 'Verify the primary source too.' }; + const routed = await requesterStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'); + expect(routed).toMatchObject({ status: 'accepted' }); + + /** The same invocation reaching the owner directly replays that result rather than + * queueing a second steer, so local and routed callers agree. */ + await expect( + ownerStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'), + ).resolves.toEqual(routed); + expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(1); + + /** Reusing one invocation id for different content is a caller error, not a retry. */ + await expect( + requesterStore.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Something else entirely.' }, + 'invocation-1', + ), + ).resolves.toMatchObject({ status: 'invalid' }); + expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(1); + + finish({ content: 'Cross-replica result.' }); + await waitForSettled(ownerStore, config.scopeId, started); + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + requesterStore.destroyTaskControlTransport(), + ]); + }); + + it('fails a child closed when its owner address cannot be published', async () => { + const userId = 'unregistered-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const transport = new TestTaskControlTransport(hub); + transport.registerTask = async () => { + throw new Error('registration failed'); + }; + const store = new SubagentThreadTaskStore(methods); + await store.configureTaskControlTransport(transport); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(async () => ({ content: 'never reached' })); + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + /** An unaddressable child cannot be polled, controlled, or cancelled, so no + * provider work may start behind a failed registration. */ + expect(run).not.toHaveBeenCalled(); + expect(store.get(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'error', + }); + await store.destroyTaskControlTransport(); + }); + + it('returns a lost result to the same poll invocation and refuses a different one', async () => { + const userId = 'durable-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const requesterStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Cross-replica result.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(ownerStore, config.scopeId, started); + + await expect(requesterStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { status: 'completed', result: 'Cross-replica result.' }, + ); + + /** The owner's one-shot result is gone, but the child's durable thread still holds + * it, so the invocation that already collected it recovers its own result. */ + await expect(requesterStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { status: 'completed', result: 'Cross-replica result.' }, + ); + + /** A different invocation is told it was collected rather than handed a copy. */ + await expect(requesterStore.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject( + { status: 'claimed' }, + ); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + requesterStore.destroyTaskControlTransport(), + ]); + }); + + it('bounds a result recovered from its durable child message', async () => { + const userId = 'large-result-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'x'.repeat(150_000) }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject({ + status: 'completed', + }); + + /** The durable message keeps the child's untruncated output, so recovering it + * must apply the same bound a routed response would have. */ + const recovered = await store.claimTask(config.scopeId, taskId, 'poll-1'); + expect(recovered.status).toBe('completed'); + if (recovered.status === 'completed') { + expect(recovered.result.length).toBeLessThanOrEqual(100_000); + } + }); + + it('recovers a durable result after the owning process and registration are gone', async () => { + const userId = 'restarted-owner-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const ownerStore = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Recovered without owner memory.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(ownerStore, config.scopeId, started); + + /** A fresh store has neither the in-memory task nor a Redis owner registration. */ + const restartedStore = new SubagentThreadTaskStore(methods); + const unrelatedParentConversationId = randomUUID(); + await saveParent(userId, unrelatedParentConversationId); + const unrelatedConfig = buildSubagentThreadTaskConfig(restartedStore, { + userId, + parentConversationId: unrelatedParentConversationId, + }); + await expect( + restartedStore.claimTask(unrelatedConfig.scopeId, taskId, 'wrong-parent-poll'), + ).resolves.toEqual({ status: 'not_found' }); + + await expect(restartedStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { + status: 'completed', + result: 'Recovered without owner memory.', + }, + ); + await expect(restartedStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { + status: 'completed', + result: 'Recovered without owner memory.', + }, + ); + await expect(restartedStore.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject( + { + status: 'claimed', + }, + ); + }); + + it('tells a second invocation a retained result was already collected', async () => { + const userId = 'duplicate-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Only one caller may hold this.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + const threadId = requireThreadId(started); + + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject({ + status: 'completed', + result: 'Only one caller may hold this.', + }); + + /** An owner replaying a retained response would hand the same terminal result to + * another invocation; the durable record decides, so that one is told it was + * already collected rather than being given a second copy. */ + const replayingStore = new SubagentThreadTaskStore(methods); + await replayingStore.configureTaskControlTransport( + replayTransport({ + status: 'completed', + task: { ...threadSnapshot(taskId), threadId, status: 'completed' }, + result: 'Only one caller may hold this.', + }), + ); + await expect(replayingStore.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject( + { status: 'claimed' }, + ); + + /** The invocation that already holds it still recovers its own result. */ + await expect(replayingStore.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject( + { status: 'completed', result: 'Only one caller may hold this.' }, + ); + + await replayingStore.destroyTaskControlTransport(); + }); + + it('refuses to build a store the host wired without a required method', () => { + const { claimSubagentTaskResult: _omitted, ...incomplete } = methods; + + /** The host wires this from JavaScript, so a missing method has to fail at + * startup rather than as an unavailable result the first time a task settles. */ + expect(() => + createSubagentThreadTaskStore( + incomplete as unknown as Parameters[0], + ), + ).toThrow('claimSubagentTaskResult'); + expect(() => createSubagentThreadTaskStore(methods)).not.toThrow(); + }); + + it('renews its own fence while a long deletion is still running', async () => { + const userId = 'long-deletion-user'; + const renewals: string[] = []; + const store = new SubagentThreadTaskStore(methods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 30, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + renewOwnerAdmission: async (_userId: string, token: string) => { + renewals.push(token); + return true; + }, + releaseOwnerAdmission: async () => undefined, + }); + const listLeases = jest.spyOn(methods, 'listActiveSubagentThreadLeases').mockResolvedValue([]); + try { + await store.withOwnerDeletionFence(userId, undefined, async () => { + /** A deletion outlasting its 90ms fence window must not let the fence lapse. */ + await new Promise((resolve) => setTimeout(resolve, 300)); + return 'deleted'; + }); + } finally { + listLeases.mockRestore(); + } + + expect(renewals.length).toBeGreaterThan(0); + expect(new Set(renewals).size).toBe(1); + }); + + it('re-fences and drains again after a fence gap during deletion', async () => { + const userId = 'deletion-gap-user'; + let recoveryAllowed = false; + const fenceOwnerAdmission = jest.fn(async () => undefined); + const listActiveSubagentThreadLeases = jest.fn(async () => []); + const testMethods = { ...methods, listActiveSubagentThreadLeases }; + const renewOwnerAdmission = jest.fn(async () => { + if (!recoveryAllowed) { + throw new Error('database temporarily unavailable'); + } + return false; + }); + const store = new SubagentThreadTaskStore(testMethods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 30, + ownerFenceGraceMs: 60, + fenceOwnerAdmission, + renewOwnerAdmission, + releaseOwnerAdmission: async () => undefined, + }); + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + await new Promise((resolve) => setTimeout(resolve, 150)); + recoveryAllowed = true; + return 'deleted'; + }), + ).resolves.toBe('deleted'); + + expect(fenceOwnerAdmission).toHaveBeenCalledTimes(2); + expect(listActiveSubagentThreadLeases).toHaveBeenCalledTimes(2); + }); + + it('does not report deletion success when post-gap recovery fails', async () => { + const userId = 'deletion-gap-failure-user'; + const listActiveSubagentThreadLeases = jest.fn(async () => []); + const testMethods = { ...methods, listActiveSubagentThreadLeases }; + const store = new SubagentThreadTaskStore(testMethods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 30, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + renewOwnerAdmission: async () => { + throw new Error('database unavailable'); + }, + releaseOwnerAdmission: async () => undefined, + }); + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + await new Promise((resolve) => setTimeout(resolve, 150)); + return 'deleted'; + }), + ).rejects.toThrow('database unavailable'); + }); + + it('cancels a grandchild whose own conversation the cascade removed', async () => { + const userId = 'cascade-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods, { maxThreadDepth: 3 }); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const childConversationId = randomUUID(); + await saveParent(userId, childConversationId, { + subagentThread: { + rootConversationId: parentConversationId, + parentConversationId, + parentAgentId: 'parent-agent', + subagentType: 'researcher', + depth: 1, + }, + }); + /** The grandchild runs inside the child's scope, which a plan naming only the + * deleted root never covers. */ + const config = buildSubagentThreadTaskConfig(ownerStore, { + userId, + parentConversationId: childConversationId, + }); + let finish = (_value: { content: string }): void => undefined; + const running = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start(taskRequest(config.scopeId, { run: async () => running })); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + + const plan = await deletingStore.planCancellationForConversations(userId, [ + parentConversationId, + ]); + await expect( + deletingStore.cancelPlan(plan, [parentConversationId, childConversationId]), + ).resolves.toBeGreaterThanOrEqual(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + finish({ content: 'late' }); + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('reports a result as unavailable when its collection cannot be recorded', async () => { + const userId = 'unrecordable-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ content: 'Recorded before it is handed over.' }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + + /** Handing the result over without recording its claimant would let another + * invocation collect the same one-shot output once the database recovers. */ + const claimResult = jest + .spyOn(methods, 'claimSubagentTaskResult') + .mockRejectedValueOnce(new Error('database unavailable')); + try { + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + } finally { + claimResult.mockRestore(); + } + + /** The result stays unclaimed, so a later poll still collects it exactly once. */ + await expect(store.claimTask(config.scopeId, taskId, 'poll-1')).resolves.toMatchObject({ + status: 'completed', + result: 'Recorded before it is handed over.', + }); + await expect(store.claimTask(config.scopeId, taskId, 'poll-2')).resolves.toMatchObject({ + status: 'claimed', + }); + }); + + it('keeps a live task’s control invocation when settled tasks fill the window', async () => { + const userId = 'invocation-eviction-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { + maxControlInvocations: 2, + completedTtlMs: 20, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const running = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const live = store.start(taskRequest(config.scopeId, { run: async () => running })); + const liveTaskId = requireAccepted(live).task.taskId; + const settled = store.start( + taskRequest(config.scopeId, { run: async () => ({ content: 'done' }) }), + ); + const settledTaskId = requireAccepted(settled).task.taskId; + await waitForSettled(store, config.scopeId, settled); + + const steer = { action: 'queue' as const, message: 'Verify the primary source too.' }; + const applied = store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-live'); + expect(applied).toMatchObject({ status: 'accepted' }); + store.controlInvocation(config.scopeId, settledTaskId, steer, 'invocation-settled'); + + for (let attempt = 0; attempt < 100; attempt += 1) { + if (store.get(config.scopeId, settledTaskId) == null) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + expect(store.get(config.scopeId, settledTaskId)).toBeUndefined(); + + /** The window is full, so admitting another invocation sweeps the records of tasks + * this store no longer holds. The live task's record survives, so a caller + * retrying it replays instead of steering that child a second time. */ + store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-later'); + expect(store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-live')).toEqual( + applied, + ); + expect(store.get(config.scopeId, liveTaskId)?.pendingControls).toBe(2); + + /** With every remaining record belonging to a live task, a further invocation is + * refused rather than displacing one: applying it unrecorded would let its own + * retry apply the command twice. */ + expect( + store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-third'), + ).toMatchObject({ status: 'invalid' }); + expect(store.get(config.scopeId, liveTaskId)?.pendingControls).toBe(2); + + finish({ content: 'done' }); + await waitForSettled(store, config.scopeId, live); + }); + + it('caps the merged local and remote task list the poll tool reads', async () => { + const userId = 'merged-list-cap-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + /** The base store caps concurrent runs twice over — ten per scope and a hundred + * across the store — and this test is about what the merge returns rather than + * about admission, so both are raised to admit every task it starts. */ + const store = new SubagentThreadTaskStore(methods, { + maxRunningPerScope: 150, + maxRunningTotal: 150, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const remote = Array.from({ length: 150 }, (_unused, index) => + threadSnapshot(`remote-task-${index + 1}`), + ); + await store.configureTaskControlTransport({ + ...replayTransport({ status: 'claimed', task: threadSnapshot('remote-task-1') }), + list: async () => remote, + }); + + const local = await Promise.all( + Array.from({ length: 150 }, () => store.start(taskRequest(config.scopeId))), + ); + await Promise.all(local.map((started) => waitForSettled(store, config.scopeId, started))); + expect(store.list(config.scopeId)).toHaveLength(150); + + /** Each owner's reply and the remote aggregation are bounded on their own, but the + * poll tool reads this merge — 300 distinct tasks must still arrive as 200. */ + await expect(store.listTasks(config.scopeId)).resolves.toHaveLength(200); + + await store.destroyTaskControlTransport(); + }); + + it('routes a control for a remote task while the local invocation window is full', async () => { + const userId = 'remote-control-under-load-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { maxControlInvocations: 1 }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const remoteResult: SubagentTaskControlResult = { + status: 'cancelled', + task: threadSnapshot('remote-task'), + }; + const routed = jest.fn(async () => remoteResult); + await store.configureTaskControlTransport({ + ...replayTransport({ status: 'claimed', task: threadSnapshot('remote-task') }), + control: routed, + }); + + let finish = (_value: { content: string }): void => undefined; + const running = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const live = store.start(taskRequest(config.scopeId, { run: async () => running })); + const liveTaskId = requireAccepted(live).task.taskId; + const steer = { action: 'queue' as const, message: 'Check the changelog as well.' }; + expect(store.controlInvocation(config.scopeId, liveTaskId, steer, 'local-1')).toMatchObject({ + status: 'accepted', + }); + + /** The window holds a live task's record and cannot be swept, but a task this + * replica never owned is the remote owner's to refuse or apply. */ + await expect( + store.controlTask(config.scopeId, 'remote-task', { action: 'cancel' }, 'remote-1'), + ).resolves.toEqual(remoteResult); + expect(routed).toHaveBeenCalledWith( + config.scopeId, + 'remote-task', + { action: 'cancel' }, + 'remote-1', + ); + + finish({ content: 'done' }); + await waitForSettled(store, config.scopeId, live); + await store.destroyTaskControlTransport(); + }); + + it('fails a deletion closed when the admission fence cannot be held', async () => { + const userId = 'fence-lapse-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { + ownerDrainTimeoutMs: 60, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + renewOwnerAdmission: async () => { + throw new Error('database unavailable'); + }, + releaseOwnerAdmission: async () => undefined, + }); + /** A drain that outlasts the 120ms fence window while every renewal rejects: the + * last confirmed `fencedUntil` passes and nothing is left holding admission shut. */ + const leases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockImplementationOnce(async () => { + await new Promise((resolve) => setTimeout(resolve, 200)); + return []; + }); + const deletion = jest.fn(async () => 'deleted'); + try { + await expect(store.withOwnerDeletionFence(userId, undefined, deletion)).rejects.toThrow( + 'admission fence expired', + ); + /** Nothing was removed, so the caller can retry once the fence holds again. */ + expect(deletion).not.toHaveBeenCalled(); + } finally { + leases.mockRestore(); + } + }); + + it('treats a renewal that lands after its own deadline as a lapse', async () => { + const userId = 'fence-late-renewal-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let renewals = 0; + const store = new SubagentThreadTaskStore(methods, { + ownerDrainTimeoutMs: 60, + ownerFenceGraceMs: 60, + fenceOwnerAdmission: async () => undefined, + /** Succeeds, but the first write only lands well past the 120ms deadline it was + * meant to extend — admission stood open for the difference. */ + renewOwnerAdmission: async () => { + renewals += 1; + if (renewals === 1) { + await new Promise((resolve) => setTimeout(resolve, 150)); + } + return true; + }, + releaseOwnerAdmission: async () => undefined, + }); + const leases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockImplementationOnce(async () => { + await new Promise((resolve) => setTimeout(resolve, 250)); + return []; + }); + const deletion = jest.fn(async () => 'deleted'); + try { + await expect(store.withOwnerDeletionFence(userId, undefined, deletion)).rejects.toThrow( + 'admission fence expired', + ); + /** Every renewal reported success, so a deadline restored from the write's own + * start time would have read as continuously fenced. */ + expect(renewals).toBeGreaterThan(0); + expect(deletion).not.toHaveBeenCalled(); + } finally { + leases.mockRestore(); + } + }); + + it('releases the owner fence after an in-flight renewal instead of racing it', async () => { + const userId = 'fence-renewal-race-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let releaseRenewal = (): void => undefined; + const renewalBlocked = new Promise((resolve) => { + releaseRenewal = resolve; + }); + let markRenewing = (): void => undefined; + const renewing = new Promise((resolve) => { + markRenewing = resolve; + }); + const order: string[] = []; + const fenceOwnerAdmission = jest.fn(async () => { + order.push('fence'); + }); + /** The renewal is still waiting on the database when the deletion finishes, and it + * reports the fence lost — the shape that used to leave a fresh, unreleasable one. */ + let renewalAttempts = 0; + const renewOwnerAdmission = jest.fn(async () => { + renewalAttempts += 1; + markRenewing(); + await renewalBlocked; + order.push('renew'); + /** The in-flight renewal discovers the entry missing and re-takes it; the + * recovery renewal then confirms that replacement while the second drain runs. */ + return renewalAttempts > 1; + }); + const releaseOwnerAdmission = jest.fn(async () => { + order.push('release'); + }); + const testMethods = { + ...methods, + listActiveSubagentThreadLeases: jest.fn(async () => []), + }; + const store = new SubagentThreadTaskStore(testMethods, { + ownerDrainTimeoutMs: 60, + ownerFenceGraceMs: 60, + fenceOwnerAdmission, + renewOwnerAdmission, + releaseOwnerAdmission, + }); + + let releaseDeletion = (): void => undefined; + const deletionBlocked = new Promise((resolve) => { + releaseDeletion = resolve; + }); + const fenced = store.withOwnerDeletionFence(userId, undefined, async () => { + await deletionBlocked; + return 'deleted'; + }); + await renewing; + releaseDeletion(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(order).toEqual(['fence']); + + releaseRenewal(); + await expect(fenced).resolves.toBe('deleted'); + /** The lost entry is re-taken before the recovery drain and only released after + * the in-flight renewal and recovery renewal both settle. */ + expect(order).toEqual(['fence', 'renew', 'fence', 'renew', 'release']); + expect(fenceOwnerAdmission).toHaveBeenCalledTimes(2); + }); + + it('cancels each drained task once and retries only unconfirmed deliveries', async () => { + const userId = 'drain-user'; + const parentConversationId = randomUUID(); + const store = new SubagentThreadTaskStore(methods, { + ownerDrainPollMs: 1, + ownerDrainTimeoutMs: 5_000, + }); + const lease = { taskId: 'task-1', parentConversationId, conversationId: randomUUID() }; + const listLeases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockResolvedValueOnce([lease]) + .mockResolvedValueOnce([lease]) + .mockResolvedValueOnce([lease]) + .mockResolvedValueOnce([lease]) + .mockResolvedValue([]); + const controlTask = jest + .spyOn(store, 'controlTask') + .mockRejectedValueOnce(new Error('owner unavailable')) + .mockResolvedValueOnce({ status: 'not_found' }) + .mockResolvedValue({ status: 'cancelled', task: threadSnapshot('task-1') }); + try { + await store.cancelAndDrainForOwner(userId); + + /** An unconfirmed delivery is retried under the same invocation — including a + * `not_found`, which means the owner's registration is missing while its lease + * is live — and once the owner confirms, the drain only waits for the lease. */ + expect(controlTask).toHaveBeenCalledTimes(3); + expect(new Set(controlTask.mock.calls.map((call) => call[3])).size).toBe(1); + expect(listLeases).toHaveBeenCalledTimes(5); + } finally { + listLeases.mockRestore(); + controlTask.mockRestore(); + } + }); + + it('fences owner admission around the deletion it drains for', async () => { + const userId = 'fenced-user'; + const order: string[] = []; + const tokens: string[] = []; + const renewed: string[] = []; + const released: string[] = []; + const store = new SubagentThreadTaskStore(methods, { + ownerDrainPollMs: 1, + fenceOwnerAdmission: async (_userId: string, token: string) => { + tokens.push(token); + order.push('fence'); + }, + renewOwnerAdmission: async (_userId: string, token: string) => { + renewed.push(token); + return true; + }, + releaseOwnerAdmission: async (_userId: string, token: string) => { + released.push(token); + order.push('release'); + }, + }); + const listLeases = jest + .spyOn(methods, 'listActiveSubagentThreadLeases') + .mockImplementation(async () => { + order.push('drain'); + return []; + }); + try { + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + order.push('delete'); + return 'deleted'; + }), + ).resolves.toBe('deleted'); + expect(order).toEqual(['fence', 'drain', 'delete', 'release']); + /** Only the fence this deletion took is lifted, so an overlapping deletion + * keeps admission closed until its own fence is released. */ + expect(released).toEqual(tokens); + expect(tokens[0]).toEqual(expect.any(String)); + + /** A failed deletion still lifts the fence, so one bad request cannot leave the + * account unable to run subagents. */ + order.length = 0; + await expect( + store.withOwnerDeletionFence(userId, undefined, async () => { + throw new Error('deletion failed'); + }), + ).rejects.toThrow('deletion failed'); + expect(order).toEqual(['fence', 'drain', 'release']); + } finally { + listLeases.mockRestore(); + } + }); + + it('routes conversation-deletion cancellation to a remote task owner', async () => { + const userId = 'routed-delete-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + await Promise.resolve(); + + const plan = await deletingStore.planCancellationForConversations(userId, [ + parentConversationId, + ]); + await expect(deletingStore.cancelPlan(plan)).resolves.toBe(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('routes cancellation for a deleted child thread to its remote owner', async () => { + const userId = 'routed-child-delete-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + for (let attempt = 0; attempt < 200; attempt += 1) { + if ((await methods.getConvo(userId, threadId)) != null) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + expect(await methods.getConvo(userId, threadId)).not.toBeNull(); + + /** The parent survives this deletion, so the child's own thread is the only target. */ + const plan = await deletingStore.planCancellationForConversations(userId, [threadId]); + await expect(deletingStore.cancelPlan(plan)).resolves.toBe(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('cancels a child admitted after the deletion snapshot from its durable lease', async () => { + const userId = 'lease-cancel-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods); + const deletingStore = new SubagentThreadTaskStore(methods); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + for (let attempt = 0; attempt < 200; attempt += 1) { + const leases = await methods.listActiveSubagentThreadLeases({ + user: userId, + now: new Date(), + }); + if (leases.length > 0) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + + /** Production ordering: the plan is resolved first, the cascade is deleted, and + * only then is the plan replayed against the owner directory. */ + const plan = await deletingStore.planCancellationForConversations(userId, [ + parentConversationId, + ]); + await methods.deleteConvos(userId, { conversationId: parentConversationId }); + await expect( + deletingStore.cancelPlan(plan, [parentConversationId, requireThreadId(started)]), + ).resolves.toBeGreaterThanOrEqual(1); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + + it('drains only active lease addresses when deleting every conversation across replicas', async () => { + const userId = 'routed-owner-drain-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const hub = new TestTaskRoutingHub(); + const ownerStore = new SubagentThreadTaskStore(methods, { ownerDrainPollMs: 5 }); + const deletingStore = new SubagentThreadTaskStore(methods, { ownerDrainPollMs: 5 }); + await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const started = ownerStore.start( + taskRequest(config.scopeId, { + run: async (runtime) => { + markEntered(); + return new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }); + }, + }), + ); + await entered; + + await deletingStore.cancelAndDrainForOwner(userId); + await waitForSettled(ownerStore, config.scopeId, started); + expect(ownerStore.get(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'cancelled', + }); + + await Promise.all([ + ownerStore.destroyTaskControlTransport(), + deletingStore.destroyTaskControlTransport(), + ]); + }); + it('bounds durable delegation depth to one by default', async () => { const userId = 'depth-user'; const rootConversationId = randomUUID(); diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index dc0f047706..415ef9533f 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -8,24 +8,36 @@ import { } from '@librechat/agents/langchain/messages'; import type { InMemorySubagentTaskStoreOptions, + SubagentTaskClaim, SubagentTaskConfig, SubagentTaskControlCommand, SubagentTaskControlResult, SubagentTaskRuntime, + SubagentTaskSnapshot, SubagentTaskStartRequest, SubagentTaskStartResult, } from '@librechat/agents'; import type { AllMethods, + IActiveSubagentThreadLease, IConversation, IMessage, MessageMethods, ConversationMethods, + SubagentTaskResultClaim, } from '@librechat/data-schemas'; import type { BaseMessage, StoredMessage } from '@librechat/agents/langchain/messages'; +import type { SubagentTaskControlTransport } from './subagentTaskRouting'; import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; +import { + boundedClaim, + boundedTaskList, + controlFingerprint, + SubagentTaskOwnerUnavailableError, +} from './subagentTaskRouting'; import { createSubagentAttemptKey, createSubagentThreadId } from './subagentThreadIds'; import { runWithDetachedSubagentUsage } from './subagentTaskContext'; +import { createConcurrencyLimiter } from '~/utils/promise'; import { aggregateEmittedUsage } from './usage'; const SCOPE_VERSION = 1; @@ -33,10 +45,29 @@ const DEFAULT_MAX_THREAD_DEPTH = 1; const DEFAULT_LEASE_TTL_MS = 30_000; const DEFAULT_LEASE_HEARTBEAT_MS = 10_000; const DEFAULT_OWNER_DRAIN_TIMEOUT_MS = 45_000; +/** Keeps the admission fence alive across the deletion that follows the drain. */ +const OWNER_FENCE_GRACE_MS = 5 * 60_000; const DEFAULT_OWNER_DRAIN_POLL_MS = 100; +/** Matches the deletion drain batch so cancellation cannot burst Redis. */ +const DELETION_CANCEL_CONCURRENCY = 32; +/** Bounds retained control invocations; one entry per applied command. */ +const MAX_CONTROL_INVOCATIONS = 4_096; + +/** A cancellation target set resolved before the conversations are removed. */ +export interface SubagentCancellationPlan { + userId: string; + tenantId?: string; + conversationIds: string[]; + scopes: Array<{ scopeId: string; threadIds: string[] | null }>; + leases: IActiveSubagentThreadLease[]; +} +/** Three missed 10-second transport heartbeats retire a crashed owner. */ +const DEFAULT_TASK_ROUTING_TTL_MS = 30_000; const MAX_TRANSCRIPT_BYTES = 12 * 1024 * 1024; const TRANSCRIPT_SELECT = 'messageId parentMessageId text createdAt +subagentTranscript +subagentTask'; +const DURABLE_RESULT_SELECT = + 'messageId conversationId sender text createdAt updatedAt +subagentTask'; class SubagentThreadPublicError extends Error {} class SubagentThreadDeletedError extends SubagentThreadPublicError {} @@ -44,11 +75,13 @@ class SubagentThreadDeletedError extends SubagentThreadPublicError {} type SubagentThreadMethods = Pick< AllMethods, | 'acquireSubagentThreadLease' + | 'claimSubagentTaskResult' | 'countActiveSubagentThreadLeases' | 'deleteConvos' | 'deleteMessages' | 'getConvo' | 'getMessages' + | 'listActiveSubagentThreadLeases' | 'reserveSubagentThread' | 'releaseSubagentThreadLease' | 'renewSubagentThreadLease' @@ -88,6 +121,8 @@ interface TaskThreadLease { shared?: { token: string; lost: boolean; + /** Epoch ms this lease is durable until, advanced only by a confirmed renewal. */ + expiresAt: number; heartbeat?: ReturnType; heartbeatInFlight?: Promise; }; @@ -99,7 +134,13 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor leaseHeartbeatMs?: number; ownerDrainTimeoutMs?: number; ownerDrainPollMs?: number; + taskRoutingTtlMs?: number; isOwnerActive?: (userId: string) => Promise; + maxControlInvocations?: number; + ownerFenceGraceMs?: number; + fenceOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; + renewOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; + releaseOwnerAdmission?: (userId: string, token: string) => Promise; } function positiveInteger(value: number | undefined, fallback: number): number { @@ -142,6 +183,10 @@ function parseScope(scopeId: string): SubagentThreadScope { }; } +function serializeScope(scope: Omit): string { + return JSON.stringify({ version: SCOPE_VERSION, ...scope }); +} + function matchesTenant(actual: string | undefined, expected: string | undefined): boolean { return actual === expected; } @@ -284,21 +329,65 @@ function publicFailureDetail(error: unknown): string { : 'The child run could not be completed.'; } +/** Rebuilds the terminal claim a recovered durable result stands for. */ +function recoveredClaim( + message: IMessage, + claim: Extract, +): SubagentTaskClaim | undefined { + const status = message.subagentTask?.status; + const content = message.text ?? ''; + /** A durable child message keeps the untruncated output, so recovering one applies + * the same bounds a routed response would have. */ + if (status === 'completed') { + return boundedClaim({ status: 'completed', task: claim.task, result: content }); + } + if (status === 'error' || status === 'cancelled') { + return boundedClaim({ status, task: claim.task, error: content }); + } + return undefined; +} + +function drainKey(parentConversationId: string, taskId: string): string { + return `${parentConversationId}\u0000${taskId}`; +} + function safeErrorMessage(error: unknown): string { return `Subagent task failed: ${publicFailureDetail(error).slice(0, 2_000)}`; } -/** Persists view-only logical child threads with process-local controls and a shared execution fence. */ +/** Persists view-only logical child threads with owner-routed controls and a shared execution fence. */ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { readonly supportsThreadContinuation = true; private readonly activeThreads = new Map(); + private readonly controlInvocations = new Map< + string, + { scopeId: string; taskId: string; fingerprint: string; result: SubagentTaskControlResult } + >(); + private readonly parentPersistence = new Map>(); private readonly maxThreadDepth: number; private readonly leaseTtlMs: number; private readonly leaseHeartbeatMs: number; private readonly ownerDrainTimeoutMs: number; private readonly ownerDrainPollMs: number; + private readonly taskRoutingTtlMs: number; + private readonly maxControlInvocations: number; + private readonly ownerFenceGraceMs: number; private readonly isOwnerActive: (userId: string) => Promise; + private readonly fenceOwnerAdmission?: ( + userId: string, + token: string, + fencedUntil: Date, + ) => Promise; + + private readonly renewOwnerAdmission?: ( + userId: string, + token: string, + fencedUntil: Date, + ) => Promise; + + private readonly releaseOwnerAdmission?: (userId: string, token: string) => Promise; + private taskControlTransport?: SubagentTaskControlTransport; constructor( private readonly methods: SubagentThreadMethods, @@ -319,7 +408,37 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { DEFAULT_OWNER_DRAIN_TIMEOUT_MS, ); this.ownerDrainPollMs = positiveInteger(options.ownerDrainPollMs, DEFAULT_OWNER_DRAIN_POLL_MS); + this.taskRoutingTtlMs = positiveInteger(options.taskRoutingTtlMs, DEFAULT_TASK_ROUTING_TTL_MS); + this.maxControlInvocations = positiveInteger( + options.maxControlInvocations, + MAX_CONTROL_INVOCATIONS, + ); + this.ownerFenceGraceMs = positiveInteger(options.ownerFenceGraceMs, OWNER_FENCE_GRACE_MS); this.isOwnerActive = options.isOwnerActive ?? (async () => true); + this.fenceOwnerAdmission = options.fenceOwnerAdmission; + this.renewOwnerAdmission = options.renewOwnerAdmission; + this.releaseOwnerAdmission = options.releaseOwnerAdmission; + } + + /** Enables optional cross-replica lookup after the host's Redis service is ready. */ + async configureTaskControlTransport(transport: SubagentTaskControlTransport): Promise { + if (this.taskControlTransport != null) { + throw new Error('Subagent task control transport is already configured.'); + } + await transport.bind({ + claim: (scopeId, taskId) => super.claim(scopeId, taskId), + control: (scopeId, taskId, command, invocationId) => + this.controlInvocation(scopeId, taskId, command, invocationId), + list: (scopeId) => super.list(scopeId), + cancelScope: (scopeId, threadIds) => this.cancelForScope(scopeId, threadIds), + }); + this.taskControlTransport = transport; + } + + async destroyTaskControlTransport(): Promise { + const transport = this.taskControlTransport; + this.taskControlTransport = undefined; + await transport?.destroy(); } /** Gates child creation on the ordinary parent write without retaining request state. */ @@ -381,6 +500,15 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { if (runtime.signal.aborted) { throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); } + /** Publish the owner address before any provider work: a child running + * while unaddressable cannot be polled, controlled, or cancelled, and its + * side effects would already have happened by the time a heartbeat + * republished it. A failed registration fails the task closed instead. */ + await this.taskControlTransport?.registerTask( + request.scopeId, + runtime.taskId, + this.taskRoutingTtlMs, + ); await parentReady; const prepared = await this.prepareThread( request.scopeId, @@ -490,6 +618,285 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { return started; } + /** + * Claims locally when possible, otherwise asks the registered owning replica. + * + * A child's terminal result is durable in its own thread, so collection is recorded + * there against the polling invocation rather than kept alive in the owner's memory. + * The invocation that lost a response re-acquires its own result on the next poll; + * a different invocation is told the result was already collected. Owner-side + * retention stays a fast path, free to expire, instead of the only copy. + */ + async claimTask( + scopeId: string, + taskId: string, + invocationId?: string, + ): Promise { + const local = super.claim(scopeId, taskId); + const claim = + local.status !== 'not_found' + ? local + : ((await this.taskControlTransport?.claim(scopeId, taskId)) ?? local); + if (invocationId == null || claim.status === 'running') { + return claim; + } + if (claim.status === 'not_found') { + return this.claimDurableTaskResult(scopeId, taskId, invocationId); + } + const threadId = claim.task.threadId; + if (threadId == null || threadId === '') { + return claim; + } + /** The durable record decides who holds this one-shot result. The invocation that + * already consumed it re-acquires and is handed it again, a second invocation is + * told it was collected instead of being given a duplicate, and a task with no + * durable record to arbitrate keeps whatever the owner just answered. */ + const collected = await this.assignResultClaim( + parseScope(scopeId).userId, + threadId, + claim.task.taskId, + invocationId, + ); + if (collected.status === 'claimed') { + return { status: 'claimed', task: claim.task }; + } + if (collected.status === 'not_found') { + return claim; + } + return claim.status === 'claimed' ? (recoveredClaim(collected.message, claim) ?? claim) : claim; + } + + /** + * Recovers a terminal task after its owning process and Redis registration are gone. + * The task id locates only a candidate; durable child lineage re-establishes the + * trusted parent scope before the one-shot result is claimed. + */ + private async claimDurableTaskResult( + scopeId: string, + taskId: string, + invocationId: string, + ): Promise { + const scope = parseScope(scopeId); + let message: IMessage | undefined; + try { + [message] = await this.methods.getMessages( + { + user: scope.userId, + messageId: `${taskId}:assistant`, + 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + }, + DURABLE_RESULT_SELECT, + { limit: 1, sort: false }, + ); + } catch (error) { + logger.warn('[subagentThreads] Failed to locate a durable child result', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const threadId = message?.conversationId; + const status = message?.subagentTask?.status; + if ( + message == null || + !isNonEmptyString(threadId) || + !isNonEmptyString(message.sender) || + (status !== 'completed' && status !== 'error' && status !== 'cancelled') + ) { + return { status: 'not_found' }; + } + + let parent: IConversation | null; + let conversation: IConversation | null; + try { + [parent, conversation] = await Promise.all([ + this.methods.getConvo(scope.userId, scope.parentConversationId), + this.methods.getConvo(scope.userId, threadId), + ]); + } catch (error) { + logger.warn('[subagentThreads] Failed to verify durable child lineage', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const lineage = conversation?.subagentThread; + if ( + parent == null || + conversation == null || + lineage == null || + conversation.endpoint !== EModelEndpoint.agents || + lineage.parentConversationId !== scope.parentConversationId || + lineage.subagentType !== message.sender || + lineage.depth > this.maxThreadDepth || + !matchesTenant(parent.tenantId, scope.tenantId) || + !matchesTenant(conversation.tenantId, scope.tenantId) + ) { + return { status: 'not_found' }; + } + + const createdAt = message.createdAt?.getTime(); + const updatedAt = message.updatedAt?.getTime() ?? createdAt; + if (createdAt == null || updatedAt == null) { + return { status: 'not_found' }; + } + const task: SubagentTaskSnapshot = { + taskId, + threadId, + subagentType: lineage.subagentType, + status, + createdAt, + updatedAt, + resultAvailable: true, + resultClaimed: true, + pendingControls: 0, + ...(status === 'completed' ? {} : { error: message.text ?? '' }), + }; + const collected = await this.assignResultClaim(scope.userId, threadId, taskId, invocationId); + if (collected.status === 'not_found') { + return { status: 'not_found' }; + } + if (collected.status === 'claimed') { + return { status: 'claimed', task }; + } + return ( + recoveredClaim(collected.message, { status: 'claimed', task }) ?? { + status: 'not_found', + } + ); + } + + /** + * Assigns one durable terminal result to the invocation collecting it. A failed + * write is not an absent record: handing the result over without recording its + * claimant would let another invocation acquire the same one-shot output once the + * database recovers, so this reports the retryable path and leaves the result + * unclaimed for a later poll. + */ + private async assignResultClaim( + userId: string, + threadId: string, + taskId: string, + invocationId: string, + ): Promise { + try { + return await this.methods.claimSubagentTaskResult({ + userId, + conversationId: threadId, + taskId, + claimId: invocationId, + }); + } catch (error) { + logger.warn('[subagentThreads] Failed to record a collected child result', error); + throw new SubagentTaskOwnerUnavailableError(); + } + } + + /** + * Controls locally when possible, otherwise asks the registered owning replica. + * `invocationId` identifies one caller invocation: a routed retransmission of that + * invocation replays the owner's result, while a fresh invocation applies again even + * when its action and message are identical. + */ + async controlTask( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string = randomUUID(), + ): Promise { + const local = this.controlInvocation(scopeId, taskId, command, invocationId); + if (local.status !== 'not_found') { + return local; + } + return ( + (await this.taskControlTransport?.control(scopeId, taskId, command, invocationId)) ?? local + ); + } + + /** + * Applies one logical control exactly once for its owning task. Idempotency lives + * here rather than in the transport so a local and a routed caller of the same + * invocation agree, and it is keyed by task as well as invocation because provider + * tool-call ids repeat across runs and agents. + */ + controlInvocation( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): SubagentTaskControlResult { + const key = `${scopeId}\u0000${taskId}\u0000${invocationId}`; + const fingerprint = controlFingerprint(command); + const applied = this.controlInvocations.get(key); + if (applied != null) { + /** One invocation is one command; reusing its id for different content is a + * caller error rather than a retry, so it is refused instead of applied. */ + return applied.fingerprint === fingerprint + ? applied.result + : { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + } + if (this.get(scopeId, taskId) == null) { + /** Not this replica's task. Refusing here would keep the command from ever + * reaching its owner, so local load cannot veto a remote cancellation: the + * owner applies its own window to the routed request. */ + return this.control(scopeId, taskId, command); + } + if (!this.makeRoomForInvocation()) { + /** Every tracked invocation belongs to a task this store still holds. Applying + * this command without room to record it would let a caller retry apply it a + * second time, so it is refused before the child is touched at all. */ + logger.warn('[subagentThreads] Refused a control; live invocation records are full'); + return { + status: 'invalid', + message: 'Too many control invocations are in flight for this process; retry shortly.', + }; + } + const result = this.control(scopeId, taskId, command); + if (result.status === 'not_found') { + return result; + } + this.controlInvocations.set(key, { scopeId, taskId, fingerprint, result }); + return result; + } + + /** + * Frees invocation slots by dropping records whose task the store no longer holds: + * a settled task cannot be controlled again, so its record is worthless, while a + * live one is exactly what a caller retry needs to replay instead of applying its + * command twice. The sweep runs only when the window is full and clears every dead + * record at once, so it is amortized rather than repeated per control. + */ + private makeRoomForInvocation(): boolean { + if (this.controlInvocations.size < this.maxControlInvocations) { + return true; + } + for (const [key, invocation] of this.controlInvocations) { + if (this.get(invocation.scopeId, invocation.taskId) == null) { + this.controlInvocations.delete(key); + } + } + return this.controlInvocations.size < this.maxControlInvocations; + } + + /** Returns this process's tasks plus tasks reported by registered remote owners. */ + async listTasks(scopeId: string): Promise { + const local = super.list(scopeId); + const remote = (await this.taskControlTransport?.list(scopeId)) ?? []; + const byId = new Map(local.map((task) => [task.taskId, task])); + for (const task of remote) { + byId.set(task.taskId, task); + } + /** The remote aggregation and each owner's reply carry their own bound, but this + * merge is what the poll tool reads: without a cap here the list the model sees is + * that bound plus however many children this replica happens to own. */ + return boundedTaskList([...byId.values()]); + } + + /** Fast capability probe used while deciding whether a later turn needs the poll tool. */ + async hasTasks(scopeId: string): Promise { + if (super.list(scopeId).length > 0) { + return true; + } + return (await this.taskControlTransport?.hasTasks(scopeId)) ?? false; + } + override control( scopeId: string, taskId: string, @@ -544,6 +951,148 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ); } + /** + * Resolves every cancellation target while the conversations still exist. The plan is + * replayed after deletion, when those rows can no longer be read back, so the second + * pass only has to reach registered owners through Redis. + */ + async planCancellationForConversations( + userId: string, + conversationIds: Iterable, + tenantId?: string, + ): Promise { + const targetIds = [...new Set(conversationIds)]; + const plan: SubagentCancellationPlan = { + userId, + ...(tenantId == null ? {} : { tenantId }), + conversationIds: targetIds, + scopes: [], + leases: [], + }; + if (targetIds.length === 0 || this.taskControlTransport == null) { + return plan; + } + const targets = new Set(targetIds); + const scopeIdFor = (parentConversationId: string): string => + serializeScope({ + userId, + parentConversationId, + ...(tenantId ? { tenantId } : {}), + }); + /** Deleting a conversation takes its whole scope; a deleted child only cancels its + * own thread inside a parent scope that survives. */ + const conversations = await Promise.all( + targetIds.map((conversationId) => this.methods.getConvo(userId, conversationId)), + ); + const threadTargetsByParent = new Map>(); + for (const [index, conversation] of conversations.entries()) { + const parentConversationId = conversation?.subagentThread?.parentConversationId; + if ( + parentConversationId == null || + targets.has(parentConversationId) || + !matchesTenant(conversation?.tenantId, tenantId) + ) { + continue; + } + const threadIds = threadTargetsByParent.get(parentConversationId) ?? new Set(); + threadIds.add(targetIds[index]); + threadTargetsByParent.set(parentConversationId, threadIds); + } + plan.scopes = [ + ...targetIds.map((parentConversationId) => ({ + scopeId: scopeIdFor(parentConversationId), + threadIds: null, + })), + ...[...threadTargetsByParent].map(([parentConversationId, threadIds]) => ({ + scopeId: scopeIdFor(parentConversationId), + threadIds: [...threadIds], + })), + ]; + /** Captured now so descendants removed by the cascade stay reachable afterwards. */ + plan.leases = await this.methods.listActiveSubagentThreadLeases({ + user: userId, + now: new Date(), + ...(tenantId == null ? {} : { tenantId }), + }); + return plan; + } + + /** + * Cancels local children and replays a plan against registered remote owners. + * `removedConversationIds` extends it with the cascade a deletion reported, matched + * against leases captured before those rows were removed. + */ + async cancelPlan( + plan: SubagentCancellationPlan, + removedConversationIds: Iterable = [], + ): Promise { + const { userId, tenantId } = plan; + const planned = new Set(plan.conversationIds); + const removed = new Set(removedConversationIds); + /** A cascade can remove descendants the plan never named — a grandchild lives in + * its own parent's scope, not the deleted root's — so every removed conversation + * is cancelled as a scope of its own. */ + const targets = [...new Set([...planned, ...removed])]; + let cancelled = this.cancelForConversations(userId, targets, tenantId); + const transport = this.taskControlTransport; + if (transport == null) { + return cancelled; + } + const cancelSlot = createConcurrencyLimiter(DELETION_CANCEL_CONCURRENCY); + const cascadeScopes = [...removed] + .filter((conversationId) => !planned.has(conversationId)) + .map((parentConversationId) => ({ + scopeId: serializeScope({ + userId, + parentConversationId, + ...(tenantId ? { tenantId } : {}), + }), + threadIds: null, + })); + const scopeCancellations = [...plan.scopes, ...cascadeScopes].map((scope) => + cancelSlot(() => transport.cancelScope(scope.scopeId, scope.threadIds)), + ); + const leaseCancellations = plan.leases + .filter( + (lease) => removed.has(lease.parentConversationId) || removed.has(lease.conversationId), + ) + .map((lease) => + cancelSlot(() => + this.controlTask( + serializeScope({ + userId, + parentConversationId: lease.parentConversationId, + ...(tenantId ? { tenantId } : {}), + }), + lease.taskId, + { action: 'cancel' }, + ), + ), + ); + for (const count of await Promise.all(scopeCancellations)) { + cancelled += count; + } + for (const result of await Promise.all(leaseCancellations)) { + if (result.status === 'cancelled') { + cancelled += 1; + } + } + return cancelled; + } + + /** Cancels this process's live children for one scope, optionally narrowed to threads. */ + private cancelForScope(scopeId: string, threadIds: string[] | null): number { + const scope = parseScope(scopeId); + const targets = threadIds == null ? null : new Set(threadIds); + return this.cancelMatchingThreads( + (candidate, threadId) => + candidate.userId === scope.userId && + candidate.parentConversationId === scope.parentConversationId && + matchesTenant(candidate.tenantId, scope.tenantId) && + (targets == null || targets.has(threadId)), + ); + } + /** Cancels every active child owned by a user before a delete-all operation. */ cancelForOwner(userId: string, tenantId?: string): number { return this.cancelMatchingThreads( @@ -551,26 +1100,202 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { ); } - /** Cancels local work and waits for every replica's durable lease to drain. */ + /** + * Deletes an owner's conversations behind a durable admission fence. Draining alone + * cannot close the race: a child admitted on another replica after the drain read + * its leases would begin provider work against a parent that is about to disappear. + * Fencing first inverts that — the fence is written before any lease is read, and a + * child validates the fence after its own lease is written, so one of the two always + * observes the other. The fence expires by itself, so a process lost mid-deletion + * cannot leave the account unable to run subagents. + */ + async withOwnerDeletionFence( + userId: string, + tenantId: string | undefined, + deletion: () => Promise, + ): Promise { + const fenceWindowMs = this.ownerDrainTimeoutMs + this.ownerFenceGraceMs; + const token = randomUUID(); + /** Only a confirmed write moves this, so a run of failed renewals leaves it in the + * past and the deletion can tell that its fence is no longer guaranteed. */ + let fencedUntil = Date.now() + fenceWindowMs; + let fenceLapsed = false; + await this.fenceOwnerAdmission?.(userId, token, new Date(fencedUntil)); + /** A very large account, or a stalled database, can outlast one fence window, and + * a fence that expires mid-deletion lets another replica admit a child against + * conversations being deleted. It is renewed for as long as the work runs. */ + let releasing = false; + let inFlight: Promise | undefined; + const renewal = setInterval( + () => { + if (inFlight != null) { + return; + } + inFlight = (async () => { + const deadline = fencedUntil; + const renewedUntil = Date.now() + fenceWindowMs; + const held = await this.renewOwnerAdmission?.(userId, token, new Date(renewedUntil)); + if (held === false) { + /** The durable entry was absent, so admission may already have opened even + * when the local deadline has not passed. Reacquire for containment, but + * retain the lapse so the enclosing deletion re-drains before success. */ + fenceLapsed = true; + if (releasing) { + return; + } + /** The entry is gone — expired, or pruned by another deletion — so this + * deletion takes its fence again rather than running on unfenced. */ + await this.fenceOwnerAdmission?.(userId, token, new Date(renewedUntil)); + } + if (Date.now() >= deadline) { + /** The write only landed after the deadline it was meant to extend, so + * admission stood open in between and a child could have taken a lease the + * drain had already read past. A fence cannot be restored backwards over + * that gap, so the lapse is recorded rather than papered over. */ + fenceLapsed = true; + return; + } + fencedUntil = renewedUntil; + })() + .catch((error) => { + logger.warn('[subagentThreads] Failed to hold the owner admission fence', error); + }) + .finally(() => { + inFlight = undefined; + }); + }, + Math.max(1, Math.floor(fenceWindowMs / 3)), + ); + renewal.unref?.(); + const stopRenewal = async (): Promise => { + clearInterval(renewal); + await inFlight; + }; + const fenceHeld = (): boolean => + this.fenceOwnerAdmission == null || (!fenceLapsed && Date.now() < fencedUntil); + try { + await this.cancelAndDrainForOwner(userId, tenantId); + /** The drain can outlast the fence window when the database is unreachable, and + * renewals that keep failing leave the account open to admitting a child against + * conversations about to disappear. Nothing has been removed yet, so this fails + * closed and the caller retries once the fence can be held again. */ + if (!fenceHeld()) { + throw new Error('The subagent admission fence expired before this deletion began.'); + } + const deleted = await deletion(); + /** Settle a renewal already in flight before deciding whether deletion crossed a + * gap. Otherwise a late write can report the lapse only after this check and the + * finally block would release the fence without re-draining. */ + await stopRenewal(); + if (!fenceHeld()) { + /** The rows are gone, but the gap can leave a child another replica admitted + * while the fence was down. Re-take the fence and drain that work before this + * operation may report success. */ + logger.error( + '[subagentThreads] Owner deletion outlived its admission fence; draining children admitted in the gap', + ); + const recoveryUntil = Date.now() + fenceWindowMs; + const reheld = await this.renewOwnerAdmission?.(userId, token, new Date(recoveryUntil)); + if (reheld !== true) { + await this.fenceOwnerAdmission?.(userId, token, new Date(recoveryUntil)); + } + if (Date.now() >= recoveryUntil) { + throw new Error('The subagent admission fence expired while it was being restored.'); + } + fencedUntil = recoveryUntil; + fenceLapsed = false; + await this.cancelAndDrainForOwner(userId, tenantId); + if (!fenceHeld()) { + throw new Error('The subagent admission fence expired while recovering this deletion.'); + } + } + return deleted; + } finally { + releasing = true; + await stopRenewal(); + /** `clearInterval` stops only future passes. A renewal still waiting on the + * database would otherwise find its fence released, read that as expiry, and + * write a fresh one that nothing is left to lift. */ + /** Only this deletion's own fence is lifted: an overlapping deletion that took a + * later one keeps admission closed until it finishes. */ + await this.releaseOwnerAdmission?.(userId, token).catch((error) => { + logger.warn('[subagentThreads] Failed to release the owner admission fence', error); + }); + } + } + + /** + * Cancels local work and waits for every replica's durable lease to drain. Each task + * is cancelled under one invocation held for the whole drain and only while its + * owner has not answered: a fresh invocation per poll would retain a replay entry on + * the owner for every pass, and a task already reported cancelled needs no second + * command, only its lease to disappear. + */ async cancelAndDrainForOwner(userId: string, tenantId?: string): Promise { this.cancelForOwner(userId, tenantId); const deadline = Date.now() + this.ownerDrainTimeoutMs; + const invocations = new Map(); + const answered = new Set(); while (true) { - const active = await this.methods.countActiveSubagentThreadLeases({ + const activeLeases = await this.methods.listActiveSubagentThreadLeases({ user: userId, now: new Date(), ...(tenantId == null ? {} : { tenantId }), }); - if (active === 0) { + if (activeLeases.length === 0) { return; } if (Date.now() >= deadline) { throw new Error('Timed out draining detached subagent tasks for account deletion.'); } + const unanswered = activeLeases.filter( + ({ parentConversationId, taskId }) => !answered.has(drainKey(parentConversationId, taskId)), + ); + for (let index = 0; index < unanswered.length; index += DELETION_CANCEL_CONCURRENCY) { + await Promise.all( + unanswered + .slice(index, index + DELETION_CANCEL_CONCURRENCY) + .map(({ parentConversationId, taskId }) => + this.cancelDrainedTask( + { userId, parentConversationId, taskId, tenantId }, + invocations, + answered, + ), + ), + ); + } await new Promise((resolve) => setTimeout(resolve, this.ownerDrainPollMs)); } } + /** Sends one drained task's cancellation, retrying only unconfirmed deliveries. */ + private async cancelDrainedTask( + target: { userId: string; parentConversationId: string; taskId: string; tenantId?: string }, + invocations: Map, + answered: Set, + ): Promise { + const { userId, parentConversationId, taskId, tenantId } = target; + const key = drainKey(parentConversationId, taskId); + const invocationId = invocations.get(key) ?? randomUUID(); + invocations.set(key, invocationId); + const scopeId = serializeScope({ + userId, + parentConversationId, + ...(tenantId == null ? {} : { tenantId }), + }); + try { + const result = await this.controlTask(scopeId, taskId, { action: 'cancel' }, invocationId); + /** Only the owner confirming the task is stopped ends the commands for it. A + * `not_found` means its registration is missing while its lease is still live — + * an unconfirmed delivery, retried once the owner republishes itself. */ + if (result.status === 'cancelled' || result.status === 'not_running') { + answered.add(key); + } + } catch (error) { + logger.warn('[subagentThreads] Retrying an unconfirmed child cancellation', error); + } + } + private startSharedLeaseHeartbeat( scopeId: string, scope: SubagentThreadScope, @@ -645,19 +1370,32 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { return false; } try { + const deadline = shared.expiresAt; const now = new Date(); + const renewedUntil = now.getTime() + this.leaseTtlMs; const renewed = await this.methods.renewSubagentThreadLease({ user: scope.userId, conversationId: threadId, token: shared.token, now, - expiresAt: new Date(now.getTime() + this.leaseTtlMs), + expiresAt: new Date(renewedUntil), ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), }); if (!renewed) { shared.lost = true; + return false; } - return renewed; + if (Date.now() >= deadline) { + /** The renewal filter compares against the `now` captured before the call, so a + * write that only lands after this lease had expired still succeeds and moves + * the row forward. An owner drain reading active leases in that gap saw this + * thread as free, so the executor stops rather than run past a deletion that + * may already have stepped over it. */ + shared.lost = true; + return false; + } + shared.expiresAt = renewedUntil; + return true; } catch (error) { shared.lost = true; logger.warn('[subagentThreads] Lost the shared child-thread lease', error); @@ -779,7 +1517,11 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { 'This child thread is already being continued by another run.', ); } - lease.shared = { token: sharedToken, lost: false }; + lease.shared = { + token: sharedToken, + lost: false, + expiresAt: now.getTime() + this.leaseTtlMs, + }; this.startSharedLeaseHeartbeat(scopeId, scope, threadId, lease); /** Account deletion can fence the owner after the optimistic probe but before * this lease exists. Once the lease is visible, revalidate so deletion either @@ -1226,6 +1968,22 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { } } +const REQUIRED_THREAD_METHODS = [ + 'acquireSubagentThreadLease', + 'claimSubagentTaskResult', + 'countActiveSubagentThreadLeases', + 'deleteConvos', + 'deleteMessages', + 'getConvo', + 'getMessages', + 'listActiveSubagentThreadLeases', + 'releaseSubagentThreadLease', + 'renewSubagentThreadLease', + 'reserveSubagentThread', + 'saveConvo', + 'saveMessage', +] as const; + export function createSubagentThreadTaskStore( methods: Pick< ConversationMethods, @@ -1233,14 +1991,27 @@ export function createSubagentThreadTaskStore( | 'countActiveSubagentThreadLeases' | 'deleteConvos' | 'getConvo' + | 'listActiveSubagentThreadLeases' | 'releaseSubagentThreadLease' | 'reserveSubagentThread' | 'renewSubagentThreadLease' | 'saveConvo' > & - Pick, + Pick< + MessageMethods, + 'claimSubagentTaskResult' | 'deleteMessages' | 'getMessages' | 'saveMessage' + >, options?: SubagentThreadTaskStoreOptions, ): SubagentThreadTaskStore { + /** The host wires this from JavaScript, where the parameter type checks nothing. A + * method missing there would otherwise surface as a routed failure at claim time, + * long after startup, so the omission is caught here instead. */ + const missing = REQUIRED_THREAD_METHODS.filter( + (name) => typeof (methods as Record)[name] !== 'function', + ); + if (missing.length > 0) { + throw new Error(`Subagent thread task store is missing methods: ${missing.join(', ')}`); + } return new SubagentThreadTaskStore(methods, options); } @@ -1250,6 +2021,6 @@ export function buildSubagentThreadTaskConfig( ): SubagentTaskConfig { return { store, - scopeId: JSON.stringify({ version: SCOPE_VERSION, ...scope }), + scopeId: serializeScope(scope), }; } diff --git a/packages/api/src/cache/redisUtils.spec.ts b/packages/api/src/cache/redisUtils.spec.ts new file mode 100644 index 0000000000..08abbd2753 --- /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 de37c8ba5c..7c824d4ebe 100644 --- a/packages/api/src/cache/redisUtils.ts +++ b/packages/api/src/cache/redisUtils.ts @@ -1,7 +1,49 @@ -import type { RedisClientType, RedisClusterType } from '@redis/client'; import { logger } from '@librechat/data-schemas'; +import type { ClusterOptions, RedisOptions, Cluster, Redis } from 'ioredis'; +import type { RedisClientType, RedisClusterType } from '@redis/client'; import { cacheConfig } from './cacheConfig'; +/** + * Duplicates an ioredis connection with option overrides. `Cluster.duplicate` reads its + * first argument as an optional startup-node list and its second as the overrides, + * unlike `Redis.duplicate`, so options passed positionally to a cluster are silently + * dropped and the duplicate quietly inherits the original's behaviour. + */ +export function duplicateIoRedisClient( + client: Redis | Cluster, + options: RedisOptions & ClusterOptions = {}, +): Redis | Cluster { + if (client.isCluster) { + const duplicate = (client as Cluster).duplicate([], options); + if (options.enableOfflineQueue !== false) { + return duplicate; + } + let clusterHasBeenReady = duplicate.status === 'ready'; + duplicate.once('ready', () => { + clusterHasBeenReady = true; + }); + /** ioredis deliberately forces `enableOfflineQueue: true` on every Cluster node + * after applying `redisOptions`. It needs that queue while a new node discovers + * topology, so changing it at `+node` prevents the cluster from ever becoming + * ready. Initial nodes switch once connected; nodes discovered after the cluster + * was usable fail fast immediately, including during a slot-owner replacement. */ + const disableNodeOfflineQueue = (node: Redis): void => { + const disable = (): void => { + node.options.enableOfflineQueue = false; + }; + if (node.status === 'ready' || clusterHasBeenReady) { + disable(); + } else { + node.once('ready', disable); + } + }; + duplicate.on('+node', disableNodeOfflineQueue); + duplicate.nodes('all').forEach(disableNodeOfflineQueue); + return duplicate; + } + return (client as Redis).duplicate(options); +} + /** * Efficiently deletes multiple Redis keys with support for both cluster and single-node modes. * diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index 1c69d79d6d..b87de135c7 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 79b9d84ad9..0703e5129c 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 772938d9fa..d9f16f327b 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 097fe99810..add99d0fc6 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -28,6 +28,7 @@ let bulkSaveMessages: ReturnType['bulkSaveMessages' let updateMessageText: ReturnType['updateMessageText']; let deleteMessagesSince: ReturnType['deleteMessagesSince']; let recordMessage: ReturnType['recordMessage']; +let claimSubagentTaskResult: ReturnType['claimSubagentTaskResult']; beforeAll(async () => { mongoServer = await MongoMemoryServer.create(); @@ -47,6 +48,7 @@ beforeAll(async () => { updateMessageText = methods.updateMessageText; deleteMessagesSince = methods.deleteMessagesSince; recordMessage = methods.recordMessage; + claimSubagentTaskResult = methods.claimSubagentTaskResult; await mongoose.connect(mongoUri); }); @@ -1549,4 +1551,81 @@ describe('Message Operations', () => { expect(doc?.tenantId).toBeUndefined(); }); }); + describe('claimSubagentTaskResult', () => { + const terminalResult = async (taskId: string, conversationId: string, status: string) => + saveMessage({ userId: 'user123' }, { + messageId: `${taskId}:assistant`, + conversationId, + text: 'child result', + subagentTask: { attemptKey: `${taskId}:attempt`, status }, + } as Partial); + + it('hands one terminal result to a single polling invocation', async () => { + const taskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(taskId, conversationId, 'completed'); + + const first = await claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + claimId: 'poll-1', + }); + expect(first.status).toBe('acquired'); + expect(first.status === 'acquired' && first.message.text).toBe('child result'); + + /** The same invocation retrying recovers the result it never received. */ + const retried = await claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId, + claimId: 'poll-1', + }); + expect(retried.status).toBe('acquired'); + + /** Another invocation is told it was collected instead of handed a copy. */ + await expect( + claimSubagentTaskResult({ userId: 'user123', conversationId, taskId, claimId: 'poll-2' }), + ).resolves.toEqual({ status: 'claimed' }); + }); + + it('reports a result that is missing or still running as not found', async () => { + const runningTaskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(runningTaskId, conversationId, 'running'); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: runningTaskId, + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + + await expect( + claimSubagentTaskResult({ + userId: 'user123', + conversationId, + taskId: uuidv4(), + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + }); + + it('never hands one owner’s result to another user', async () => { + const taskId = uuidv4(); + const conversationId = uuidv4(); + await terminalResult(taskId, conversationId, 'completed'); + + await expect( + claimSubagentTaskResult({ + userId: 'other-user', + conversationId, + taskId, + claimId: 'poll-1', + }), + ).resolves.toEqual({ status: 'not_found' }); + }); + }); }); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index a864acc576..6f427bb820 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -48,6 +48,11 @@ interface MessageQueryOptions { sort?: Record | false; } +export type SubagentTaskResultClaim = + | { status: 'not_found' } + | { status: 'claimed' } + | { status: 'acquired'; message: IMessage }; + export interface MessageMethods { saveMessage( ctx: { userId: string; isTemporary?: boolean; interfaceConfig?: AppConfig['interfaceConfig'] }, @@ -82,6 +87,12 @@ export interface MessageMethods { message: Partial & { newMessageId?: string }, metadata?: { context?: string }, ): Promise>; + claimSubagentTaskResult(params: { + userId: string; + conversationId: string; + taskId: string; + claimId: string; + }): Promise; deleteMessagesSince( userId: string, params: { messageId: string; conversationId: string }, @@ -518,6 +529,62 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } } + /** + * Assigns one durable terminal child result to the polling invocation that collects + * it. The same invocation may re-acquire, so a poll whose response was lost recovers + * the result it never received; a different invocation is told it was already + * collected rather than handed a second copy. + */ + async function claimSubagentTaskResult({ + userId, + conversationId, + taskId, + claimId, + }: { + userId: string; + conversationId: string; + taskId: string; + claimId: string; + }): Promise { + if ( + taskId.length === 0 || + taskId.length > 256 || + conversationId.length === 0 || + conversationId.length > 256 || + claimId.length === 0 || + claimId.length > 128 + ) { + throw new TypeError('Invalid subagent task result claim'); + } + const Message = mongoose.models.Message as Model; + const filter = { + user: userId, + conversationId, + messageId: `${taskId}:assistant`, + 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + }; + const acquired = await Message.findOneAndUpdate( + { + ...filter, + $or: [ + { 'subagentTask.resultClaim': { $exists: false } }, + { 'subagentTask.resultClaim.claimId': claimId }, + ], + }, + { $set: { 'subagentTask.resultClaim': { claimId, claimedAt: new Date() } } }, + { + new: true, + timestamps: false, + projection: { messageId: 1, conversationId: 1, text: 1, subagentTask: 1 }, + }, + ).lean(); + if (acquired != null) { + return { status: 'acquired', message: acquired }; + } + const existing = await Message.exists(filter); + return existing == null ? { status: 'not_found' } : { status: 'claimed' }; + } + /** * Deletes messages in a conversation since a specific message. */ @@ -655,6 +722,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateMessageText, updateToolCallResult, updateMessage, + claimSubagentTaskResult, deleteMessagesSince, getMessages, getMessage, diff --git a/packages/data-schemas/src/methods/user.methods.spec.ts b/packages/data-schemas/src/methods/user.methods.spec.ts index 9540026232..d7be36e814 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 2d324bfcdc..2840c458b3 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 890f1fe2cd..8282cbf059 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -144,6 +144,14 @@ const messageSchema: Schema = new Schema( enum: ['running', 'completed', 'error', 'cancelled'], required: true, }, + resultClaim: { + type: { + claimId: { type: String, required: true }, + claimedAt: { type: Date, required: true }, + }, + _id: false, + default: undefined, + }, }, _id: false, select: false, diff --git a/packages/data-schemas/src/schema/user.ts b/packages/data-schemas/src/schema/user.ts index 79a9a64195..8d13c1ec79 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 5fa5fbf0a9..387720cf9f 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 9b1d4492f4..541a5f1d1f 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -53,6 +53,11 @@ export interface IMessage extends Document { attemptKey: string; requestFingerprint?: string; status: 'running' | 'completed' | 'error' | 'cancelled'; + /** Records which polling invocation collected this terminal result. */ + resultClaim?: { + claimId: string; + claimedAt: Date; + }; }; contextMeta?: { calibrationRatio?: number; diff --git a/packages/data-schemas/src/types/user.ts b/packages/data-schemas/src/types/user.ts index 169209cf4a..8370fb63b6 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;