diff --git a/CONTEXT.md b/CONTEXT.md index 59e4f83bb3..80f2c15602 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -1,4 +1,5 @@ # Domain language - **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. - **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/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index 34cfc04200..b3b900e5c7 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -1055,6 +1055,15 @@ describe('BaseClient', () => { anotherExistingField: 'anotherValue', temperature: 0.7, modelLabel: 'GPT-3.5', + subagentThread: { + rootConversationId: 'root-conversation', + parentConversationId: 'parent-conversation', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool-call', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, }; getConvo.mockResolvedValue(existingConvo); @@ -1089,6 +1098,7 @@ describe('BaseClient', () => { // Only check that someExistingField is in unsetFields expect(saveOptions.unsetFields).toHaveProperty('someExistingField', 1); + expect(saveOptions.unsetFields).not.toHaveProperty('subagentThread'); // Mock saveConvo to return the expected fields saveConvo.mockImplementation((req, fields) => { diff --git a/api/package.json b/api/package.json index 1717ed810a..4ffb1a5237 100644 --- a/api/package.json +++ b/api/package.json @@ -46,7 +46,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.3", + "@librechat/agents": "^3.6.6", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index b8fecb4f8e..9821f47db7 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -25,6 +25,7 @@ const { verifyEmail, resendVerificationEmail } = require('~/server/services/Auth const { getMCPManager, getFlowStateManager, getMCPServersRegistry } = require('~/config'); const { invalidateCachedTools } = require('~/server/services/Config/getCachedTools'); const { processDeleteRequest } = require('~/server/services/Files/process'); +const subagentThreadTaskStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); const { drainAgentTriggerDeliveriesForUser, prepareAgentTriggerUserPurge, @@ -393,6 +394,7 @@ const deleteUserController = async (req, res) => { await prepareAgentTriggerUserPurge(user.id, triggerDeletionFence, user.tenantId); } await drainAgentTriggerDeliveriesForUser(user.id); + await subagentThreadTaskStore.cancelAndDrainForOwner(user.id, user.tenantId); const activeAgentRuns = await GenerationJobManager.getCleanupBlockingJobIdsForUser( user.id, user.tenantId, diff --git a/api/server/controllers/UserController.spec.js b/api/server/controllers/UserController.spec.js index b21fe904e3..c9a3f6eb44 100644 --- a/api/server/controllers/UserController.spec.js +++ b/api/server/controllers/UserController.spec.js @@ -7,6 +7,7 @@ const mockDrainAgentTriggerDeliveriesForUser = jest.fn().mockResolvedValue(undef const mockPrepareAgentTriggerUserPurge = jest.fn().mockResolvedValue(undefined); const mockCancelAgentTriggerUserPurge = jest.fn().mockResolvedValue(true); const mockPurgeAgentTriggerDeliveriesForUser = jest.fn().mockResolvedValue(undefined); +const mockCancelAndDrainSubagentThreads = jest.fn().mockResolvedValue(undefined); jest.mock('@librechat/data-schemas', () => { const actual = jest.requireActual('@librechat/data-schemas'); @@ -95,6 +96,10 @@ jest.mock('~/server/services/Agents/triggers', () => ({ purgeAgentTriggerDeliveriesForUser: (...args) => mockPurgeAgentTriggerDeliveriesForUser(...args), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({ + cancelAndDrainForOwner: (...args) => mockCancelAndDrainSubagentThreads(...args), +})); + jest.mock('~/server/services/Files/process', () => ({ processDeleteRequest: jest.fn().mockResolvedValue({ deletedFileIds: [], failedFileIds: [] }), })); @@ -355,6 +360,7 @@ describe('deleteUserController', () => { undefined, ); expect(mockDrainAgentTriggerDeliveriesForUser).toHaveBeenCalledWith(userId.toString()); + expect(mockCancelAndDrainSubagentThreads).toHaveBeenCalledWith(userId.toString(), undefined); expect(beginAgentTriggerUserDeletion.mock.invocationCallOrder[0]).toBeLessThan( mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0], ); @@ -362,6 +368,9 @@ describe('deleteUserController', () => { mockDrainAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0], ); expect(mockDrainAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0]).toBeLessThan( + mockCancelAndDrainSubagentThreads.mock.invocationCallOrder[0], + ); + expect(mockCancelAndDrainSubagentThreads.mock.invocationCallOrder[0]).toBeLessThan( deleteMessages.mock.invocationCallOrder[0], ); expect(deleteMessages.mock.invocationCallOrder[0]).toBeLessThan( @@ -421,6 +430,29 @@ describe('deleteUserController', () => { expect(deleteUserById).not.toHaveBeenCalled(); }); + it('fails closed before data cleanup when detached subagents do not drain', async () => { + const userId = new mongoose.Types.ObjectId(); + mockCancelAndDrainSubagentThreads.mockRejectedValueOnce(new Error('child drain timed out')); + const req = { + user: { + id: userId.toString(), + _id: userId, + email: 'active-child@test.com', + tenantId: 'tenant-1', + }, + }; + + await deleteUserController(req, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(500); + expect(deleteMessages).not.toHaveBeenCalled(); + expect(cancelAgentTriggerUserDeletion).toHaveBeenCalledWith( + userId.toString(), + expect.any(Date), + ); + expect(deleteUserById).not.toHaveBeenCalled(); + }); + it('should remove the user from all groups via $pullAll', async () => { const userId = new mongoose.Types.ObjectId(); const userIdStr = userId.toString(); diff --git a/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js b/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js index 545283ce26..d843005c85 100644 --- a/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js +++ b/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js @@ -97,6 +97,10 @@ jest.mock('~/server/services/Agents/triggers', () => ({ purgeAgentTriggerDeliveriesForUser: jest.fn(), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({ + cancelAndDrainForOwner: jest.fn(), +})); + jest.mock('~/server/services/Config', () => ({ getAppConfig: (...args) => mockGetAppConfig(...args), })); diff --git a/api/server/controllers/__tests__/deleteUser.spec.js b/api/server/controllers/__tests__/deleteUser.spec.js index 9af915b772..f6777172a6 100644 --- a/api/server/controllers/__tests__/deleteUser.spec.js +++ b/api/server/controllers/__tests__/deleteUser.spec.js @@ -27,6 +27,7 @@ const mockCancelAgentTriggerUserPurge = jest.fn(); const mockPurgeAgentTriggerDeliveriesForUser = jest.fn(); const mockBeginAgentTriggerUserDeletion = jest.fn(); const mockCancelAgentTriggerUserDeletion = jest.fn(); +const mockCancelAndDrainSubagentThreads = jest.fn(); jest.mock('@librechat/data-schemas', () => ({ logger: { error: jest.fn(), info: jest.fn() }, @@ -122,6 +123,10 @@ jest.mock('~/server/services/Agents/triggers', () => ({ purgeAgentTriggerDeliveriesForUser: (...args) => mockPurgeAgentTriggerDeliveriesForUser(...args), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({ + cancelAndDrainForOwner: (...args) => mockCancelAndDrainSubagentThreads(...args), +})); + jest.mock('~/server/services/Config', () => ({ getAppConfig: jest.fn(), })); @@ -165,6 +170,7 @@ function stubDeletionMocks() { mockPurgeAgentTriggerDeliveriesForUser.mockResolvedValue(); mockBeginAgentTriggerUserDeletion.mockResolvedValue('acquired'); mockCancelAgentTriggerUserDeletion.mockResolvedValue(true); + mockCancelAndDrainSubagentThreads.mockResolvedValue(); } beforeEach(() => { diff --git a/api/server/controllers/__tests__/maybeUninstallOAuthMCP.spec.js b/api/server/controllers/__tests__/maybeUninstallOAuthMCP.spec.js index 30d687bf64..5e6e1d5444 100644 --- a/api/server/controllers/__tests__/maybeUninstallOAuthMCP.spec.js +++ b/api/server/controllers/__tests__/maybeUninstallOAuthMCP.spec.js @@ -108,6 +108,10 @@ jest.mock('~/server/services/Agents/triggers', () => ({ purgeAgentTriggerDeliveriesForUser: jest.fn(), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({ + cancelAndDrainForOwner: jest.fn(), +})); + jest.mock('~/server/services/Config', () => ({ getAppConfig: jest.fn(), })); diff --git a/api/server/controllers/agents/__tests__/client.subagentUsage.spec.js b/api/server/controllers/agents/__tests__/client.subagentUsage.spec.js new file mode 100644 index 0000000000..21c15e809d --- /dev/null +++ b/api/server/controllers/agents/__tests__/client.subagentUsage.spec.js @@ -0,0 +1,73 @@ +const mockGetMultiplier = jest.fn(() => 1); +const mockGetCacheMultiplier = jest.fn(() => 1); + +jest.mock('~/models', () => ({ + getMultiplier: (...args) => mockGetMultiplier(...args), + getCacheMultiplier: (...args) => mockGetCacheMultiplier(...args), +})); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { debug: jest.fn(), error: jest.fn(), warn: jest.fn(), info: jest.fn() }, +})); + +const AgentClient = require('../client'); + +describe('AgentClient#buildSubagentUsageEmitter', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('uses a lifecycle-safe snapshot after the parent client is disposed', async () => { + const write = jest.fn(); + const usageEmitSink = []; + const pendingSubagentEmits = []; + const endpointTokenConfig = { input: 2, output: 3 }; + const self = { + options: { + res: { write }, + req: { user: { id: 'user-1' } }, + endpointTokenConfig, + endpointTokenConfigByAgentId: new Map([['child-agent', endpointTokenConfig]]), + }, + responseMessageId: 'response-1', + jobCreatedAt: 1234, + usageEmitSink, + pendingSubagentEmits, + subagentUsageSeq: 4, + }; + const emit = AgentClient.prototype.buildSubagentUsageEmitter.call(self, { + interfaceConfig: { contextCost: true }, + }); + + /** Mirror the fields cleared by disposeClient before the detached child + * finishes; the callback must not read any of them. */ + self.options = null; + self.responseMessageId = null; + self.jobCreatedAt = null; + self.usageEmitSink = null; + self.pendingSubagentEmits = null; + + const usage = { + input_tokens: 10, + output_tokens: 5, + total_tokens: 15, + model: 'child-model', + provider: 'custom', + agentId: 'child-agent', + }; + await emit(usage); + + expect(usageEmitSink).toEqual([ + expect.objectContaining({ + runId: 'response-1:1234', + seq: 5, + usage_type: 'subagent', + cost: expect.any(Number), + }), + ]); + expect(usage.cost).toBe(usageEmitSink[0].cost); + expect(write).toHaveBeenCalledTimes(1); + expect(pendingSubagentEmits).toHaveLength(1); + await expect(pendingSubagentEmits[0]).resolves.toBeUndefined(); + }); +}); diff --git a/api/server/controllers/agents/__tests__/responses.unit.spec.js b/api/server/controllers/agents/__tests__/responses.unit.spec.js index 943411d3e7..f3500996d6 100644 --- a/api/server/controllers/agents/__tests__/responses.unit.spec.js +++ b/api/server/controllers/agents/__tests__/responses.unit.spec.js @@ -159,6 +159,8 @@ jest.mock('@librechat/api', () => ({ getTransactionsConfig: mockGetTransactionsConfig, recordCollectedUsage: mockRecordCollectedUsage, createSubagentUsageSink: jest.fn().mockReturnValue(jest.fn()), + CHILD_THREAD_READ_ONLY_ERROR: + 'This subagent thread is view-only. Continue it from its parent agent or create a separate chat.', getLangfuseTraceMessageFields: jest.fn().mockResolvedValue({ langfuseSampled: true, langfuseDestinationIds: ['destination-1'], @@ -639,6 +641,41 @@ describe('createResponse controller', () => { ); }); + it('rejects a remote response continuation of a view-only subagent thread', async () => { + const { + validateResponseRequest, + sendResponsesErrorResponse, + CHILD_THREAD_READ_ONLY_ERROR, + } = require('@librechat/api'); + const { getConvo, saveConvo, saveMessage } = require('~/models'); + validateResponseRequest.mockReturnValueOnce({ + request: { + model: 'agent-123', + input: 'Mutate the child.', + stream: false, + store: true, + previous_response_id: 'child-thread', + }, + }); + getConvo.mockResolvedValueOnce({ + conversationId: 'child-thread', + user: 'user-123', + subagentThread: { parentConversationId: 'parent-thread' }, + }); + + await createResponse(req, res); + + expect(sendResponsesErrorResponse).toHaveBeenCalledWith( + res, + 409, + CHILD_THREAD_READ_ONLY_ERROR, + 'invalid_request', + 'conversation_read_only', + ); + expect(saveConvo).not.toHaveBeenCalled(); + expect(saveMessage).not.toHaveBeenCalled(); + }); + it('should return 500 when getConvo throws a DB error', async () => { const { validateResponseRequest, sendResponsesErrorResponse } = require('@librechat/api'); const { getConvo } = require('~/models'); diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 70fdec95b0..3da58904e8 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -21,6 +21,7 @@ const { applyContextToAgent, isMemoryAgentEnabled, recordCollectedUsage, + createDetachedSubagentUsageRecorder, sendEvent, computeUsageCostUSD, aggregateEmittedUsage, @@ -225,6 +226,11 @@ class AgentClient extends BaseClient { * these before returning — otherwise job cleanup can race the persist. * @type {Promise[]} */ this.pendingSubagentEmits = []; + /** Stable per-generation sequence for subagent usage events. Detached + * usage is billed outside `collectedUsage`, so array length is no longer + * a valid sequence source. @type {number} */ + this.subagentUsageSeq = + usageEmitSink?.filter((event) => event?.usage_type === 'subagent').length ?? 0; /** @type {AgentClientOptions} */ this.options = Object.assign({ endpoint: options.endpoint }, clientOptions); /** @type {string} */ @@ -397,6 +403,30 @@ class AgentClient extends BaseClient { }; } + /** + * Registers the parent conversation write as a child-dispatch prerequisite. + * The store retains only the persistence promise, never this request-scoped client. + * @param {string} message + * @param {Record} [opts] + */ + async sendMessage(message, opts = {}) { + const subagentTasks = this.options?.subagentTasks; + const store = subagentTasks?.store; + if (typeof store?.registerParentPersistence !== 'function') { + return super.sendMessage(message, opts); + } + const getReqData = opts.getReqData; + return super.sendMessage(message, { + ...opts, + getReqData: (data = {}) => { + getReqData?.(data); + if (data.userMessagePromise instanceof Promise) { + store.registerParentPersistence(subagentTasks.scopeId, data.userMessagePromise); + } + }, + }); + } + setOptions(_options) {} /** @@ -2431,39 +2461,67 @@ class AgentClient extends BaseClient { * @returns {((usage: UsageMetadata) => void) | undefined} */ buildSubagentUsageEmitter(appConfig) { - const res = this.options.res; - const streamId = this.options.req?._resumableStreamId || null; + /** Detached children can report usage after `disposeClient` has cleared the + * parent client. Snapshot every value the emitter needs now; the returned + * callback must not dereference mutable client state. */ + const options = this.options; + const res = options.res; + const streamId = options.req?._resumableStreamId || null; if (!res && !streamId) { return undefined; } const includeCost = appConfig?.interfaceConfig?.contextCost === true; + const responseMessageId = this.responseMessageId; + const jobCreatedAt = this.jobCreatedAt; + const usageEmitSink = this.usageEmitSink; + const pendingSubagentEmits = this.pendingSubagentEmits; + const endpointTokenConfig = options.endpointTokenConfig; + const endpointTokenConfigByAgentId = + options.endpointTokenConfigByAgentId instanceof Map + ? new Map(options.endpointTokenConfigByAgentId) + : options.endpointTokenConfigByAgentId; + let subagentUsageSeq = this.subagentUsageSeq; return (usage) => { + subagentUsageSeq += 1; + const cache_creation = + usage.input_token_details?.cache_creation ?? usage.cache_creation_input_tokens; + const cache_read = usage.input_token_details?.cache_read ?? usage.cache_read_input_tokens; const data = { input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, total_tokens: usage.total_tokens, - input_token_details: this.subagentCacheDetails(usage), + input_token_details: + cache_creation == null && cache_read == null ? undefined : { cache_creation, cache_read }, model: usage.model, provider: usage.provider, usage_type: 'subagent', - runId: this.responseMessageId, - /** Unique per collected entry (post-push length) for resume dedupe */ - seq: this.collectedUsage.length, + runId: jobCreatedAt != null ? `${responseMessageId}:${jobCreatedAt}` : responseMessageId, + /** Unique per child call for reconnect/resume dedupe. */ + seq: subagentUsageSeq, /** Price with the SUBAGENT's own endpoint token config (its endpoint may * differ from the parent's); `usage.agentId` is tagged by the sink. */ cost: includeCost ? computeUsageCostUSD( usage, { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier }, - this.resolveAgentEndpointTokenConfig(usage), + resolveAgentTokenConfig({ + agentId: usage?.agentId, + byAgentId: endpointTokenConfigByAgentId, + fallback: endpointTokenConfig, + }), ) : undefined, }; + if (data.cost != null) { + /** The detached task collector persists this same usage object on the + * child message after the emitter has attached authoritative cost. */ + usage.cost = data.cost; + } /** Fold into the response's usage rollup (synchronously, regardless of * emit success) so the persisted total matches the live session, which * also folds subagent usage into its cost/totals. */ - if (this.usageEmitSink) { - this.usageEmitSink.push(data); + if (usageEmitSink) { + usageEmitSink.push(data); } /** The sink fires this without awaiting, so retain the promise and flush * it in chatCompletion's finally — emitChunk persists (HSET) before @@ -2478,7 +2536,7 @@ class AgentClient extends BaseClient { event: UsageEvents.ON_TOKEN_USAGE, data, }, - { expectedCreatedAt: this.jobCreatedAt }, + { expectedCreatedAt: jobCreatedAt }, ); } else { sendEvent(res, { event: UsageEvents.ON_TOKEN_USAGE, data }); @@ -2487,20 +2545,45 @@ class AgentClient extends BaseClient { logger.warn('[AgentClient] Failed to emit subagent usage', err); } })(); - this.pendingSubagentEmits.push(emit); + pendingSubagentEmits.push(emit); return emit; }; } - /** Normalizes a subagent usage event's cache token details for emission. */ - subagentCacheDetails(usage) { - const cache_creation = - usage.input_token_details?.cache_creation ?? usage.cache_creation_input_tokens; - const cache_read = usage.input_token_details?.cache_read ?? usage.cache_read_input_tokens; - if (cache_creation == null && cache_read == null) { - return undefined; - } - return { cache_creation, cache_read }; + /** + * Detached children may outlive the parent turn's one-time billing flush. + * Bill each detached model call on the SDK's awaited usage path; foreground + * children continue to batch with the parent turn. + * @param {AppConfig['balance']} balance + * @param {AppConfig['transactions']} transactions + * @returns {(usage: UsageMetadata) => Promise} + */ + buildDetachedSubagentUsageRecorder(balance, transactions) { + const options = this.options; + const billing = { + user: this.user ?? options?.req?.user?.id, + conversationId: this.conversationId, + messageId: this.responseMessageId, + model: this.model ?? options?.agent?.model_parameters?.model, + endpointTokenConfig: options?.endpointTokenConfig, + endpointTokenConfigByAgentId: options?.endpointTokenConfigByAgentId, + }; + return createDetachedSubagentUsageRecorder( + { + spendTokens: db.spendTokens, + spendStructuredTokens: db.spendStructuredTokens, + pricing: { + getMultiplier: db.getMultiplier, + getCacheMultiplier: db.getCacheMultiplier, + }, + bulkWriteOps: { + insertMany: db.bulkInsertTransactions, + updateBalance: db.updateBalance, + }, + isPrincipalActive: db.isAgentTriggerPrincipalActive, + }, + { ...billing, balance, transactions }, + ); } /** @@ -3170,17 +3253,18 @@ class AgentClient extends BaseClient { summarizationConfig: appConfig?.summarization, appConfig, tokenCounter, - /** Bills subagent child-run model calls — child graphs execute - * outside the streamEvents loop, so ModelEndHandler never sees - * them. Entries land in collectedUsage tagged - * `usage_type: 'subagent'` and are spent by recordCollectedUsage. + /** Bills subagent child-run model calls — foreground usage joins + * the parent batch, while detached usage is recorded per call and + * persisted with its child result because it may outlive this turn. * The sink also streams each as an `on_token_usage` event so the * gauge's session cost/totals include billed subagent usage (the * `subagent` tag keeps it out of the live context meter). */ subagentUsageSink: createSubagentUsageSink( this.collectedUsage, this.buildSubagentUsageEmitter(appConfig), + this.buildDetachedSubagentUsageRecorder(balanceConfig, transactionsConfig), ), + subagentTasks: this.options.subagentTasks, }).then((createdRun) => { if (!createdRun) { throw new Error('Failed to create run'); @@ -3568,7 +3652,9 @@ class AgentClient extends BaseClient { subagentUsageSink: createSubagentUsageSink( this.collectedUsage, this.buildSubagentUsageEmitter(appConfig), + this.buildDetachedSubagentUsageRecorder(balanceConfig, transactionsConfig), ), + subagentTasks: this.options.subagentTasks, }); if (!run) { diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 5405ad01fc..227a8a4ad1 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -2,6 +2,9 @@ const mockCreateRun = jest.fn(); const mockCaptureAgentCheckpointGeneration = jest.fn(); const mockDeleteAgentCheckpoint = jest.fn(); const mockIsHITLEnabled = jest.fn().mockReturnValue(false); +const mockRecordCollectedUsage = jest.fn(); +const mockDetachedUsageRecorder = jest.fn(); +const mockCreateDetachedSubagentUsageRecorder = jest.fn(() => mockDetachedUsageRecorder); const mockBuildAgentScopedContext = jest.fn((...args) => jest.requireActual('@librechat/api').buildAgentScopedContext(...args), ); @@ -15,6 +18,7 @@ const mockFormatAgentMessages = jest.fn(() => ({ const { Providers } = require('@librechat/agents'); const { Constants, ContentTypes, EModelEndpoint } = require('librechat-data-provider'); const { GenerationJobManager, createStreamServices } = require('@librechat/api'); +const BaseClient = require('~/app/clients/BaseClient'); const AgentClient = require('./client'); const { resolveConfigServers } = require('~/server/services/MCP'); @@ -43,6 +47,8 @@ jest.mock('@librechat/api', () => ({ countFormattedMessageTokens: jest.fn(() => 42), countTokens: jest.fn((text) => Math.ceil(String(text ?? '').length / 4)), createTokenCounter: jest.fn(() => jest.fn(() => 0)), + createDetachedSubagentUsageRecorder: (...args) => + mockCreateDetachedSubagentUsageRecorder(...args), captureAgentCheckpointGeneration: (...args) => mockCaptureAgentCheckpointGeneration(...args), deleteAgentCheckpoint: (...args) => mockDeleteAgentCheckpoint(...args), decrementPendingRequest: jest.fn(async () => {}), @@ -57,8 +63,106 @@ jest.mock('@librechat/api', () => ({ }), loadAgent: jest.fn(), maybePrewarmCodeSandbox: jest.fn(), + recordCollectedUsage: (...args) => mockRecordCollectedUsage(...args), })); +describe('AgentClient - detached subagent usage', () => { + it('records each detached call from an immutable snapshot after parent disposal', async () => { + mockRecordCollectedUsage.mockClear(); + mockCreateDetachedSubagentUsageRecorder.mockClear(); + mockDetachedUsageRecorder.mockClear(); + const client = Object.create(AgentClient.prototype); + const childTokenConfig = { input: 1, output: 2 }; + client.user = 'user-123'; + client.conversationId = 'conversation-123'; + client.responseMessageId = 'response-123'; + client.model = 'primary-model'; + client.options = { + req: { user: { id: 'request-user' } }, + agent: { model_parameters: { model: 'fallback-model' } }, + endpointTokenConfig: { input: 3, output: 4 }, + endpointTokenConfigByAgentId: new Map([['agent-child', childTokenConfig]]), + }; + const balance = { enabled: true }; + const transactions = { enabled: true }; + const usage = { + usage_type: 'subagent', + input_tokens: 100, + output_tokens: 20, + agentId: 'agent-child', + }; + + const recordUsage = client.buildDetachedSubagentUsageRecorder(balance, transactions); + client.user = null; + client.conversationId = null; + client.responseMessageId = null; + client.model = null; + client.options = null; + + await recordUsage(usage); + + expect(mockCreateDetachedSubagentUsageRecorder).toHaveBeenCalledTimes(1); + const [deps, billing] = mockCreateDetachedSubagentUsageRecorder.mock.calls[0]; + expect(deps).toEqual({ + spendTokens: expect.any(Function), + spendStructuredTokens: expect.any(Function), + pricing: { + getMultiplier: expect.any(Function), + getCacheMultiplier: expect.any(Function), + }, + bulkWriteOps: { + insertMany: expect.any(Function), + updateBalance: expect.any(Function), + }, + isPrincipalActive: expect.any(Function), + }); + expect(billing).toEqual({ + user: 'user-123', + conversationId: 'conversation-123', + model: 'primary-model', + messageId: 'response-123', + balance, + transactions, + endpointTokenConfig: { input: 3, output: 4 }, + endpointTokenConfigByAgentId: expect.any(Map), + }); + expect(billing.endpointTokenConfigByAgentId.get('agent-child')).toBe(childTokenConfig); + expect(mockDetachedUsageRecorder).toHaveBeenCalledWith(usage); + }); +}); + +describe('AgentClient - subagent parent persistence', () => { + it('registers the parent user-message write before detached child dispatch can proceed', async () => { + const userMessagePromise = Promise.resolve({ + message: { messageId: 'parent-user-message', conversationId: 'parent-conversation' }, + }); + const registerParentPersistence = jest.fn(); + const upstreamGetReqData = jest.fn(); + const baseSend = jest + .spyOn(BaseClient.prototype, 'sendMessage') + .mockImplementation(async (_message, opts) => { + opts.getReqData({ userMessagePromise }); + return { ok: true }; + }); + const client = Object.create(AgentClient.prototype); + client.options = { + subagentTasks: { + scopeId: 'trusted-parent-scope', + store: { registerParentPersistence }, + }, + }; + + await client.sendMessage('Start the parent turn.', { getReqData: upstreamGetReqData }); + + expect(upstreamGetReqData).toHaveBeenCalledWith({ userMessagePromise }); + expect(registerParentPersistence).toHaveBeenCalledWith( + 'trusted-parent-scope', + userMessagePromise, + ); + baseSend.mockRestore(); + }); +}); + describe('AgentClient - label settlement', () => { it('drains a trailing fill enqueued by an in-flight reasoning revision', async () => { const client = Object.create(AgentClient.prototype); @@ -230,9 +334,16 @@ jest.mock('~/server/services/MCP', () => ({ })); jest.mock('~/models', () => ({ + bulkInsertTransactions: jest.fn(), + getCacheMultiplier: jest.fn(), getAgent: jest.fn(), + getMultiplier: jest.fn(), getRoleByName: jest.fn(), getFormattedMemories: jest.fn(), + isAgentTriggerPrincipalActive: jest.fn().mockResolvedValue(true), + spendStructuredTokens: jest.fn(), + spendTokens: jest.fn(), + updateBalance: jest.fn(), })); // Mock getMCPManager diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 60f3bcf9de..fea9763940 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -70,10 +70,11 @@ function getInitializationFailure(error) { }; } -function resolveConversationCreatedAt({ userId, conversationId, isNewConvo }) { +function resolveConversationCreatedAt({ userId, conversationId, isNewConvo, conversation }) { return resolveConversationAnchor({ isNewConversation: isNewConvo, - loadConversation: () => getConvo(userId, conversationId), + loadConversation: () => + conversation !== undefined ? Promise.resolve(conversation) : getConvo(userId, conversationId), onLoadError: (error) => { logger.warn('[AgentController] Failed to resolve conversation timestamp anchor', { conversationId, @@ -443,6 +444,9 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit userId, conversationId, isNewConvo, + conversation: Object.prototype.hasOwnProperty.call(req, 'resolvedConversation') + ? req.resolvedConversation + : undefined, }); if ( diff --git a/api/server/controllers/agents/responses.js b/api/server/controllers/agents/responses.js index 9949259c88..3fd5c45d3d 100644 --- a/api/server/controllers/agents/responses.js +++ b/api/server/controllers/agents/responses.js @@ -54,6 +54,7 @@ const { createAggregatorEventHandlers, getLangfuseTraceMessageFields, stripActivityLabelParts, + CHILD_THREAD_READ_ONLY_ERROR, } = require('@librechat/api'); const { createResponsesToolEndCallback, @@ -368,9 +369,22 @@ const executeResponse = async (envelope, { req, res }) => { 'invalid_request', ); } - if (!(await db.getConvo(principal.userId, request.previous_response_id))) { + const previousConversation = await db.getConvo( + principal.userId, + request.previous_response_id, + ); + if (!previousConversation) { return sendResponsesErrorResponse(res, 404, 'Conversation not found', 'not_found'); } + if (previousConversation.subagentThread != null) { + return sendResponsesErrorResponse( + res, + 409, + CHILD_THREAD_READ_ONLY_ERROR, + 'invalid_request', + 'conversation_read_only', + ); + } } const conversationId = request.previous_response_id ?? uuidv4(); diff --git a/api/server/middleware/validate/subagentThreadTurn.js b/api/server/middleware/validate/subagentThreadTurn.js new file mode 100644 index 0000000000..07374a325b --- /dev/null +++ b/api/server/middleware/validate/subagentThreadTurn.js @@ -0,0 +1,8 @@ +const { createSubagentThreadTurnGuard } = require('@librechat/api'); +const subagentThreadTaskStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); +const db = require('~/models'); + +module.exports = createSubagentThreadTurnGuard({ + getConvo: db.getConvo, + store: subagentThreadTaskStore, +}); diff --git a/api/server/routes/__test-utils__/convos-route-mocks.js b/api/server/routes/__test-utils__/convos-route-mocks.js index 72a758c87e..d6ac3ca0da 100644 --- a/api/server/routes/__test-utils__/convos-route-mocks.js +++ b/api/server/routes/__test-utils__/convos-route-mocks.js @@ -121,4 +121,9 @@ module.exports = { })), assistantEndpoint: () => ({ initializeClient: jest.fn() }), + + subagentThreadStore: () => ({ + cancelForConversations: jest.fn(), + cancelForOwner: jest.fn(), + }), }; diff --git a/api/server/routes/__tests__/convos-duplicate-ratelimit.spec.js b/api/server/routes/__tests__/convos-duplicate-ratelimit.spec.js index a75c11ccba..08d71735a0 100644 --- a/api/server/routes/__tests__/convos-duplicate-ratelimit.spec.js +++ b/api/server/routes/__tests__/convos-duplicate-ratelimit.spec.js @@ -38,6 +38,9 @@ jest.mock('~/server/routes/files/multer', () => require(MOCKS).multerSetup()); jest.mock('multer', () => require(MOCKS).multerLib()); jest.mock('~/server/services/Endpoints/azureAssistants', () => require(MOCKS).assistantEndpoint()); jest.mock('~/server/services/Endpoints/assistants', () => require(MOCKS).assistantEndpoint()); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => + require(MOCKS).subagentThreadStore(), +); describe('POST /api/convos/duplicate - Rate Limiting', () => { let app; diff --git a/api/server/routes/__tests__/convos.spec.js b/api/server/routes/__tests__/convos.spec.js index 1645a4b8dc..ec5b23f263 100644 --- a/api/server/routes/__tests__/convos.spec.js +++ b/api/server/routes/__tests__/convos.spec.js @@ -18,6 +18,9 @@ jest.mock('~/server/routes/files/multer', () => require(MOCKS).multerSetup()); jest.mock('multer', () => require(MOCKS).multerLib()); jest.mock('~/server/services/Endpoints/azureAssistants', () => require(MOCKS).assistantEndpoint()); jest.mock('~/server/services/Endpoints/assistants', () => require(MOCKS).assistantEndpoint()); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => + require(MOCKS).subagentThreadStore(), +); describe('Convos Routes', () => { let app; @@ -28,6 +31,7 @@ describe('Convos Routes', () => { deleteAllSharedLinksWithCleanup, deleteConvoSharedLinksWithCleanup, } = require('@librechat/api'); + const subagentThreadStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); beforeAll(() => { convosRouter = require('../convos'); @@ -61,6 +65,7 @@ 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); }); it('should delete all conversations, tool calls, and shared links for a user', async () => { @@ -234,6 +239,43 @@ describe('Convos Routes', () => { }); describe('DELETE /', () => { + it('cancels root and descendant leases and cleans every cascaded conversation', async () => { + deleteConvos.mockResolvedValue({ + deletedCount: 2, + conversationIds: ['parent-conversation', 'child-conversation'], + }); + deleteToolCalls.mockResolvedValue({ deletedCount: 1 }); + deleteConvoSharedLinksWithCleanup.mockResolvedValue({ deletedCount: 1 }); + + const response = await request(app) + .delete('/api/convos') + .send({ + arg: { conversationId: 'parent-conversation' }, + }); + + expect(response.status).toBe(201); + expect(subagentThreadStore.cancelForConversations).toHaveBeenNthCalledWith( + 1, + 'test-user-123', + ['parent-conversation'], + undefined, + ); + expect(subagentThreadStore.cancelForConversations).toHaveBeenNthCalledWith( + 2, + 'test-user-123', + ['parent-conversation', 'child-conversation'], + undefined, + ); + expect(deleteToolCalls.mock.calls.map((call) => call[1])).toEqual([ + 'parent-conversation', + 'child-conversation', + ]); + expect(deleteConvoSharedLinksWithCleanup.mock.calls.map((call) => call[1])).toEqual([ + 'parent-conversation', + 'child-conversation', + ]); + }); + it('should delete a single conversation, tool calls, and associated shared links', async () => { const mockConversationId = 'conv-123'; const mockDbResponse = { diff --git a/api/server/routes/__tests__/messages-content-edit.spec.js b/api/server/routes/__tests__/messages-content-edit.spec.js index 51b97a4a8c..4c12efff2e 100644 --- a/api/server/routes/__tests__/messages-content-edit.spec.js +++ b/api/server/routes/__tests__/messages-content-edit.spec.js @@ -12,8 +12,12 @@ jest.mock('@librechat/api', () => ({ sendFeedbackScore: jest.fn().mockResolvedValue(undefined), traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), mergeQuotedTextForCount: jest.fn((text) => text), + CHILD_THREAD_READ_ONLY_ERROR: 'Child thread is view-only.', + isSubagentThreadWriteBlocked: jest.fn().mockResolvedValue(false), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({})); + jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), logger: { @@ -65,6 +69,7 @@ describe('PUT /:conversationId/:messageId content edit', () => { it('preserves content-part metadata when editing its text', async () => { getMessages.mockResolvedValue([ { + conversationId: 'conversation-1', tokenCount: 10, content: [ { @@ -101,6 +106,7 @@ describe('PUT /:conversationId/:messageId content edit', () => { it('clears the generated reasoning title when its reasoning text is edited', async () => { getMessages.mockResolvedValue([ { + conversationId: 'conversation-1', tokenCount: 10, content: [ { @@ -152,6 +158,7 @@ describe('PUT /:conversationId/:messageId content edit', () => { getMessages.mockResolvedValue([ { + conversationId: 'conversation-1', tokenCount: 10, content: [ { diff --git a/api/server/routes/__tests__/messages-delete.spec.js b/api/server/routes/__tests__/messages-delete.spec.js index 44ffaf56d7..0661b94d96 100644 --- a/api/server/routes/__tests__/messages-delete.spec.js +++ b/api/server/routes/__tests__/messages-delete.spec.js @@ -13,8 +13,12 @@ jest.mock('@librechat/api', () => ({ countTokens: jest.fn().mockResolvedValue(10), sendFeedbackScore: jest.fn().mockResolvedValue(undefined), traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), + CHILD_THREAD_READ_ONLY_ERROR: 'Child thread is view-only.', + isSubagentThreadWriteBlocked: jest.fn().mockResolvedValue(false), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({})); + jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), logger: { diff --git a/api/server/routes/__tests__/messages-feedback.spec.js b/api/server/routes/__tests__/messages-feedback.spec.js index 6325e9bc14..92aac83a1a 100644 --- a/api/server/routes/__tests__/messages-feedback.spec.js +++ b/api/server/routes/__tests__/messages-feedback.spec.js @@ -10,8 +10,12 @@ jest.mock('@librechat/api', () => ({ countTokens: jest.fn().mockResolvedValue(10), sendFeedbackScore: jest.fn().mockResolvedValue(undefined), traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), + CHILD_THREAD_READ_ONLY_ERROR: 'Child thread is view-only.', + isSubagentThreadWriteBlocked: jest.fn().mockResolvedValue(false), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({})); + jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), logger: { diff --git a/api/server/routes/__tests__/messages-get-real-validation.spec.js b/api/server/routes/__tests__/messages-get-real-validation.spec.js index bafb54977f..cc4c3956cb 100644 --- a/api/server/routes/__tests__/messages-get-real-validation.spec.js +++ b/api/server/routes/__tests__/messages-get-real-validation.spec.js @@ -32,8 +32,12 @@ jest.mock('@librechat/api', () => ({ getJob: jest.fn(), }, isPendingActionStale: jest.fn(() => false), + CHILD_THREAD_READ_ONLY_ERROR: 'Child thread is view-only.', + isSubagentThreadWriteBlocked: jest.fn().mockResolvedValue(false), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({})); + jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), logger: { diff --git a/api/server/routes/__tests__/messages-get.spec.js b/api/server/routes/__tests__/messages-get.spec.js index bf3daea5bf..2e26f20ee3 100644 --- a/api/server/routes/__tests__/messages-get.spec.js +++ b/api/server/routes/__tests__/messages-get.spec.js @@ -11,8 +11,12 @@ jest.mock('@librechat/api', () => ({ countTokens: jest.fn().mockResolvedValue(10), sendFeedbackScore: jest.fn().mockResolvedValue(undefined), traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), + CHILD_THREAD_READ_ONLY_ERROR: 'Child thread is view-only.', + isSubagentThreadWriteBlocked: jest.fn().mockResolvedValue(false), })); +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({})); + jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), logger: { diff --git a/api/server/routes/__tests__/messages-subagent-thread.spec.js b/api/server/routes/__tests__/messages-subagent-thread.spec.js new file mode 100644 index 0000000000..9d1560b278 --- /dev/null +++ b/api/server/routes/__tests__/messages-subagent-thread.spec.js @@ -0,0 +1,154 @@ +const express = require('express'); +const request = require('supertest'); + +const mockIsSubagentThreadWriteBlocked = jest.fn(); + +jest.mock('@librechat/agents', () => ({ sleep: jest.fn() })); + +jest.mock('@librechat/api', () => ({ + unescapeLaTeX: jest.fn((value) => value), + countTokens: jest.fn().mockResolvedValue(1), + sendFeedbackScore: jest.fn().mockResolvedValue(undefined), + traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), + mergeQuotedTextForCount: jest.fn((text) => text), + CHILD_THREAD_READ_ONLY_ERROR: 'This subagent thread is view-only.', + isSubagentThreadWriteBlocked: (...args) => mockIsSubagentThreadWriteBlocked(...args), +})); + +jest.mock('@librechat/data-schemas', () => ({ + ...jest.requireActual('@librechat/data-schemas'), + logger: { + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + }, +})); + +jest.mock('librechat-data-provider', () => ({ + ...jest.requireActual('librechat-data-provider'), +})); + +jest.mock('~/models', () => ({ + getConvo: jest.fn(), + getMessage: jest.fn(), + getMessages: jest.fn(), + saveConvo: jest.fn(), + saveMessage: jest.fn(), + updateMessage: jest.fn(), + deleteMessages: jest.fn(), + getConvosQueried: jest.fn(), + searchMessages: jest.fn(), + getMessagesByCursor: jest.fn(), +})); + +jest.mock('~/server/services/Endpoints/agents/subagentThreadStore', () => ({ + isThreadActiveForOwner: jest.fn(), +})); + +jest.mock('~/server/services/Artifacts/update', () => ({ + findAllArtifacts: jest.fn(), + replaceArtifactContent: jest.fn(), +})); + +jest.mock('~/server/middleware', () => ({ + requireJwtAuth: (req, _res, next) => next(), + validateMessageReq: (req, _res, next) => next(), + configMiddleware: (req, _res, next) => next(), + sendValidationResponse: jest.fn(), + prepareMessageRequestValidation: jest.fn(), +})); + +describe('message mutation policy for durable subagent threads', () => { + let app; + const db = require('~/models'); + + beforeAll(() => { + const messagesRouter = require('../messages'); + app = express(); + app.use(express.json()); + app.use((req, _res, next) => { + req.user = { id: 'owner-user', tenantId: 'tenant-a' }; + next(); + }); + app.use('/api/messages', messagesRouter); + }); + + beforeEach(() => { + jest.clearAllMocks(); + mockIsSubagentThreadWriteBlocked.mockResolvedValue(true); + db.getMessage.mockResolvedValue({ + messageId: 'message-1', + conversationId: 'child-conversation', + isCreatedByUser: false, + content: [], + }); + db.getMessages.mockResolvedValue([ + { + messageId: 'message-1', + conversationId: 'child-conversation', + isCreatedByUser: false, + content: [], + }, + ]); + }); + + it('blocks every transcript-affecting message route through one shared policy', async () => { + const responses = await Promise.all([ + request(app) + .post('/api/messages/child-conversation') + .send({ messageId: 'new-message', text: 'write' }), + request(app) + .put('/api/messages/child-conversation/message-1') + .send({ text: 'edit', model: 'gpt-5' }), + request(app).delete('/api/messages/child-conversation/message-1'), + request(app).post('/api/messages/branch').send({ + messageId: 'message-1', + agentId: 'agent-1', + }), + request(app).post('/api/messages/artifact/message-1').send({ + index: 0, + original: 'before', + updated: 'after', + }), + ]); + + expect(responses.map((response) => response.status)).toEqual([409, 409, 409, 409, 409]); + for (const response of responses) { + expect(response.body).toEqual({ error: 'This subagent thread is view-only.' }); + } + expect(mockIsSubagentThreadWriteBlocked).toHaveBeenCalledTimes(5); + expect(mockIsSubagentThreadWriteBlocked).toHaveBeenCalledWith( + expect.objectContaining({ getConvo: db.getConvo }), + { + userId: 'owner-user', + conversationId: 'child-conversation', + tenantId: 'tenant-a', + }, + ); + expect(db.saveMessage).not.toHaveBeenCalled(); + expect(db.saveConvo).not.toHaveBeenCalled(); + expect(db.updateMessage).not.toHaveBeenCalled(); + expect(db.deleteMessages).not.toHaveBeenCalled(); + }); + + it('authorizes edits against the message owner conversation, not a writable URL', async () => { + mockIsSubagentThreadWriteBlocked.mockResolvedValue(true); + db.getMessages.mockResolvedValue([ + { + messageId: 'message-1', + conversationId: 'child-conversation', + isCreatedByUser: true, + content: [], + }, + ]); + + const response = await request(app) + .put('/api/messages/ordinary-conversation/message-1') + .send({ text: 'forged edit', model: 'gpt-5' }); + + expect(response.status).toBe(404); + expect(mockIsSubagentThreadWriteBlocked).not.toHaveBeenCalled(); + expect(db.updateMessage).not.toHaveBeenCalled(); + }); +}); diff --git a/api/server/routes/agents/chat.js b/api/server/routes/agents/chat.js index 1cd52c6d8c..9b466ad1b1 100644 --- a/api/server/routes/agents/chat.js +++ b/api/server/routes/agents/chat.js @@ -17,6 +17,7 @@ const { canAccessAgentFromBody, } = require('~/server/middleware'); const { initializeClient } = require('~/server/services/Endpoints/agents'); +const guardSubagentThreadTurn = require('~/server/middleware/validate/subagentThreadTurn'); const AgentController = require('~/server/controllers/agents/request'); const ResumeController = require('~/server/controllers/agents/resume'); const addTitle = require('~/server/services/Endpoints/agents/title'); @@ -73,6 +74,7 @@ router.use(moderateText); router.use(checkAgentAccess); router.use(checkAgentResourceAccess); router.use(validateConvoAccess); +router.use(guardSubagentThreadTurn); router.use(buildEndpointOption); const controller = async (req, res, next) => { diff --git a/api/server/routes/assistants/chatV1.js b/api/server/routes/assistants/chatV1.js index 67bfc007a6..0888a78e1d 100644 --- a/api/server/routes/assistants/chatV1.js +++ b/api/server/routes/assistants/chatV1.js @@ -8,6 +8,7 @@ const { buildEndpointOption, } = require('~/server/middleware'); const validateConvoAccess = require('~/server/middleware/validate/convoAccess'); +const guardSubagentThreadTurn = require('~/server/middleware/validate/subagentThreadTurn'); const validateAssistant = require('~/server/middleware/assistants/validate'); const chatController = require('~/server/controllers/assistants/chatV1'); @@ -27,6 +28,7 @@ router.post( buildEndpointOption, validateAssistant, validateConvoAccess, + guardSubagentThreadTurn, setHeaders, chatController, ); diff --git a/api/server/routes/assistants/chatV2.js b/api/server/routes/assistants/chatV2.js index 4612743e47..d2a1c64d92 100644 --- a/api/server/routes/assistants/chatV2.js +++ b/api/server/routes/assistants/chatV2.js @@ -8,6 +8,7 @@ const { buildEndpointOption, } = require('~/server/middleware'); const validateConvoAccess = require('~/server/middleware/validate/convoAccess'); +const guardSubagentThreadTurn = require('~/server/middleware/validate/subagentThreadTurn'); const validateAssistant = require('~/server/middleware/assistants/validate'); const chatController = require('~/server/controllers/assistants/chatV2'); @@ -27,6 +28,7 @@ router.post( buildEndpointOption, validateAssistant, validateConvoAccess, + guardSubagentThreadTurn, setHeaders, chatController, ); diff --git a/api/server/routes/convos.js b/api/server/routes/convos.js index bcd83f0778..87a9d0a719 100644 --- a/api/server/routes/convos.js +++ b/api/server/routes/convos.js @@ -22,6 +22,7 @@ const { forkConversation, duplicateConversation } = require('~/server/utils/impo const { storage, importFileFilter } = require('~/server/routes/files/multer'); const requireJwtAuth = require('~/server/middleware/requireJwtAuth'); const { importConversations } = require('~/server/utils/import'); +const subagentThreadTaskStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); const getLogStores = require('~/cache/getLogStores'); const db = require('~/models'); @@ -149,16 +150,32 @@ router.delete('/', configMiddleware, async (req, res) => { } try { + const tenantId = + typeof req.user.tenantId === 'string' && req.user.tenantId !== '' + ? req.user.tenantId + : undefined; + if (filter.conversationId) { + subagentThreadTaskStore.cancelForConversations( + req.user.id, + [filter.conversationId], + tenantId, + ); + } const dbResponse = await db.deleteConvos(req.user.id, filter); + const deletedConversationIds = + dbResponse.conversationIds ?? (filter.conversationId ? [filter.conversationId] : []); + subagentThreadTaskStore.cancelForConversations(req.user.id, deletedConversationIds, tenantId); // HITL: prune the deleted conversations' durable checkpoints — a paused run's // checkpoint would otherwise persist until the Mongo TTL. Never throws. await deleteAgentCheckpoints( - dbResponse.conversationIds, + deletedConversationIds, req.config?.endpoints?.[EModelEndpoint.agents]?.checkpointer, ); if (filter.conversationId) { - await db.deleteToolCalls(req.user.id, filter.conversationId); - await deleteConvoSharedLinksWithCleanup(req.user.id, filter.conversationId); + await Promise.all(deletedConversationIds.map((id) => db.deleteToolCalls(req.user.id, id))); + await Promise.all( + deletedConversationIds.map((id) => deleteConvoSharedLinksWithCleanup(req.user.id, id)), + ); } res.status(201).json(dbResponse); } catch (error) { @@ -169,6 +186,12 @@ router.delete('/', configMiddleware, async (req, res) => { router.delete('/all', configMiddleware, async (req, res) => { try { + subagentThreadTaskStore.cancelForOwner( + req.user.id, + typeof req.user.tenantId === 'string' && req.user.tenantId !== '' + ? req.user.tenantId + : undefined, + ); const dbResponse = await db.deleteConvos(req.user.id, {}); // HITL: prune ALL the deleted conversations' durable checkpoints in one bulk pass. await deleteAgentCheckpoints( diff --git a/api/server/routes/messages.js b/api/server/routes/messages.js index d95a755015..eb87e7a05c 100644 --- a/api/server/routes/messages.js +++ b/api/server/routes/messages.js @@ -13,7 +13,10 @@ const { sendFeedbackScore, traceIdForMessage, mergeQuotedTextForCount, + CHILD_THREAD_READ_ONLY_ERROR, + isSubagentThreadWriteBlocked, } = require('@librechat/api'); +const subagentThreadTaskStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); const { findAllArtifacts, replaceArtifactContent } = require('~/server/services/Artifacts/update'); const { requireJwtAuth, @@ -27,6 +30,24 @@ const db = require('~/models'); const router = express.Router(); router.use(requireJwtAuth); +async function rejectSubagentThreadWrite(req, res, conversationId) { + const blocked = await isSubagentThreadWriteBlocked( + { getConvo: db.getConvo, store: subagentThreadTaskStore }, + { + userId: req.user.id, + conversationId, + ...(typeof req.user.tenantId === 'string' && req.user.tenantId !== '' + ? { tenantId: req.user.tenantId } + : {}), + }, + ); + if (!blocked) { + return false; + } + res.status(409).json({ error: CHILD_THREAD_READ_ONLY_ERROR }); + return true; +} + router.get('/', async (req, res) => { try { const user = req.user.id ?? ''; @@ -134,6 +155,10 @@ router.post('/branch', async (req, res) => { return res.status(404).json({ error: 'Source message not found' }); } + if (await rejectSubagentThreadWrite(req, res, sourceMessage.conversationId)) { + return; + } + if (sourceMessage.isCreatedByUser) { return res.status(400).json({ error: 'Cannot branch from user messages' }); } @@ -215,6 +240,10 @@ router.post('/artifact/:messageId', async (req, res) => { return res.status(404).json({ error: 'Message not found' }); } + if (await rejectSubagentThreadWrite(req, res, message.conversationId)) { + return; + } + const artifacts = findAllArtifacts(message); if (index >= artifacts.length) { return res.status(400).json({ error: 'Artifact index out of bounds' }); @@ -315,6 +344,9 @@ router.get('/:conversationId', prepareMessageRequestValidation, async (req, res) router.post('/:conversationId', validateMessageReq, async (req, res) => { try { + if (await rejectSubagentThreadWrite(req, res, req.params.conversationId)) { + return; + } const message = { ...req.body, conversationId: req.params.conversationId }; const reqCtx = { userId: req?.user?.id, @@ -365,6 +397,18 @@ router.get('/:conversationId/:messageId', validateMessageReq, async (req, res) = router.put('/:conversationId/:messageId', validateMessageReq, async (req, res) => { try { const { conversationId, messageId } = req.params; + const message = ( + await db.getMessages( + { messageId, user: req.user.id }, + 'conversationId content tokenCount quotes isCreatedByUser', + ) + )?.[0]; + if (!message || message.conversationId !== conversationId) { + return res.status(404).json({ error: 'Message not found' }); + } + if (await rejectSubagentThreadWrite(req, res, message.conversationId)) { + return; + } const { text, index, model } = req.body; if (index === undefined) { @@ -372,16 +416,10 @@ router.put('/:conversationId/:messageId', validateMessageReq, async (req, res) = * every send, but this edit only changes `text`. Count the merged * text+quotes so the stored `tokenCount` stays authoritative (matching the * send path); a plain text-only count under-reports by the quote block. */ - const existing = ( - await db.getMessages( - { conversationId, messageId, user: req.user.id }, - 'quotes isCreatedByUser', - ) - )?.[0]; const textToCount = mergeQuotedTextForCount( text, - existing?.quotes, - existing?.isCreatedByUser === true, + message.quotes, + message.isCreatedByUser === true, ); const tokenCount = await countTokens(textToCount, model); const result = await db.updateMessage(req?.user?.id, { messageId, text, tokenCount }); @@ -392,13 +430,6 @@ router.put('/:conversationId/:messageId', validateMessageReq, async (req, res) = return res.status(400).json({ error: 'Invalid index' }); } - const message = ( - await db.getMessages({ conversationId, messageId, user: req.user.id }, 'content tokenCount') - )?.[0]; - if (!message) { - return res.status(404).json({ error: 'Message not found' }); - } - const existingContent = message.content; if (!Array.isArray(existingContent) || index >= existingContent.length) { return res.status(400).json({ error: 'Invalid index' }); @@ -510,6 +541,9 @@ router.put( router.delete('/:conversationId/:messageId', validateMessageReq, async (req, res) => { try { const { conversationId, messageId } = req.params; + if (await rejectSubagentThreadWrite(req, res, conversationId)) { + return; + } await db.deleteMessages({ messageId, conversationId, user: req.user.id }); res.status(204).send(); } catch (error) { diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 93b248aa41..46be766f7c 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -20,6 +20,7 @@ const { collectCodeExecutionProfileRoutes, getLazySubagentConfigId, createStatefulCodeEnvironmentPolicyError, + buildSubagentThreadTaskConfig, } = require('@librechat/api'); const { ResourceType, @@ -60,6 +61,7 @@ const { getModelsConfig } = require('~/server/controllers/ModelController'); const { checkPermission, findAccessibleResources } = require('~/server/services/PermissionService'); const AgentClient = require('~/server/controllers/agents/client'); const { processAddedConvo } = require('./addedConvo'); +const subagentThreadTaskStore = require('./subagentThreadStore'); const { logViolation } = require('~/cache'); const db = require('~/models'); @@ -148,6 +150,16 @@ const initializeClient = async ({ throw new Error('Endpoint option not provided'); } const appConfig = req.config; + /** The normal controller resolves this once for timestamp anchoring. Reuse + * that trusted document for child-thread execution policy; resume and direct + * callers fall back to the same owner-scoped lookup. */ + const conversationId = req.body?.conversationId; + let requestConversationPromise = Promise.resolve(null); + if (Object.prototype.hasOwnProperty.call(req, 'resolvedConversation')) { + requestConversationPromise = Promise.resolve(req.resolvedConversation); + } else if (typeof conversationId === 'string' && conversationId !== '') { + requestConversationPromise = db.getConvo(req.user.id, conversationId); + } const startupTelemetry = getAgentStartupTelemetry(req); /** @type {string | null} */ @@ -390,25 +402,6 @@ const initializeClient = async ({ /** @type {Array} */ const usageEmitSink = []; - const eventHandlers = getDefaultHandlers({ - res, - contentParts, - stepMap, - toolInputValidationErrors, - toolExecuteOptions, - summarizationOptions, - aggregateContent, - toolEndCallback, - collectedUsage, - collectedThoughtSignatures, - streamId, - jobCreatedAt, - subagentAggregatorsByToolCallId, - usageCost, - contextUsageSink, - usageEmitSink, - }); - const [ memoryAvailable, accessibleSkillIds, @@ -416,6 +409,7 @@ const initializeClient = async ({ skillCreateAllowed, { skillStates, defaultActiveOnShare }, { primaryAgent, modelsConfig }, + requestConversation, ] = await Promise.all([ memoryAvailablePromise, accessibleSkillIdsPromise, @@ -423,6 +417,7 @@ const initializeClient = async ({ skillCreateAllowedPromise, skillStatesPromise, validatedPrimaryAgentPromise, + requestConversationPromise, ]); delete endpointOption.agent; @@ -433,8 +428,6 @@ const initializeClient = async ({ const loadTools = createToolLoader(signal, streamId, true, jobCreatedAt); /** @type {Array} */ const requestFiles = req.body.files ?? []; - /** @type {string} */ - const conversationId = req.body.conversationId; /** @type {string | undefined} */ const parentMessageId = req.body.parentMessageId; /** @@ -693,7 +686,11 @@ const initializeClient = async ({ // spawn tool, not handoff edges. Explicit children are advertised as inert, // VIEW-checked descriptors; model, tool, MCP, file, and skill initialization // happens only when the SDK selects one. + const atSubagentThreadDepthLimit = !subagentThreadTaskStore.canCreateChildThread( + requestConversation?.subagentThread?.depth ?? 0, + ); const subagentsCapabilityEnabled = enabledCapabilities.has(AgentCapabilities.subagents); + const subagentsAvailableForRun = subagentsCapabilityEnabled && !atSubagentThreadDepthLimit; /** Track skipped ids locally so repeated failures short-circuit within * the subagent loading loop. Seeded from the discovery helper's skip * list so agents that already failed handoff loading don't get retried. */ @@ -997,7 +994,7 @@ const initializeClient = async ({ }; const buildLazySubagentDescriptors = async (agent, depth = 0, ancestors = new Set()) => { - if (!subagentsCapabilityEnabled || !agent.subagents?.enabled) { + if (!subagentsAvailableForRun || !agent.subagents?.enabled) { return []; } if (agent.subagents.allowSelf !== false) { @@ -1165,7 +1162,7 @@ const initializeClient = async ({ async function resolveGraphSubagentsFor(config, graphSignal = signal) { throwIfAborted(graphSignal); const definitions = - subagentsCapabilityEnabled && config.subagents?.enabled === true + subagentsAvailableForRun && config.subagents?.enabled === true ? (config.subagents.graphs ?? []) : []; const resolvedGraphs = []; @@ -1212,15 +1209,53 @@ const initializeClient = async ({ await resolveGraphSubagentsFor(config); } - primaryConfig.subagents = subagentsCapabilityEnabled ? primaryConfig.subagents : undefined; + /** Build detached execution only for an attributable owner/thread. New + * tasks still require a spawnable child, while an existing process-local + * 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. */ + const hasSpawnableSubagent = rootSubagentConfigs.some( + (config) => + config.subagents?.enabled === true && + (config.subagents.allowSelf !== false || + (config.subagentAgentConfigs?.length ?? 0) > 0 || + (config.lazySubagentConfigs?.length ?? 0) > 0 || + (config.subagentGraphConfigs?.length ?? 0) > 0), + ); + const trustedSubagentTasks = + backgroundToolsAvailable && + typeof req.user?.id === 'string' && + req.user.id !== '' && + typeof conversationId === 'string' && + conversationId !== '' + ? buildSubagentThreadTaskConfig(subagentThreadTaskStore, { + userId: req.user.id, + parentConversationId: conversationId, + ...(typeof req.user.tenantId === 'string' && req.user.tenantId !== '' + ? { tenantId: req.user.tenantId } + : {}), + }) + : undefined; + const hasExistingSubagentTask = + trustedSubagentTasks != null && + trustedSubagentTasks.store.list(trustedSubagentTasks.scopeId).length > 0; + const subagentTasks = + trustedSubagentTasks != null && + ((subagentsAvailableForRun && hasSpawnableSubagent) || hasExistingSubagentTask) + ? trustedSubagentTasks + : undefined; + if (subagentTasks != null) { + toolExecuteOptions.subagentTasks = subagentTasks; + } - /** If the capability is off at the endpoint level, strip `subagents` on - * every loaded config — not just the primary. `run.ts` calls + primaryConfig.subagents = subagentsAvailableForRun ? primaryConfig.subagents : undefined; + + /** If the capability is off or this durable child is at the depth limit, + * strip `subagents` on every loaded config — not just the primary. `run.ts` calls * `buildSubagentConfigs` for every agent in the array, so a handoff * agent with `subagents.enabled: true` persisted on its document would - * otherwise still expose self-spawn at runtime even though the admin - * has disabled the capability globally. */ - if (!subagentsCapabilityEnabled) { + * otherwise still expose self-spawn at runtime. */ + if (!subagentsAvailableForRun) { primaryConfig.lazySubagentConfigs = undefined; primaryConfig.subagentGraphConfigs = undefined; for (const config of agentConfigs.values()) { @@ -1306,6 +1341,25 @@ const initializeClient = async ({ fallback: usageCost.endpointTokenConfig, }); + const eventHandlers = getDefaultHandlers({ + res, + contentParts, + stepMap, + toolInputValidationErrors, + toolExecuteOptions, + summarizationOptions, + aggregateContent, + toolEndCallback, + collectedUsage, + collectedThoughtSignatures, + streamId, + jobCreatedAt, + subagentAggregatorsByToolCallId, + usageCost, + contextUsageSink, + usageEmitSink, + }); + const client = new AgentClient({ req, res, @@ -1330,6 +1384,7 @@ const initializeClient = async ({ maxContextTokens: primaryConfig.maxContextTokens, endpoint: isEphemeralAgentId(primaryConfig.id) ? primaryConfig.endpoint : EModelEndpoint.agents, subagentAggregatorsByToolCallId, + subagentTasks, /** Resolved endpoint token/pricing config so spending and cost reflect * configured rates for custom-endpoint agents instead of defaults. */ endpointTokenConfig: primaryConfig.endpointTokenConfig, diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 31753a9924..a5be3f9852 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -138,6 +138,7 @@ describe('initializeClient — processAgent ACL gate', () => { const makeReq = () => ({ user: { id: testUser._id.toString(), role: 'USER' }, body: { conversationId: 'conv_1', files: [] }, + resolvedConversation: null, config: { endpoints: {} }, _resumableStreamId: null, }); @@ -645,6 +646,7 @@ describe('initializeClient — subagent loading', () => { const makeSubagentReq = () => ({ user: { id: testUser._id.toString(), role: 'USER' }, body: { conversationId: 'conv_sub', files: [] }, + resolvedConversation: null, config: { endpoints: { agents: { @@ -714,6 +716,118 @@ describe('initializeClient — subagent loading', () => { return agent; }; + it('creates one trusted durable thread scope for detached subagents', async () => { + mockInitializeAgent.mockResolvedValue( + makePrimaryConfig({ + subagents: { enabled: true, allowSelf: true, agent_ids: [] }, + }), + ); + const req = makeSubagentReq(); + req.config.endpoints.agents.capabilities.push('run_in_background'); + + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + + expect(agentClientArgs.subagentTasks).toBe(capturedToolExecuteOptions.subagentTasks); + expect(agentClientArgs.subagentTasks.store.supportsThreadContinuation).toBe(true); + expect(JSON.parse(agentClientArgs.subagentTasks.scopeId)).toEqual({ + version: 1, + userId: testUser._id.toString(), + parentConversationId: 'conv_sub', + }); + }); + + it('keeps an existing detached task controllable after subagent config is disabled', async () => { + mockInitializeAgent.mockResolvedValue( + makePrimaryConfig({ + subagents: { enabled: true, allowSelf: true, agent_ids: [] }, + }), + ); + const initialReq = makeSubagentReq(); + initialReq.config.endpoints.agents.capabilities.push('run_in_background'); + await initializeClient({ + req: initialReq, + res: {}, + signal: new AbortController().signal, + 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, + }, + ]); + mockInitializeAgent.mockResolvedValue(makePrimaryConfig({})); + const changedReq = makeSubagentReq(); + changedReq.config.endpoints.agents.capabilities.push('run_in_background'); + + await initializeClient({ + req: changedReq, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + + expect(agentClientArgs.subagentTasks).toEqual(existingConfig); + expect(capturedToolExecuteOptions.subagentTasks).toEqual(existingConfig); + expect(agentClientArgs.agent.subagents).toBeUndefined(); + listSpy.mockRestore(); + }); + + it('disables every nested subagent path at the durable child-thread depth limit', async () => { + const primaryConfig = makePrimaryConfig({ + subagents: { enabled: true, allowSelf: true, agent_ids: [] }, + }); + mockInitializeAgent.mockResolvedValue(primaryConfig); + const req = makeSubagentReq(); + req.config.endpoints.agents.capabilities.push('run_in_background'); + req.resolvedConversation = { + conversationId: 'conv_sub', + subagentThread: { depth: 1 }, + }; + + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + + expect(agentClientArgs.subagentTasks).toBeUndefined(); + expect(capturedToolExecuteOptions.subagentTasks).toBeUndefined(); + expect(agentClientArgs.agent.subagents).toBeUndefined(); + }); + + it('keeps detached subagents disabled without the admin background capability', async () => { + mockInitializeAgent.mockResolvedValue( + makePrimaryConfig({ + subagents: { enabled: true, allowSelf: true, agent_ids: [] }, + }), + ); + + await initializeClient({ + req: makeSubagentReq(), + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + + expect(agentClientArgs.subagentTasks).toBeUndefined(); + expect(capturedToolExecuteOptions.subagentTasks).toBeUndefined(); + }); + it('defers pure-subagent MCP initialization until the descriptor is selected', async () => { const subAgent = await createAgent({ id: SUBAGENT_ID, @@ -1210,8 +1324,10 @@ describe('initializeClient — subagent loading', () => { Promise.resolve(agent.id === PRIMARY_ID ? primaryConfig : memberConfigs.get(agent.id)), ); + const req = makeSubagentReq(); + req.config.endpoints.agents.capabilities.push('run_in_background'); const { userMCPAuthMap } = await initializeClient({ - req: makeSubagentReq(), + req, res: {}, signal: new AbortController().signal, endpointOption: makeEndpointOption(), @@ -1224,6 +1340,7 @@ describe('initializeClient — subagent loading', () => { }, ]); expect(memberIds.every((id) => !agentClientArgs.agentConfigs.has(id))).toBe(true); + expect(agentClientArgs.subagentTasks).toBe(capturedToolExecuteOptions.subagentTasks); expect(userMCPAuthMap).toEqual({ server_0: { token: 'token_0' }, server_1: { token: 'token_1' }, diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.js b/api/server/services/Endpoints/agents/subagentThreadStore.js new file mode 100644 index 0000000000..ff1dd6c60b --- /dev/null +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -0,0 +1,25 @@ +const { createSubagentThreadTaskStore } = 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. */ +const subagentThreadTaskStore = createSubagentThreadTaskStore( + { + acquireSubagentThreadLease: db.acquireSubagentThreadLease, + countActiveSubagentThreadLeases: db.countActiveSubagentThreadLeases, + deleteConvos: db.deleteConvos, + deleteMessages: db.deleteMessages, + getConvo: db.getConvo, + getMessages: db.getMessages, + releaseSubagentThreadLease: db.releaseSubagentThreadLease, + reserveSubagentThread: db.reserveSubagentThread, + renewSubagentThreadLease: db.renewSubagentThreadLease, + saveConvo: db.saveConvo, + saveMessage: db.saveMessage, + }, + { + isOwnerActive: db.isAgentTriggerPrincipalActive, + }, +); + +module.exports = subagentThreadTaskStore; diff --git a/api/server/utils/import/fork.spec.js b/api/server/utils/import/fork.spec.js index 6fe5b4f8c9..e93ed9b7a3 100644 --- a/api/server/utils/import/fork.spec.js +++ b/api/server/utils/import/fork.spec.js @@ -146,6 +146,29 @@ describe('forkConversation', () => { ); }); + test('detaches subagent lineage when forking a child conversation', async () => { + getConvo.mockResolvedValue({ + ...mockConversation, + subagentThread: { + rootConversationId: 'root-conversation', + parentConversationId: 'parent-conversation', + parentToolCallId: 'parent-tool-call', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + }); + + await forkConversation({ + originalConvoId: 'abc123', + targetMessageId: '3', + requestUserId: 'user1', + option: ForkOptions.DIRECT_PATH, + }); + + expect(bulkSaveConvos.mock.calls[0][0][0]).not.toHaveProperty('subagentThread'); + }); + test('should fork conversation with branches', async () => { const result = await forkConversation({ originalConvoId: 'abc123', @@ -311,6 +334,27 @@ describe('duplicateConversation', () => { // bulkIncrementTagCounts will be called with empty array expect(bulkIncrementTagCounts).toHaveBeenCalledWith('user1', []); }); + + test('detaches subagent lineage when duplicating a child conversation', async () => { + getConvo.mockResolvedValue({ + ...mockConversation, + subagentThread: { + rootConversationId: 'root-conversation', + parentConversationId: 'parent-conversation', + parentToolCallId: 'parent-tool-call', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + }); + + await duplicateConversation({ + userId: 'user1', + conversationId: 'abc123', + }); + + expect(bulkSaveConvos.mock.calls[0][0][0]).not.toHaveProperty('subagentThread'); + }); }); describe('forkSharedConversation', () => { diff --git a/api/server/utils/import/importBatchBuilder.js b/api/server/utils/import/importBatchBuilder.js index b1856737cd..ea82923463 100644 --- a/api/server/utils/import/importBatchBuilder.js +++ b/api/server/utils/import/importBatchBuilder.js @@ -127,6 +127,7 @@ class ImportBatchBuilder { ...this.getRetentionFields(), }; convo._id && delete convo._id; + delete convo.subagentThread; this.conversations.push(convo); return { conversation: convo, messages: this.messages }; diff --git a/client/src/components/Chat/ChatView.tsx b/client/src/components/Chat/ChatView.tsx index 3d97c4fee8..c6316d2ae1 100644 --- a/client/src/components/Chat/ChatView.tsx +++ b/client/src/components/Chat/ChatView.tsx @@ -75,6 +75,12 @@ function ChatView({ index = 0, project }: { index?: number; project?: TChatProje const chatHelpers = useChatHelpers(index, conversationId); const addedChatHelpers = useAddedResponse(); + const activeConversation = + chatHelpers.conversation?.conversationId === conversationId + ? chatHelpers.conversation + : undefined; + const activeSubagentThread = activeConversation?.subagentThread; + useAdaptiveSSE(rootSubmission, chatHelpers, false, index); // Auto-resume if navigating back to conversation with active job. @@ -116,6 +122,11 @@ function ChatView({ index = 0, project }: { index?: number; project?: TChatProje : undefined; const pageHeading = isLandingPage || !conversationTitle ? localize('com_ui_new_chat') : conversationTitle; + const parentConversationId = activeSubagentThread?.parentConversationId; + /** Durable child threads are an execution record owned by their parent agent. + * Human continuation is a separate future fork/promotion flow, never an + * in-place mutation of this canonical child transcript. */ + const isSubagentThreadReadOnly = activeSubagentThread != null; return ( @@ -124,7 +135,10 @@ function ChatView({ index = 0, project }: { index?: number; project?: TChatProje

{pageHeading}

-
+
<>
{isLandingPage && } - + {isSubagentThreadReadOnly ? ( +
+ {localize('com_ui_subagent_thread_read_only')} +
+ ) : ( + + )} {!isLandingPage &&
}
diff --git a/client/src/components/Chat/Header.tsx b/client/src/components/Chat/Header.tsx index 0d936c89df..454b872e0f 100644 --- a/client/src/components/Chat/Header.tsx +++ b/client/src/components/Chat/Header.tsx @@ -5,6 +5,7 @@ import { OpenSidebar, PresetsMenu, NewChat, HeaderMenu } from './Menus'; import ModelSelector from './Menus/Endpoints/ModelSelector'; import { useGetStartupConfig } from '~/data-provider'; import ExportAndShareMenu from './ExportAndShareMenu'; +import SubagentThreadLink from './SubagentThreadLink'; import BookmarkMenu from './Menus/BookmarkMenu'; import { TemporaryChat } from './TemporaryChat'; import AddMultiConvo from './AddMultiConvo'; @@ -20,7 +21,13 @@ const defaultInterface = getConfigDefaults().interface; * reordering. Branching is CSS-only — `useMediaQuery` resolves after paint and * would pop the row a frame late on every mount. */ -function Header() { +function Header({ + parentConversationId, + readOnly = false, +}: { + parentConversationId?: string; + readOnly?: boolean; +}) { const { data: startupConfig } = useGetStartupConfig(); const navVisible = useRecoilValue(store.sidebarExpanded); @@ -59,8 +66,15 @@ function Header() { hiddenBehindNav, )} > - - {interfaceConfig.presets === true && interfaceConfig.modelSelect === true && ( + {parentConversationId != null && ( + + )} + {!readOnly && } + {!readOnly && interfaceConfig.presets === true && interfaceConfig.modelSelect === true && ( )} {hasAccessToBookmarks === true && ( diff --git a/client/src/components/Chat/Messages/Content/Parts/SubagentCall.tsx b/client/src/components/Chat/Messages/Content/Parts/SubagentCall.tsx index 8dadf8a502..1c8674c210 100644 --- a/client/src/components/Chat/Messages/Content/Parts/SubagentCall.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/SubagentCall.tsx @@ -21,11 +21,13 @@ import type { SubagentTickerLine } from '~/utils/subagentContent'; import ToolCallGroup from '~/components/Chat/Messages/Content/ToolCallGroup'; import MarkdownLite from '~/components/Chat/Messages/Content/MarkdownLite'; import ToolApproval from '~/components/Chat/Messages/Content/ToolApproval'; +import SubagentThreadLink from '~/components/Chat/SubagentThreadLink'; import { cn, groupSequentialToolCalls, parseToolName } from '~/utils'; import Container from '~/components/Chat/Messages/Content/Container'; import ToolCall from '~/components/Chat/Messages/Content/ToolCall'; import { MessageContext } from '~/Providers/MessageContext'; import MessageIcon from '~/components/Share/MessageIcon'; +import { parseSubagentBackgroundHandle } from './handle'; import { subagentProgressByToolCallId } from '~/store'; import { useAgentsMapContext } from '~/Providers'; import { useMCPServerNames } from '~/hooks/MCP'; @@ -186,6 +188,10 @@ export default function SubagentCall({ const agentsMap = useAgentsMapContext(); const [open, setOpen] = useState(false); const [promptExpanded, setPromptExpanded] = useState(false); + const backgroundHandle = useMemo( + () => parseSubagentBackgroundHandle(output, args), + [output, args], + ); const subagentType = progress?.subagentType ?? extractSubagentType(args); const isSelfSpawn = subagentType === 'self'; @@ -457,7 +463,7 @@ export default function SubagentCall({ ); } - if (output) { + if (output && backgroundHandle == null) { /** Fallback: no aggregated content parts but the backend * wrote a final tool_call output. Happens for older * subagent runs recorded before the event forwarder @@ -567,11 +573,19 @@ export default function SubagentCall({ )} >
- - {isSelfSpawn - ? localize('com_ui_subagent_dialog_title_self') - : localize('com_ui_subagent_dialog_title', { 0: subagentType })} - +
+ + {isSelfSpawn + ? localize('com_ui_subagent_dialog_title_self') + : localize('com_ui_subagent_dialog_title', { 0: subagentType })} + + {backgroundHandle != null && ( + + )} +
{localize('com_ui_subagent_dialog_description')} diff --git a/client/src/components/Chat/Messages/Content/Parts/__tests__/SubagentCall.test.tsx b/client/src/components/Chat/Messages/Content/Parts/__tests__/SubagentCall.test.tsx index 67e82c0c2b..952385af6a 100644 --- a/client/src/components/Chat/Messages/Content/Parts/__tests__/SubagentCall.test.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/__tests__/SubagentCall.test.tsx @@ -1,4 +1,5 @@ import React from 'react'; +import { MemoryRouter } from 'react-router-dom'; import { RecoilRoot, useRecoilCallback } from 'recoil'; import { render, screen, act, fireEvent, waitFor, within } from '@testing-library/react'; import type { SubagentUpdateEvent } from 'librechat-data-provider'; @@ -17,6 +18,12 @@ import { import SubagentCall, { SUBAGENT_TICKER_THROTTLE_MS } from '../SubagentCall'; import { subagentProgressByToolCallId } from '~/store/subagents'; +const mockNavigateToConvo = jest.fn(); + +jest.mock('librechat-data-provider/react-query', () => ({ + useGetConversationByIdQuery: (id: string) => ({ data: { conversationId: id } }), +})); + jest.mock('~/hooks', () => ({ useLocalize: () => @@ -34,6 +41,7 @@ jest.mock('~/hooks', () => ({ com_ui_subagent_dialog_description: 'Isolated child run.', com_ui_subagent_no_result_yet: 'No result yet.', com_ui_subagent_empty_result: 'No text.', + com_ui_subagent_open_thread: 'Open child chat', com_ui_collapse: 'Collapse', com_ui_expand: 'Expand', com_ui_subagent_ticker_writing: 'Writing', @@ -46,6 +54,7 @@ jest.mock('~/hooks', () => ({ }; return translations[key] ?? key; }, + useNavigateToConvo: () => ({ navigateToConvo: mockNavigateToConvo }), })); /** Stub the leaf content-part renderers — the tests only need to confirm @@ -637,6 +646,62 @@ describe('SubagentCall — dialog content', () => { rerender({null}); }); + it('links only an exact host-issued detached result to its durable child chat', () => { + const output = JSON.stringify({ + background_task_id: 'task-1', + subagent_thread_id: 'child-thread-1', + tool: 'subagent', + subagent_type: 'self', + status: 'running', + message: + 'Started subagent "self" background task. Poll the host background-task tool with background_task_id "task-1".', + }); + render( + + + + + , + ); + + openSubagentDialog(); + fireEvent.click(screen.getByRole('button', { name: 'Open child chat' })); + expect(mockNavigateToConvo).toHaveBeenCalledWith({ conversationId: 'child-thread-1' }); + expect(screen.queryByText(output)).not.toBeInTheDocument(); + }); + + it('does not turn model-authored foreground output into a child-chat link', () => { + const output = JSON.stringify({ + background_task_id: 'task-1', + subagent_thread_id: 'child-thread-1', + tool: 'subagent', + subagent_type: 'self', + status: 'running', + message: 'background_task_id task-1', + }); + render( + + + , + ); + + openSubagentDialog(); + expect(screen.queryByRole('link', { name: 'Open child chat' })).not.toBeInTheDocument(); + expect(screen.getByText(output)).toBeInTheDocument(); + }); + it('renders persistedContent parts when no live events are available (page-refresh flow)', () => { /** * After a refresh the Recoil atom is empty — the child's history has diff --git a/client/src/components/Chat/Messages/Content/Parts/__tests__/handle.test.ts b/client/src/components/Chat/Messages/Content/Parts/__tests__/handle.test.ts index 6e8cff7a31..366ab7a2ad 100644 --- a/client/src/components/Chat/Messages/Content/Parts/__tests__/handle.test.ts +++ b/client/src/components/Chat/Messages/Content/Parts/__tests__/handle.test.ts @@ -1,4 +1,8 @@ -import { parseBackgroundHandle, splitBackgroundAttachments } from '../handle'; +import { + parseBackgroundHandle, + parseSubagentBackgroundHandle, + splitBackgroundAttachments, +} from '../handle'; describe('parseBackgroundHandle', () => { const handle = JSON.stringify({ @@ -80,3 +84,43 @@ describe('parseBackgroundHandle', () => { expect(parseBackgroundHandle(big)).toBeNull(); }); }); + +describe('parseSubagentBackgroundHandle', () => { + const handle = JSON.stringify({ + background_task_id: 'task-1', + subagent_thread_id: 'thread-1', + tool: 'subagent', + subagent_type: 'researcher', + status: 'running', + message: + 'Started subagent "researcher" background task. Poll the host background-task tool with background_task_id "task-1".', + }); + + it('parses the exact host-issued detached-subagent handle', () => { + expect(parseSubagentBackgroundHandle(handle, { run_in_background: true })).toEqual( + expect.objectContaining({ + background_task_id: 'task-1', + subagent_thread_id: 'thread-1', + tool: 'subagent', + }), + ); + }); + + it.each([ + [undefined, { run_in_background: true }], + ['', { run_in_background: true }], + ['{"subagent_thread_id":"thread-1"}', { run_in_background: true }], + [ + '{"background_task_id":"task-1","subagent_thread_id":"thread-1","tool":"execute_code","subagent_type":"researcher","status":"running","message":"background_task_id task-1"}', + { run_in_background: true }, + ], + [ + '{"background_task_id":"task-1","subagent_thread_id":"thread-1","tool":"subagent","subagent_type":"researcher","status":"running","message":"ordinary result","extra":true}', + { run_in_background: true }, + ], + [handle, { run_in_background: false }], + [handle, undefined], + ])('rejects non-handles and spoofable near-matches: %s', (output, args) => { + expect(parseSubagentBackgroundHandle(output, args)).toBeNull(); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/Parts/handle.ts b/client/src/components/Chat/Messages/Content/Parts/handle.ts index 27fb591f42..8c89cf575b 100644 --- a/client/src/components/Chat/Messages/Content/Parts/handle.ts +++ b/client/src/components/Chat/Messages/Content/Parts/handle.ts @@ -53,6 +53,11 @@ export interface BackgroundHandle { message: string; } +export interface SubagentBackgroundHandle extends BackgroundHandle { + subagent_thread_id: string; + subagent_type: string; +} + const HANDLE_KEYS: ReadonlyArray = [ 'background_task_id', 'tool', @@ -91,3 +96,66 @@ export function parseBackgroundHandle(output?: string): BackgroundHandle | null } return null; } + +const SUBAGENT_HANDLE_KEYS: ReadonlyArray = [ + 'background_task_id', + 'subagent_thread_id', + 'tool', + 'subagent_type', + 'status', + 'message', +]; + +/** + * Extracts the host-issued durable child-thread handle from a detached + * `subagent` result. The strict shape check is intentional: model-authored + * prose or an ordinary child result that happens to mention a thread id must + * never become a navigation target. + */ +export function parseSubagentBackgroundHandle( + output?: string | null, + args?: string | Record, +): SubagentBackgroundHandle | null { + let parsedArgs: Record | undefined; + if (typeof args === 'string') { + try { + const candidate = JSON.parse(args) as unknown; + parsedArgs = + candidate != null && typeof candidate === 'object' && !Array.isArray(candidate) + ? (candidate as Record) + : undefined; + } catch { + return null; + } + } else { + parsedArgs = args; + } + if (parsedArgs?.run_in_background !== true || !output || output.length > 2_000) { + return null; + } + const trimmed = output.trim(); + if ( + !trimmed.startsWith('{') || + !trimmed.includes('"background_task_id"') || + !trimmed.includes('"subagent_thread_id"') + ) { + return null; + } + try { + const parsed = JSON.parse(trimmed) as Partial | null; + if ( + parsed != null && + Object.keys(parsed).length === SUBAGENT_HANDLE_KEYS.length && + SUBAGENT_HANDLE_KEYS.every( + (key) => typeof parsed[key] === 'string' && (parsed[key] as string).trim() !== '', + ) && + parsed.tool === 'subagent' && + (parsed.message as string).includes('background_task_id') + ) { + return parsed as SubagentBackgroundHandle; + } + } catch { + return null; + } + return null; +} diff --git a/client/src/components/Chat/Messages/HoverButtons.tsx b/client/src/components/Chat/Messages/HoverButtons.tsx index ce3e629255..b54c890043 100644 --- a/client/src/components/Chat/Messages/HoverButtons.tsx +++ b/client/src/components/Chat/Messages/HoverButtons.tsx @@ -174,6 +174,7 @@ const HoverButtons = ({ } const { isCreatedByUser, error } = message; + const isSubagentThreadReadOnly = conversation.subagentThread != null; const onEdit = () => { if (isEditing) { @@ -227,7 +228,7 @@ const HoverButtons = ({ )} {/* Edit Button */} - {isEditableEndpoint && !hideEditButton && ( + {!isSubagentThreadReadOnly && isEditableEndpoint && !hideEditButton && ( e && handleContinue(e)} title={localize('com_ui_continue')} diff --git a/client/src/components/Chat/Messages/__tests__/HoverButtons.spec.tsx b/client/src/components/Chat/Messages/__tests__/HoverButtons.spec.tsx index a11cfed65f..3657e0aed5 100644 --- a/client/src/components/Chat/Messages/__tests__/HoverButtons.spec.tsx +++ b/client/src/components/Chat/Messages/__tests__/HoverButtons.spec.tsx @@ -30,12 +30,14 @@ const userMessage = { function renderHoverButtons({ isSubmitting, message = userMessage, + conversation: targetConversation = conversation, isLast = false, latestMessageId = 'assistant-1', getCanCopy = () => hasCopyableText({ text: message.text, content: message.content }), }: { isSubmitting: boolean; message?: TMessage; + conversation?: TConversation; isLast?: boolean; latestMessageId?: string; getCanCopy?: () => boolean; @@ -55,7 +57,7 @@ function renderHoverButtons({ isLast={isLast} isEditing={false} message={message} - conversation={conversation} + conversation={targetConversation} isSubmitting={isSubmitting} enterEdit={jest.fn()} regenerate={jest.fn()} @@ -173,4 +175,40 @@ describe('HoverButtons edit affordance', () => { expect(screen.queryByTestId('copy-response-button')).toBeNull(); expect(getCanCopy).not.toHaveBeenCalled(); }); + + it('keeps child-thread history readable without model-turn controls', () => { + const assistantMessage = { + ...userMessage, + messageId: 'assistant-child', + isCreatedByUser: false, + text: 'Completed child result', + } as TMessage; + const childConversation = { + ...conversation, + conversationId: 'child-thread', + subagentThread: { + rootConversationId: 'parent-thread', + parentConversationId: 'parent-thread', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool-call', + parentAgentId: 'parent-agent', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + } as TConversation; + + const container = renderHoverButtons({ + isSubmitting: false, + message: assistantMessage, + conversation: childConversation, + isLast: true, + latestMessageId: assistantMessage.messageId, + }); + + expect(screen.getByTestId('copy-response-button')).toBeEnabled(); + expect(container.querySelector(`#edit-${assistantMessage.messageId}`)).toBeNull(); + expect(screen.queryByTestId('regenerate-generation-button')).toBeNull(); + expect(screen.queryByTestId('continue-generation-button')).toBeNull(); + }); }); diff --git a/client/src/components/Chat/SubagentThreadLink.tsx b/client/src/components/Chat/SubagentThreadLink.tsx new file mode 100644 index 0000000000..22db15754e --- /dev/null +++ b/client/src/components/Chat/SubagentThreadLink.tsx @@ -0,0 +1,64 @@ +import { useMemo } from 'react'; +import { Button } from '@librechat/client'; +import { ChevronLeft, ChevronRight } from 'lucide-react'; +import { useGetConversationByIdQuery } from 'librechat-data-provider/react-query'; +import { useLocalize, useNavigateToConvo } from '~/hooks'; +import { cn } from '~/utils'; + +const CHILD_THREAD_POLL_WINDOW_MS = 60_000; + +export default function SubagentThreadLink({ + threadId, + relation, + className, + labelClassName, +}: { + threadId: string; + relation: 'parent' | 'child'; + className?: string; + labelClassName?: string; +}) { + const localize = useLocalize(); + const { navigateToConvo } = useNavigateToConvo(); + const normalizedThreadId = threadId.trim(); + const isParent = relation === 'parent'; + const childPoll = useMemo( + () => ({ threadId: normalizedThreadId, deadline: Date.now() + CHILD_THREAD_POLL_WINDOW_MS }), + [normalizedThreadId], + ); + const { data: targetConversation } = useGetConversationByIdQuery(normalizedThreadId, { + enabled: normalizedThreadId !== '', + retry: false, + refetchInterval: (conversation) => + !isParent && + conversation == null && + normalizedThreadId === childPoll.threadId && + Date.now() < childPoll.deadline + ? 1500 + : false, + }); + if (normalizedThreadId === '' || targetConversation == null) { + return null; + } + + const label = localize( + isParent ? 'com_ui_subagent_back_to_parent' : 'com_ui_subagent_open_thread', + ); + const Icon = isParent ? ChevronLeft : ChevronRight; + + return ( + + ); +} diff --git a/client/src/components/Chat/__tests__/ChatView.subagent.spec.tsx b/client/src/components/Chat/__tests__/ChatView.subagent.spec.tsx new file mode 100644 index 0000000000..7eae1f93c4 --- /dev/null +++ b/client/src/components/Chat/__tests__/ChatView.subagent.spec.tsx @@ -0,0 +1,121 @@ +import React from 'react'; +import { render, screen } from '@testing-library/react'; +import ChatView from '../ChatView'; + +let mockConversation: Record | null; + +jest.mock('react-router-dom', () => ({ + useParams: () => ({ conversationId: 'child-thread' }), +})); + +jest.mock('recoil', () => ({ + useRecoilValue: () => false, +})); + +jest.mock('react-hook-form', () => ({ + useForm: () => ({}), +})); + +jest.mock('@librechat/client', () => ({ + Spinner: () =>
, +})); + +jest.mock('librechat-data-provider', () => ({ + Constants: { NEW_CONVO: 'new', SEARCH: 'search' }, + buildTree: ({ messages }: { messages: unknown[] }) => messages, +})); + +jest.mock('~/hooks', () => ({ + useAddedResponse: () => ({}), + useResumeOnLoad: jest.fn(), + useAdaptiveSSE: jest.fn(), + useQueueDrain: jest.fn(), + useLocalize: () => (key: string) => key, + useChatHelpers: () => ({ + conversation: mockConversation, + getMessages: jest.fn(), + ask: jest.fn(), + }), +})); + +jest.mock('~/Providers', () => { + const Passthrough = ({ children }: { children: React.ReactNode }) => <>{children}; + const Context = { Provider: Passthrough }; + return { + ChatContext: Context, + AddedChatContext: Context, + ChatFormProvider: Passthrough, + useFileMapContext: () => new Map(), + }; +}); + +jest.mock('~/data-provider', () => ({ + useGetMessagesByConvoId: () => ({ + data: [{ messageId: 'message-1' }], + isLoading: false, + isFetching: false, + }), +})); + +jest.mock('../Input/ConversationStarters', () => () => null); +jest.mock('../Messages/MessagesView', () => () =>
); +jest.mock('../Presentation', () => ({ children }: { children: React.ReactNode }) => ( + <>{children} +)); +jest.mock('../Input/ChatForm', () => () =>
); +jest.mock('../Landing', () => () =>
); +jest.mock('../Footer', () => () =>
); +jest.mock('../Header', () => ({ readOnly }: { readOnly?: boolean }) => ( +
+)); + +jest.mock('~/utils', () => ({ + cn: (...values: Array) => values.filter(Boolean).join(' '), +})); + +jest.mock('~/store', () => ({ + __esModule: true, + default: { + submissionByIndex: jest.fn(), + isSubmittingFamily: jest.fn(), + centerFormOnLanding: {}, + }, +})); + +describe('ChatView child-thread execution identity', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('renders a child without a standalone identity as view-only', () => { + mockConversation = { + conversationId: 'child-thread', + title: 'Graph child', + subagentThread: { + parentConversationId: 'parent-thread', + }, + }; + + render(); + + expect(screen.queryByTestId('chat-form')).not.toBeInTheDocument(); + expect(screen.getByRole('note')).toHaveTextContent('com_ui_subagent_thread_read_only'); + expect(screen.getByTestId('header')).toHaveAttribute('data-read-only', 'true'); + }); + + it('keeps a saved-agent child view-only after it settles', () => { + mockConversation = { + conversationId: 'child-thread', + title: 'Saved agent child', + subagentThread: { + parentConversationId: 'parent-thread', + }, + }; + + render(); + + expect(screen.queryByTestId('chat-form')).not.toBeInTheDocument(); + expect(screen.getByRole('note')).toHaveTextContent('com_ui_subagent_thread_read_only'); + expect(screen.getByTestId('header')).toHaveAttribute('data-read-only', 'true'); + }); +}); diff --git a/client/src/components/Chat/__tests__/SubagentThreadLink.test.tsx b/client/src/components/Chat/__tests__/SubagentThreadLink.test.tsx new file mode 100644 index 0000000000..6e04154dd5 --- /dev/null +++ b/client/src/components/Chat/__tests__/SubagentThreadLink.test.tsx @@ -0,0 +1,104 @@ +import React from 'react'; +import { fireEvent, render, screen } from '@testing-library/react'; +import SubagentThreadLink from '../SubagentThreadLink'; + +const mockUseGetConversationByIdQuery = jest.fn< + { data: { conversationId: string } | undefined }, + [string, Record?] +>(() => ({ data: undefined })); +const mockNavigateToConvo = jest.fn(); + +jest.mock('librechat-data-provider/react-query', () => ({ + useGetConversationByIdQuery: (id: string, config?: Record) => + mockUseGetConversationByIdQuery(id, config), +})); + +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string) => + key === 'com_ui_subagent_back_to_parent' ? 'Back to parent chat' : 'Open child chat', + useNavigateToConvo: () => ({ navigateToConvo: mockNavigateToConvo }), +})); + +jest.mock('lucide-react', () => ({ + ChevronLeft: () => , + ChevronRight: () => , +})); + +describe('SubagentThreadLink', () => { + const renderLink = (element: React.ReactElement) => render(element); + + beforeEach(() => { + mockUseGetConversationByIdQuery.mockReset(); + mockUseGetConversationByIdQuery.mockReturnValue({ data: undefined }); + mockNavigateToConvo.mockReset(); + }); + + it('loads a parent chat and navigates through the conversation state helper', () => { + const parent = { conversationId: 'parent-thread' }; + mockUseGetConversationByIdQuery.mockReturnValue({ data: parent }); + renderLink(); + + fireEvent.click(screen.getByRole('button', { name: 'Back to parent chat' })); + expect(mockNavigateToConvo).toHaveBeenCalledWith(parent); + expect(screen.getByTestId('left-icon')).toBeInTheDocument(); + expect(mockUseGetConversationByIdQuery).toHaveBeenCalledWith( + 'parent-thread', + expect.objectContaining({ enabled: true }), + ); + }); + + it('links a parent tool result only after the child conversation is durable', () => { + mockUseGetConversationByIdQuery.mockReturnValue({ + data: { conversationId: 'child/thread' }, + }); + renderLink(); + + fireEvent.click(screen.getByRole('button', { name: 'Open child chat' })); + expect(mockNavigateToConvo).toHaveBeenCalledWith({ conversationId: 'child/thread' }); + expect(screen.getByTestId('right-icon')).toBeInTheDocument(); + }); + + it('hides a provisional child link while polling for durable creation', () => { + const { container } = renderLink( + , + ); + + expect(container).toBeEmptyDOMElement(); + expect(mockUseGetConversationByIdQuery).toHaveBeenCalledWith( + 'provisional-child', + expect.objectContaining({ enabled: true, retry: false }), + ); + const config = mockUseGetConversationByIdQuery.mock.calls[0][1] as { + refetchInterval: (conversation: unknown) => number | false; + }; + expect(config.refetchInterval(undefined)).toBe(1500); + }); + + it('stops polling for a child that never became durable', () => { + const now = jest.spyOn(Date, 'now').mockReturnValue(10_000); + renderLink(); + const config = mockUseGetConversationByIdQuery.mock.calls[0][1] as { + refetchInterval: (conversation: unknown) => number | false; + }; + + now.mockReturnValue(70_000); + expect(config.refetchInterval(undefined)).toBe(false); + now.mockRestore(); + }); + + it('does not render an empty thread selector', () => { + const { container } = renderLink(); + expect(container).toBeEmptyDOMElement(); + }); + + it('passes the complete fetched child record into conversation navigation', () => { + const child = { conversationId: 'child-thread', title: 'Research child' }; + mockUseGetConversationByIdQuery.mockReturnValue({ + data: child, + }); + renderLink(); + + fireEvent.click(screen.getByRole('button', { name: 'Open child chat' })); + expect(mockNavigateToConvo).toHaveBeenCalledWith(child); + }); +}); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index ec08f6bf40..23f5c4740f 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -2098,6 +2098,7 @@ "com_ui_stop": "Stop", "com_ui_storage": "Storage", "com_ui_storage_filter_sort": "Filter and Sort by Storage", + "com_ui_subagent_back_to_parent": "Back to parent chat", "com_ui_subagent_cancelled": "Cancelled agent", "com_ui_subagent_complete": "Ran agent", "com_ui_subagent_dialog_description": "Isolated-context child run. Activity and final result below.", @@ -2106,6 +2107,8 @@ "com_ui_subagent_empty_result": "No text returned.", "com_ui_subagent_errored": "Agent errored", "com_ui_subagent_no_result_yet": "Still running — no final result yet.", + "com_ui_subagent_open_thread": "Open child chat", + "com_ui_subagent_thread_read_only": "This child thread is view-only here. Its parent agent owns this execution and can continue it with the saved thread history.", "com_ui_subagent_running": "Running agent", "com_ui_subagent_scroll_to_bottom": "Scroll to latest", "com_ui_subagent_ticker_error": "Error", diff --git a/package-lock.json b/package-lock.json index 4a95464f4e..15d6c5225a 100644 --- a/package-lock.json +++ b/package-lock.json @@ -63,7 +63,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.3", + "@librechat/agents": "^3.6.6", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", @@ -10628,9 +10628,9 @@ } }, "node_modules/@librechat/agents": { - "version": "3.6.3", - "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.3.tgz", - "integrity": "sha512-NibUZX9+spmjURqzswiB/MGMqam/YoJswwiQ+5ugmUgyta9ES6IrQ1/9LYZgk2g8IWSNGDXflPy21ronsxEl7Q==", + "version": "3.6.6", + "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.6.tgz", + "integrity": "sha512-YJqZ4Dsw/+dIu/Kbp3oMAc36zFXIrecHCcxYHdlFZA7WNnBLlBk2RwDSvf2WZUP0KfPACdx5K5LZOhTG4s+7ZQ==", "license": "MIT", "dependencies": { "@anthropic-ai/sdk": "^0.115.0", @@ -42875,7 +42875,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.3", + "@librechat/agents": "^3.6.6", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/package.json b/packages/api/package.json index bcdd8bbdcb..4f9e0319fa 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -113,7 +113,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "^4.3.3", - "@librechat/agents": "^3.6.3", + "@librechat/agents": "^3.6.6", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/src/agents/__tests__/run-codeTools.test.ts b/packages/api/src/agents/__tests__/run-codeTools.test.ts index f977e465dd..8c170fc035 100644 --- a/packages/api/src/agents/__tests__/run-codeTools.test.ts +++ b/packages/api/src/agents/__tests__/run-codeTools.test.ts @@ -1,3 +1,5 @@ +import type { SubagentTaskConfig } from '@librechat/agents'; +import { CHECK_BACKGROUND_TASK_NAME } from '~/agents/background'; import { createRun } from '~/agents/run'; /** @@ -56,7 +58,7 @@ jest.mock('~/agents/checkpointer', () => ({ getAgentCheckpointer: jest.fn().mockResolvedValue({}), })); -import { Run } from '@librechat/agents'; +import { InMemorySubagentTaskStore, Run } from '@librechat/agents'; function makeAgent(overrides?: Record) { return { @@ -72,12 +74,16 @@ function makeAgent(overrides?: Record) { }; } -async function captureRunConfig(agent = makeAgent()): Promise> { +async function captureRunConfig( + agent = makeAgent(), + subagentTasks?: SubagentTaskConfig, +): Promise> { await createRun({ agents: [agent] as never, signal: new AbortController().signal, streaming: true, streamUsage: true, + subagentTasks, }); const createMock = Run.create as jest.Mock; expect(createMock).toHaveBeenCalledTimes(1); @@ -113,4 +119,37 @@ describe('createRun code-tool eager/session wiring', () => { .agents; expect(agentInput.codeSessionKey).toBe(codeSessionKey); }); + + it('registers detached task controls only on a spawn-capable parent', async () => { + const subagentTasks: SubagentTaskConfig = { + store: new InMemorySubagentTaskStore(), + scopeId: 'owner:parent-thread', + }; + const runConfig = await captureRunConfig( + makeAgent({ + subagents: { enabled: true, allowSelf: true }, + toolDefinitions: [], + toolRegistry: new Map(), + }), + subagentTasks, + ); + const [agentInput] = (runConfig.graphConfig as { agents: Array> }) + .agents; + const parentDefinitions = agentInput.toolDefinitions as Array<{ name: string }>; + const [selfConfig] = agentInput.subagentConfigs as Array<{ + agentInputs?: { + toolDefinitions?: Array<{ name: string }>; + toolRegistry?: Map; + }; + }>; + + expect(runConfig.subagentTasks).toBe(subagentTasks); + expect(parentDefinitions.map((definition) => definition.name)).toContain( + CHECK_BACKGROUND_TASK_NAME, + ); + expect( + selfConfig.agentInputs?.toolDefinitions?.map((definition) => definition.name), + ).not.toContain(CHECK_BACKGROUND_TASK_NAME); + expect(selfConfig.agentInputs?.toolRegistry?.has(CHECK_BACKGROUND_TASK_NAME)).toBe(false); + }); }); diff --git a/packages/api/src/agents/__tests__/run-summarization.test.ts b/packages/api/src/agents/__tests__/run-summarization.test.ts index 7345d4e01b..31797efea4 100644 --- a/packages/api/src/agents/__tests__/run-summarization.test.ts +++ b/packages/api/src/agents/__tests__/run-summarization.test.ts @@ -7,6 +7,7 @@ import { } from 'librechat-data-provider'; import type { SummarizationConfig, TEndpoint } from 'librechat-data-provider'; import type { BaseMessage } from '@langchain/core/messages'; +import type { SubagentTaskConfig } from '@librechat/agents'; import type { AppConfig } from '@librechat/data-schemas'; import { createRun } from '~/agents/run'; @@ -83,7 +84,7 @@ jest.mock('~/agents/checkpointer', () => ({ getAgentCheckpointer: jest.fn().mockResolvedValue({}), })); -import { Run, buildChildInputs } from '@librechat/agents'; +import { Run, buildChildInputs, InMemorySubagentTaskStore } from '@librechat/agents'; /** Minimal RunAgent factory */ function makeAgent( @@ -162,6 +163,7 @@ async function callAndCapture( appConfig?: AppConfig; messages?: BaseMessage[]; discoveredToolNames?: string[]; + subagentTasks?: SubagentTaskConfig; } = {}, ) { const agents = opts.agents ?? [makeAgent()]; @@ -175,6 +177,7 @@ async function callAndCapture( appConfig: opts.appConfig, messages: opts.messages, discoveredToolNames: opts.discoveredToolNames, + subagentTasks: opts.subagentTasks, streaming: true, streamUsage: true, }); @@ -1010,6 +1013,20 @@ describe('subagentConfigs', () => { expect(agents[0].subagentConfigs).toBeUndefined(); }); + it('keeps the poll tool available for existing tasks after spawning is disabled', async () => { + const agents = await callAndCapture({ + subagentTasks: { + store: new InMemorySubagentTaskStore(), + scopeId: 'existing-task-scope', + }, + }); + + expect(agents[0].subagentConfigs).toBeUndefined(); + expect(agents[0].toolDefinitions).toEqual( + expect.arrayContaining([expect.objectContaining({ name: 'check_background_task' })]), + ); + }); + it('adds self-spawn when enabled and allowSelf defaults to true', async () => { const agents = await callAndCapture({ agents: [makeAgent({ subagents: { enabled: true } })], diff --git a/packages/api/src/agents/background.spec.ts b/packages/api/src/agents/background.spec.ts index 14c773e6d7..4c604c8b98 100644 --- a/packages/api/src/agents/background.spec.ts +++ b/packages/api/src/agents/background.spec.ts @@ -1,5 +1,6 @@ import { logger } from '@librechat/data-schemas'; -import type { LCTool, LCToolRegistry } from '@librechat/agents'; +import { InMemorySubagentTaskStore } from '@librechat/agents'; +import type { LCTool, LCToolRegistry, SubagentTaskConfig } from '@librechat/agents'; import { isBackgroundEligibleToolName, isBackgroundRequested, @@ -28,6 +29,20 @@ const mcpDef = (name: string): LCTool => parameters: { type: 'object', properties: { q: { type: 'string' } }, required: ['q'] }, }) as unknown as LCTool; +async function waitForSubagentTaskToSettle( + store: InMemorySubagentTaskStore, + scopeId: string, + taskId: string, +): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (store.get(scopeId, taskId)?.status !== 'running') { + return; + } + await new Promise((resolve) => setTimeout(resolve, 1)); + } + throw new Error('Timed out waiting for the detached subagent task.'); +} + describe('isBackgroundEligibleToolName', () => { it('excludes direct-path, host-special, and machinery tools', () => { for (const name of [ @@ -1138,6 +1153,120 @@ describe('runCheckBackgroundTask (singleton)', () => { }), ); }); + + it('polls and one-shot claims a detached subagent result', async () => { + const store = new InMemorySubagentTaskStore(); + const subagentTasks: SubagentTaskConfig = { store, scopeId: 'owner:parent-thread' }; + const started = store.start({ + scopeId: subagentTasks.scopeId, + idempotencyKey: 'parent-run:parent-agent:call-1', + parentRunId: 'parent-run', + parentAgentId: 'parent-agent', + parentToolCallId: 'call-1', + input: 'Research this.', + subagentKind: 'agent', + subagentType: 'researcher', + run: async () => ({ content: 'finished research' }), + }); + if (!started.accepted) { + throw new Error('Expected subagent task to start.'); + } + await waitForSubagentTaskToSettle(store, subagentTasks.scopeId, started.task.taskId); + + const first = JSON.parse( + runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { background_task_id: started.task.taskId }, + subagentTasks, + }), + ); + expect(first).toEqual( + expect.objectContaining({ + background_task_id: started.task.taskId, + subagent_thread_id: started.task.threadId, + tool: 'subagent', + status: 'completed', + result: 'finished research', + }), + ); + + const second = JSON.parse( + runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { background_task_id: started.task.taskId }, + subagentTasks, + }), + ); + expect(second).toEqual(expect.objectContaining({ status: 'claimed', result_claimed: true })); + expect(second.result).toBeUndefined(); + }); + + it('routes parent control actions only to detached subagent tasks', async () => { + const store = new InMemorySubagentTaskStore(); + const subagentTasks: SubagentTaskConfig = { store, scopeId: 'owner:parent-thread' }; + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start({ + scopeId: subagentTasks.scopeId, + idempotencyKey: 'parent-run:parent-agent:call-2', + parentRunId: 'parent-run', + parentAgentId: 'parent-agent', + parentToolCallId: 'call-2', + input: 'Keep working.', + subagentKind: 'agent', + subagentType: 'researcher', + run: async () => result, + }); + if (!started.accepted) { + throw new Error('Expected subagent task to start.'); + } + await Promise.resolve(); + + const queued = JSON.parse( + runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { + background_task_id: started.task.taskId, + action: 'queue', + message: 'Also verify the source.', + }, + subagentTasks, + }), + ); + expect(queued).toEqual( + expect.objectContaining({ status: 'accepted', control_id: expect.any(String) }), + ); + + const cancelledMessage = JSON.parse( + runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { + background_task_id: started.task.taskId, + action: 'cancel_message', + control_id: queued.control_id, + }, + subagentTasks, + }), + ); + expect(cancelledMessage.status).toBe('accepted'); + + const cancelledTask = JSON.parse( + runCheckBackgroundTask({ + userId: 'owner', + conversationId: 'parent-thread', + args: { background_task_id: started.task.taskId, action: 'cancel' }, + subagentTasks, + }), + ); + expect(cancelledTask.status).toBe('cancelled'); + finish({ content: 'late result' }); + }); }); describe('stripBackgroundFromToolRegistry', () => { diff --git a/packages/api/src/agents/background.ts b/packages/api/src/agents/background.ts index 44c8908c18..0f06c2ad6d 100644 --- a/packages/api/src/agents/background.ts +++ b/packages/api/src/agents/background.ts @@ -34,7 +34,16 @@ import { randomUUID } from 'node:crypto'; import { logger } from '@librechat/data-schemas'; import { Constants as AgentConstants } from '@librechat/agents'; import { Tools, Constants, imageGenTools } from 'librechat-data-provider'; -import type { LCTool, LCToolRegistry, JsonSchemaType } from '@librechat/agents'; +import type { + LCTool, + LCToolRegistry, + JsonSchemaType, + SubagentTaskClaim, + SubagentTaskConfig, + SubagentTaskSnapshot, + SubagentTaskControlCommand, + SubagentTaskControlResult, +} from '@librechat/agents'; import type { AgentToolOptions } from 'librechat-data-provider'; import type { CapabilityToolNames } from './selection'; import { @@ -286,9 +295,9 @@ export function stripBackgroundFromToolRegistry( return next; } -const CHECK_BACKGROUND_TASK_DESCRIPTION = `Check the status and retrieve the result of tool calls previously dispatched in the background (with run_in_background: true). +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 conversation. A task is only finished when its status is "completed" or "error" — never assume completion without polling. Results are not pushed to you; you must call this tool to collect them. Background tasks persist on this server across turns, so you can collect a result in a later turn; they do not survive a server restart.`; +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.`; const CHECK_BACKGROUND_TASK_PARAMETERS: JsonSchemaType = Object.freeze({ type: 'object', @@ -296,7 +305,20 @@ const CHECK_BACKGROUND_TASK_PARAMETERS: JsonSchemaType = Object.freeze 0 ? { pending_controls: task.pendingControls } : {}), + ...(task.error == null ? {} : { error: task.error }), + ...(options.controlId == null ? {} : { control_id: options.controlId }), + }; +} + +function serializeSubagentClaim(claim: SubagentTaskClaim): SerializedSubagentTask | undefined { + if (claim.status === 'not_found') { + return undefined; + } + if (claim.status === 'completed') { + return serializeSubagentSnapshot(claim.task, { includeResult: claim.result }); + } + if (claim.status === 'error' || claim.status === 'cancelled') { + return { + ...serializeSubagentSnapshot(claim.task, { status: claim.status }), + error: claim.error, + }; + } + return serializeSubagentSnapshot(claim.task, { status: claim.status }); +} + +function serializeSubagentControl( + result: SubagentTaskControlResult, +): SerializedSubagentTask | { status: string; message?: string } | undefined { + if (result.status === 'not_found') { + return undefined; + } + if (result.status === 'invalid') { + return { status: result.status, message: result.message }; + } + return serializeSubagentSnapshot(result.task, { + status: result.status, + ...(result.status === 'accepted' && result.controlId != null + ? { controlId: result.controlId } + : {}), + }); +} + +function buildSubagentControlCommand( + args: Record, + action: string, +): SubagentTaskControlCommand | undefined { + if (action === 'cancel') { + return { action: 'cancel' }; + } + if (action === 'cancel_message') { + return typeof args.control_id === 'string' + ? { action: 'cancel_message', controlId: args.control_id } + : undefined; + } + if (action === 'steer' || action === 'queue' || action === 'interrupt') { + return typeof args.message === 'string' ? { action, message: args.message } : undefined; + } + return undefined; +} + /** Executes a `check_background_task` call and returns the ToolMessage content. */ export function runCheckBackgroundTask(params: { userId: string; conversationId: string; args: unknown; + subagentTasks?: SubagentTaskConfig; }): string { const { userId, conversationId } = params; - const rawId = coerceArgsObject(params.args)?.background_task_id; + const args = coerceArgsObject(params.args) ?? {}; + const rawId = args.background_task_id; const taskId = typeof rawId === 'string' && rawId.trim() !== '' ? rawId.trim() : undefined; + const action = typeof args.action === 'string' && args.action !== '' ? args.action : 'poll'; if (taskId) { const task = backgroundTaskRegistry.get(userId, conversationId, taskId); - if (!task) { - return JSON.stringify({ - status: 'not_found', - background_task_id: taskId, - message: 'No background task with that id exists in this conversation.', - }); + if (task != null) { + if (action !== 'poll') { + return JSON.stringify({ + status: 'invalid', + background_task_id: taskId, + message: 'Control actions are supported only for subagent tasks.', + }); + } + return JSON.stringify(serializeTask(task, { includeResult: true })); } - return JSON.stringify(serializeTask(task, { includeResult: true })); + + 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); + } + } 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 controlled = serializeSubagentControl( + subagentTasks.store.control(subagentTasks.scopeId, taskId, command), + ); + if (controlled != null) { + return JSON.stringify(controlled); + } + } + } + + return JSON.stringify({ + status: 'not_found', + background_task_id: taskId, + message: 'No background task with that id exists in this thread.', + }); + } + + if (action !== 'poll') { + return JSON.stringify({ + status: 'invalid', + message: 'A background_task_id is required for control actions.', + }); } const tasks = backgroundTaskRegistry.list(userId, conversationId); - logger.debug(`[background] check_background_task listed ${tasks.length} task(s)`); + const subagentTasks = + params.subagentTasks?.store + .list(params.subagentTasks.scopeId) + .map((task) => serializeSubagentSnapshot(task)) ?? []; + logger.debug( + `[background] check_background_task listed ${tasks.length + subagentTasks.length} task(s)`, + ); return JSON.stringify({ - tasks: tasks.map((task) => serializeTask(task, { includeResult: false })), + tasks: [ + ...tasks.map((task) => serializeTask(task, { includeResult: false })), + ...subagentTasks, + ], }); } diff --git a/packages/api/src/agents/guard.spec.ts b/packages/api/src/agents/guard.spec.ts new file mode 100644 index 0000000000..0673ac1071 --- /dev/null +++ b/packages/api/src/agents/guard.spec.ts @@ -0,0 +1,149 @@ +import express from 'express'; +import request from 'supertest'; +import type { AllMethods, IConversation } from '@librechat/data-schemas'; +import { + CHILD_THREAD_READ_ONLY_ERROR, + createSubagentThreadTurnGuard, + isSubagentThreadWriteBlocked, +} from './guard'; +import { createSubagentThreadId } from './subagentThreadIds'; +import { SubagentThreadTaskStore } from './subagentThreads'; + +function childConversation(): IConversation { + return { + conversationId: 'child-conversation', + endpoint: 'agents', + title: 'Child', + agent_id: 'child-agent', + subagentThread: { + rootConversationId: 'parent-conversation', + parentConversationId: 'parent-conversation', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool-call', + subagentType: 'child-agent', + subagentKind: 'agent', + depth: 1, + }, + } as IConversation; +} + +function makeStore(): SubagentThreadTaskStore { + const unused = jest.fn(); + return new SubagentThreadTaskStore({ + acquireSubagentThreadLease: unused as AllMethods['acquireSubagentThreadLease'], + 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'], + releaseSubagentThreadLease: unused as AllMethods['releaseSubagentThreadLease'], + reserveSubagentThread: unused as AllMethods['reserveSubagentThread'], + renewSubagentThreadLease: unused as AllMethods['renewSubagentThreadLease'], + saveConvo: unused as AllMethods['saveConvo'], + saveMessage: unused as AllMethods['saveMessage'], + }); +} + +function createApp(getConvo: AllMethods['getConvo'], store: SubagentThreadTaskStore) { + const app = express(); + app.use(express.json()); + app.use((req, _res, next) => { + req.user = { id: 'user-1', tenantId: 'tenant-1' }; + next(); + }); + app.post('/chat', createSubagentThreadTurnGuard({ getConvo, store }), (req, res) => { + res.json({ + ok: true, + resolvedConversationId: (req as typeof req & { resolvedConversation?: IConversation | null }) + .resolvedConversation?.conversationId, + }); + }); + return app; +} + +describe('subagent child-thread write policy', () => { + it('allows an ordinary conversation and new-conversation requests', async () => { + const getConvo = jest.fn().mockResolvedValue({ + conversationId: 'ordinary-conversation', + endpoint: 'agents', + }); + const store = makeStore(); + const app = createApp(getConvo, store); + + const ordinary = await request(app) + .post('/chat') + .send({ conversationId: 'ordinary-conversation' }); + const fresh = await request(app).post('/chat').send({ conversationId: 'new' }); + + expect(ordinary.status).toBe(200); + expect(ordinary.body).toEqual({ + ok: true, + resolvedConversationId: 'ordinary-conversation', + }); + expect(fresh.status).toBe(200); + expect(getConvo).toHaveBeenCalledTimes(1); + }); + + it('rejects every model-bound human turn against a durable child conversation', async () => { + const store = makeStore(); + const response = await request( + createApp(jest.fn().mockResolvedValue(childConversation()), store), + ) + .post('/chat') + .send({ conversationId: 'child-conversation', agent_id: 'child-agent' }); + + expect(response.status).toBe(409); + expect(response.body).toEqual({ error: CHILD_THREAD_READ_ONLY_ERROR }); + }); + + it('rejects a provisional child before its conversation becomes durable', async () => { + const store = makeStore(); + jest.spyOn(store, 'isThreadActiveForOwner').mockReturnValue(true); + const getConvo = jest.fn().mockResolvedValue(null); + + const response = await request(createApp(getConvo, store)) + .post('/chat') + .send({ conversationId: 'provisional-child' }); + + expect(response.status).toBe(409); + expect(response.body).toEqual({ error: CHILD_THREAD_READ_ONLY_ERROR }); + expect(store.isThreadActiveForOwner).toHaveBeenCalledWith( + 'user-1', + 'provisional-child', + 'tenant-1', + ); + expect(getConvo).not.toHaveBeenCalled(); + }); + + it('rejects a reserved provisional child on a different API worker', async () => { + const store = makeStore(); + const getConvo = jest.fn().mockResolvedValue(null); + const reservedThreadId = createSubagentThreadId('scope', 'attempt'); + + const response = await request(createApp(getConvo, store)) + .post('/chat') + .send({ conversationId: reservedThreadId }); + + expect(response.status).toBe(409); + expect(response.body).toEqual({ error: CHILD_THREAD_READ_ONLY_ERROR }); + expect(getConvo).not.toHaveBeenCalled(); + }); + + it('keeps the shared policy owner-scoped and treats child lineage as immutable', async () => { + const store = makeStore(); + const getConvo = jest.fn().mockResolvedValue(childConversation()); + + await expect( + isSubagentThreadWriteBlocked( + { getConvo, store }, + { + userId: 'owner', + conversationId: 'child-conversation', + tenantId: 'tenant-1', + }, + ), + ).resolves.toBe(true); + + expect(getConvo).toHaveBeenCalledWith('owner', 'child-conversation'); + }); +}); diff --git a/packages/api/src/agents/guard.ts b/packages/api/src/agents/guard.ts new file mode 100644 index 0000000000..c185bdaa4c --- /dev/null +++ b/packages/api/src/agents/guard.ts @@ -0,0 +1,106 @@ +import { Constants } from 'librechat-data-provider'; +import type { ConversationMethods, IConversation } from '@librechat/data-schemas'; +import type { NextFunction, Request, RequestHandler, Response } from 'express'; +import type { SubagentThreadTaskStore } from './subagentThreads'; +import { isReservedSubagentThreadId } from './subagentThreadIds'; + +export const CHILD_THREAD_READ_ONLY_ERROR = + 'This subagent thread is view-only. Continue it from its parent agent or create a separate chat.'; + +interface SubagentTurnBody { + conversationId?: unknown; + arg?: { conversationId?: unknown }; +} + +interface SubagentTurnUser { + id?: string; + _id?: string | { toString(): string }; + tenantId?: string; +} + +export interface SubagentThreadWriteGuardDeps { + getConvo: ConversationMethods['getConvo']; + store: SubagentThreadTaskStore; +} + +export interface SubagentThreadWriteTarget { + userId: string; + conversationId: string; + tenantId?: string; +} + +interface SubagentThreadWriteResolution { + blocked: boolean; + conversation?: IConversation | null; +} + +interface ResolvedConversationRequest extends Request { + resolvedConversation?: IConversation | null; +} + +async function resolveSubagentThreadWrite( + { getConvo, store }: SubagentThreadWriteGuardDeps, + { userId, conversationId, tenantId }: SubagentThreadWriteTarget, +): Promise { + /** New child IDs are returned synchronously by the SDK before Mongo creation can + * finish. Their reserved UUID namespace closes that brief window on every replica. */ + if (isReservedSubagentThreadId(conversationId)) { + return { blocked: true }; + } + if (store.isThreadActiveForOwner(userId, conversationId, tenantId)) { + return { blocked: true }; + } + const conversation = await getConvo(userId, conversationId); + return { blocked: conversation?.subagentThread != null, conversation }; +} + +/** Applies the same immutable-child policy to every server write adapter. */ +export async function isSubagentThreadWriteBlocked( + deps: SubagentThreadWriteGuardDeps, + target: SubagentThreadWriteTarget, +): Promise { + return (await resolveSubagentThreadWrite(deps, target)).blocked; +} + +/** Rejects model-bound turns for durable or provisionally-created child threads. */ +export function createSubagentThreadTurnGuard(deps: SubagentThreadWriteGuardDeps): RequestHandler { + return async (request: Request, res: Response, next: NextFunction): Promise => { + const body = request.body as SubagentTurnBody | undefined; + const user = request.user as SubagentTurnUser | undefined; + const candidateConversationId = body?.conversationId ?? body?.arg?.conversationId; + if ( + typeof candidateConversationId !== 'string' || + candidateConversationId === '' || + candidateConversationId === Constants.NEW_CONVO + ) { + next(); + return; + } + const rawUserId = user?.id ?? user?._id; + if (rawUserId == null) { + next(); + return; + } + const userId = String(rawUserId); + const tenantId = + typeof user?.tenantId === 'string' && user.tenantId !== '' ? user.tenantId : undefined; + + try { + const resolved = await resolveSubagentThreadWrite(deps, { + userId, + conversationId: candidateConversationId, + ...(tenantId == null ? {} : { tenantId }), + }); + if (resolved.conversation !== undefined) { + (request as ResolvedConversationRequest).resolvedConversation = resolved.conversation; + } + if (!resolved.blocked) { + next(); + return; + } + res.status(409).json({ error: CHILD_THREAD_READ_ONLY_ERROR }); + } catch (error) { + next(error); + } + }; +} diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 3cda8f3837..c99e4aea24 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -10,6 +10,7 @@ import type { ToolCallRequest, ToolExecuteResult, ToolExecuteBatchRequest, + SubagentTaskConfig, } from '@librechat/agents'; import type { StructuredToolInterface } from '@librechat/agents/langchain/tools'; import type { ValidationIssue } from '@librechat/data-schemas'; @@ -82,6 +83,8 @@ export interface ToolExecuteOptions { /** Additional configurable properties to merge (e.g., userMCPAuthMap) */ configurable?: Record; }>; + /** Trusted detached-subagent task scope for polling and parent controls. */ + subagentTasks?: SubagentTaskConfig; /** Callback to process tool artifacts (code output files, file citations, etc.) */ toolEndCallback?: ToolEndCallback; /** @@ -3816,7 +3819,8 @@ function buildToolCallConfig( } export function createToolExecuteHandler(options: ToolExecuteOptions): EventHandler { - const { loadTools, toolEndCallback, persistBackgroundCodeResult, emitAttachment } = options; + const { loadTools, toolEndCallback, persistBackgroundCodeResult, emitAttachment, subagentTasks } = + options; return { handle: async (_event: string, data: ToolExecuteBatchRequest) => { @@ -3890,16 +3894,17 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand | string[] | undefined; const backgroundEnabledForRun = (backgroundToolNames?.length ?? 0) > 0; + const backgroundControlEnabled = backgroundEnabledForRun || subagentTasks != null; const backgroundToolSet: ReadonlySet = backgroundEnabledForRun ? new Set(backgroundToolNames) : EMPTY_BACKGROUND_TOOL_SET; - const backgroundReq = backgroundEnabledForRun + const backgroundReq = backgroundControlEnabled ? (mergedConfigurable?.req as ServerRequest | undefined) : undefined; - const backgroundUserId = backgroundEnabledForRun + const backgroundUserId = backgroundControlEnabled ? resolveBackgroundUserId(mergedConfigurable) : ''; - const backgroundConversationId = backgroundEnabledForRun + const backgroundConversationId = backgroundControlEnabled ? (((metadata as Record)?.thread_id as string | undefined) ?? (mergedConfigurable?.thread_id as string | undefined) ?? (backgroundReq?.body as { conversationId?: string } | undefined)?.conversationId ?? @@ -4083,11 +4088,12 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand const results: ToolExecuteResult[] = await Promise.all( toolCalls.map(async (tc: ToolCallRequest) => { - if (backgroundEnabledForRun && tc.name === CHECK_BACKGROUND_TASK_NAME) { + if (backgroundControlEnabled && tc.name === CHECK_BACKGROUND_TASK_NAME) { const pollContent = runCheckBackgroundTask({ userId: backgroundUserId, conversationId: backgroundConversationId, args: tc.args, + subagentTasks, }); /** Deliver a completed task's artifact through THIS live poll * turn (once): the tool's own turn finalized before the diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index a91ec8fee0..95d1b34b53 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -13,6 +13,7 @@ export * from './errors'; export * from './envelope'; export * from './execution'; export * from './handlers'; +export * from './guard'; export * from './harvest'; export * from './initialize'; export * from './legacy'; @@ -30,6 +31,7 @@ export * from './responses'; export * from './skills'; export * from './phases'; export * from './startup'; +export * from './subagentThreads'; export * from './skillConfigurable'; export * from './skillFiles'; export * from './codeFilesSession'; diff --git a/packages/api/src/agents/run.ts b/packages/api/src/agents/run.ts index 27634dad32..5a14ab853f 100644 --- a/packages/api/src/agents/run.ts +++ b/packages/api/src/agents/run.ts @@ -26,6 +26,7 @@ import type { RunConfig, IState, LCTool, + SubagentTaskConfig, } from '@librechat/agents'; import type { Agent, @@ -44,6 +45,7 @@ import type { SubagentUsageEvent } from '~/agents/usage'; import type * as t from '~/types'; import { CHECK_BACKGROUND_TASK_NAME, + registerBackgroundTaskTool, stripBackgroundFromToolRegistry, stripBackgroundFromToolDefinitions, } from '~/agents/background'; @@ -1045,6 +1047,7 @@ function buildSubagentConfigs( ancestors: Set = new Set(), depth = 0, prebuiltGraphInputs?: ReadonlyMap, + detachedTasksEnabled = false, ): SubagentConfigEntry[] { if (!agent.subagents?.enabled) { return []; @@ -1065,8 +1068,12 @@ function buildSubagentConfigs( * invocations would forward to tools that never declared it. The * resolver keeps a provided `agentInputs` even with `self: true`. */ - const hasBackground = (agent.backgroundToolNames?.length ?? 0) > 0; + const hasBackground = detachedTasksEnabled || (agent.backgroundToolNames?.length ?? 0) > 0; const hasInjectedIntent = (agent.intentToolNames?.length ?? 0) > 0; + const sanitizedToolRegistry = stripIntentFromToolRegistry( + stripBackgroundFromToolRegistry(agentInput.toolRegistry, agent.backgroundToolNames), + agent.intentToolNames, + ); configs.push({ self: true, type: SELF_SUBAGENT_TYPE, @@ -1085,10 +1092,13 @@ function buildSubagentConfigs( ), agent.intentToolNames, ), - toolRegistry: stripIntentFromToolRegistry( - stripBackgroundFromToolRegistry(agentInput.toolRegistry, agent.backgroundToolNames), - agent.intentToolNames, - ), + /** `registerBackgroundTaskTool` mutates the parent registry after + * configs are built. Detach its self-child snapshot so the host + * poll tool cannot appear there through that shared Map. */ + toolRegistry: + detachedTasksEnabled && sanitizedToolRegistry != null + ? new Map(sanitizedToolRegistry) + : sanitizedToolRegistry, }, } : {}), @@ -1241,6 +1251,7 @@ export async function createRun({ calibrationRatio, appConfig, subagentUsageSink, + subagentTasks, steering, activityLabel, activityPhase, @@ -1294,6 +1305,8 @@ export async function createRun({ * Switch to the `RunConfig` pick once the dependency is bumped. */ subagentUsageSink?: (event: SubagentUsageEvent) => void; + /** Host-owned detached-subagent task store and trusted parent-thread scope. */ + subagentTasks?: SubagentTaskConfig; /** * The run-scoped steer-drain hook (a `PostToolBatch` callback built via * `createSteerDrainHook`). Registered on the run's hook registry independent @@ -1608,12 +1621,19 @@ export async function createRun({ undefined, 0, prebuiltGraphInputs, + subagentTasks != null, ); if (subagentConfigs.length > 0) { agentInput.subagentConfigs = subagentConfigs; /** Seed the SDK countdown that bounds nested delegation across isolated child graphs. */ agentInput.maxSubagentDepth = MAX_SUBAGENT_DEPTH; } + if (subagentTasks != null) { + agentInput.toolDefinitions = registerBackgroundTaskTool({ + toolRegistry: agentInput.toolRegistry, + toolDefinitions: agentInput.toolDefinitions, + }).toolDefinitions; + } agentInputs.push(agentInput); } @@ -1777,6 +1797,7 @@ export async function createRun({ calibrationRatio, indexTokenCountMap, subagentUsageSink, + subagentTasks, // Exclude side-effecting / large-free-form-arg tools from eager execution. // Eager speculatively runs a tool mid-stream; for a big streamed arg (a // file body, a bash heredoc, a code block) the accumulated args can diverge diff --git a/packages/api/src/agents/subagentTaskContext.ts b/packages/api/src/agents/subagentTaskContext.ts new file mode 100644 index 0000000000..30eb6ed4dc --- /dev/null +++ b/packages/api/src/agents/subagentTaskContext.ts @@ -0,0 +1,28 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; +import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; + +/** + * Detached subagents outlive the parent turn that launched them. This + * request-local collector lets the shared SDK usage sink recognize those + * calls without retaining a request object or changing the SDK task-store + * contract. AsyncLocalStorage follows the detached executor's promise chain + * and naturally isolates concurrent child tasks. + */ +const detachedUsageStorage = new AsyncLocalStorage(); + +export function runWithDetachedSubagentUsage( + usage: UsageMetadata[], + run: () => Promise, +): Promise { + return detachedUsageStorage.run(usage, run); +} + +/** Records one detached usage item and reports whether a task context owned it. */ +export function collectDetachedSubagentUsage(usage: UsageMetadata): boolean { + const collector = detachedUsageStorage.getStore(); + if (collector == null) { + return false; + } + collector.push(usage); + return true; +} diff --git a/packages/api/src/agents/subagentThreadIds.ts b/packages/api/src/agents/subagentThreadIds.ts new file mode 100644 index 0000000000..db387a0d50 --- /dev/null +++ b/packages/api/src/agents/subagentThreadIds.ts @@ -0,0 +1,31 @@ +import { createHash } from 'node:crypto'; + +/** RFC 9562 v8 UUIDs with a fixed `b` variant nibble form the host-reserved namespace. */ +const RESERVED_SUBAGENT_THREAD_ID = + /^[0-9a-f]{8}-[0-9a-f]{4}-8[0-9a-f]{3}-b[0-9a-f]{3}-[0-9a-f]{12}$/i; + +function sha256(value: string): string { + return createHash('sha256').update(value).digest('hex'); +} + +/** Deterministic so the same initial tool-call retry converges across API replicas. */ +export function createSubagentThreadId(scopeId: string, idempotencyKey: string): string { + const hash = sha256( + `librechat:subagent-thread:v1\u0000${scopeId.trim()}\u0000${idempotencyKey.trim()}`, + ); + return `${hash.slice(0, 8)}-${hash.slice(8, 12)}-8${hash.slice(13, 16)}-b${hash.slice( + 17, + 20, + )}-${hash.slice(20, 32)}`; +} + +export function isReservedSubagentThreadId(conversationId: string): boolean { + return RESERVED_SUBAGENT_THREAD_ID.test(conversationId); +} + +/** Opaque database key; raw SDK idempotency material never needs to be persisted. */ +export function createSubagentAttemptKey(scopeId: string, idempotencyKey: string): string { + return sha256( + `librechat:subagent-attempt:v1\u0000${scopeId.trim()}\u0000${idempotencyKey.trim()}`, + ); +} diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts new file mode 100644 index 0000000000..b75ec3b7ee --- /dev/null +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -0,0 +1,1331 @@ +import mongoose from 'mongoose'; +import { randomUUID } from 'node:crypto'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import { Constants, EModelEndpoint } from 'librechat-data-provider'; +import { createMethods, createModels, logger } from '@librechat/data-schemas'; +import { AIMessage, HumanMessage } from '@librechat/agents/langchain/messages'; +import type { + SubagentTaskRuntime, + SubagentTaskStartRequest, + SubagentTaskStartResult, +} from '@librechat/agents'; +import type { AllMethods, IConversation, IMessage } from '@librechat/data-schemas'; +import type { BaseMessage } from '@librechat/agents/langchain/messages'; +import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; +import { buildSubagentThreadTaskConfig, SubagentThreadTaskStore } from './subagentThreads'; +import { createSubagentAttemptKey } from './subagentThreadIds'; +import { createSubagentUsageSink } from './usage'; + +let mongod: MongoMemoryServer; +let methods: AllMethods; +let loggerErrorSpy: jest.SpyInstance; + +function taskRequest( + scopeId: string, + overrides: Partial = {}, +): SubagentTaskStartRequest { + const input = overrides.input ?? 'Investigate the issue.'; + return { + scopeId, + idempotencyKey: overrides.idempotencyKey ?? randomUUID(), + parentRunId: overrides.parentRunId ?? randomUUID(), + parentAgentId: overrides.parentAgentId ?? 'parent-agent', + parentToolCallId: overrides.parentToolCallId ?? randomUUID(), + ...(overrides.requestFingerprint == null + ? {} + : { requestFingerprint: overrides.requestFingerprint }), + input, + subagentKind: overrides.subagentKind ?? 'agent', + subagentType: overrides.subagentType ?? 'researcher-agent', + run: + overrides.run ?? + (async (_runtime: SubagentTaskRuntime, initialMessages = []) => ({ + content: 'Completed the investigation.', + messages: [ + ...initialMessages, + new HumanMessage(input), + new AIMessage('Completed the investigation.'), + ], + })), + ...(overrides.threadId == null ? {} : { threadId: overrides.threadId }), + }; +} + +async function waitForSettled( + store: SubagentThreadTaskStore, + scopeId: string, + started: SubagentTaskStartResult, +): Promise { + const accepted = requireAccepted(started); + for (let attempt = 0; attempt < 200; attempt += 1) { + const task = store.get(scopeId, accepted.task.taskId); + if (task != null && task.status !== 'running') { + return; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error('Timed out waiting for the subagent task.'); +} + +function requireAccepted( + started: SubagentTaskStartResult, +): Extract { + if (!started.accepted) { + throw new Error('Expected the task to be accepted.'); + } + return started; +} + +function requireThreadId(started: SubagentTaskStartResult): string { + const accepted = requireAccepted(started); + if (!accepted.task.threadId) { + throw new Error('Expected the accepted task to expose its durable thread id.'); + } + return accepted.task.threadId; +} + +async function saveParent( + userId: string, + conversationId: string, + overrides: Record = {}, +): Promise { + const saved = await methods.saveConvo( + { userId }, + { + conversationId, + endpoint: EModelEndpoint.agents, + title: 'Parent thread', + agent_id: 'parent-agent', + ...overrides, + }, + ); + if (saved == null || 'message' in saved) { + throw new Error('Failed to save parent conversation.'); + } + return saved; +} + +beforeAll(async () => { + loggerErrorSpy = jest.spyOn(logger, 'error').mockImplementation(() => logger); + mongod = await MongoMemoryServer.create(); + createModels(mongoose); + methods = createMethods(mongoose); + await mongoose.connect(mongod.getUri()); +}); + +afterAll(async () => { + await mongoose.disconnect(); + await mongod.stop(); + loggerErrorSpy.mockRestore(); +}); + +beforeEach(async () => { + await Promise.all([ + (mongoose.models.Message as mongoose.Model).deleteMany({}), + (mongoose.models.Conversation as mongoose.Model).deleteMany({}), + ]); +}); + +describe('SubagentThreadTaskStore', () => { + it('maps one logical SDK thread to a durable, view-only LibreChat conversation', async () => { + const userId = 'user-1'; + 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)); + await waitForSettled(store, config.scopeId, started); + const threadId = requireThreadId(started); + + const conversation = await methods.getConvo(userId, threadId); + expect(conversation).toMatchObject({ + conversationId: threadId, + endpoint: EModelEndpoint.agents, + agent_id: 'researcher-agent', + subagentThread: { + rootConversationId: parentConversationId, + parentConversationId, + parentAgentId: 'parent-agent', + subagentType: 'researcher-agent', + subagentKind: 'agent', + depth: 1, + }, + }); + expect(conversation?.subagentThread).not.toHaveProperty('userRunnable'); + const messages = await methods.getMessages( + { user: userId, conversationId: threadId }, + '+subagentTranscript', + ); + expect(messages.map((message) => message.text)).toEqual([ + 'Investigate the issue.', + 'Completed the investigation.', + ]); + expect(messages[1].subagentTranscript).toMatchObject({ + taskId: requireAccepted(started).task.taskId, + mode: 'append', + }); + }); + + it('waits for initial parent persistence before creating the first child', async () => { + const userId = 'parent-gate-user'; + const parentConversationId = randomUUID(); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let releaseParent = (_value: unknown): void => undefined; + const parentPersistence = new Promise((resolve) => { + releaseParent = resolve; + }); + const run = jest.fn(taskRequest(config.scopeId).run); + store.registerParentPersistence(config.scopeId, parentPersistence); + + const started = store.start(taskRequest(config.scopeId, { run })); + await new Promise((resolve) => setTimeout(resolve, 20)); + + expect(run).not.toHaveBeenCalled(); + expect(await methods.getConvo(userId, requireThreadId(started))).toBeNull(); + + await saveParent(userId, parentConversationId); + releaseParent({ + message: { messageId: 'parent-message', conversationId: parentConversationId }, + }); + await waitForSettled(store, config.scopeId, started); + + expect(run).toHaveBeenCalledTimes(1); + expect(await methods.getConvo(userId, requireThreadId(started))).not.toBeNull(); + }); + + it('fails without leaving an orphan when parent persistence rejects', async () => { + const userId = 'parent-gate-failure-user'; + const parentConversationId = randomUUID(); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(taskRequest(config.scopeId).run); + store.registerParentPersistence( + config.scopeId, + Promise.reject(new Error('parent write failed')), + ); + + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).not.toHaveBeenCalled(); + expect(await methods.getConvo(userId, requireThreadId(started))).toBeNull(); + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'error', + }); + }); + + it('retains a resolved-but-unsaved parent failure until a valid write supersedes it', async () => { + const userId = 'parent-empty-result-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const rejectedRun = jest.fn(taskRequest(config.scopeId).run); + store.registerParentPersistence(config.scopeId, Promise.resolve({})); + await new Promise((resolve) => setImmediate(resolve)); + + const rejected = store.start(taskRequest(config.scopeId, { run: rejectedRun })); + await waitForSettled(store, config.scopeId, rejected); + + expect(rejectedRun).not.toHaveBeenCalled(); + expect(await methods.getConvo(userId, requireThreadId(rejected))).toBeNull(); + + store.registerParentPersistence( + config.scopeId, + Promise.resolve({ + message: { messageId: 'next-parent-message', conversationId: parentConversationId }, + }), + ); + const acceptedRun = jest.fn(taskRequest(config.scopeId).run); + const accepted = store.start(taskRequest(config.scopeId, { run: acceptedRun })); + await waitForSettled(store, config.scopeId, accepted); + + expect(acceptedRun).toHaveBeenCalledTimes(1); + }); + + it('continues with canonical transcript only and ignores non-canonical visible rows', async () => { + const userId = 'canonical-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const first = store.start(taskRequest(config.scopeId)); + await waitForSettled(store, config.scopeId, first); + const threadId = requireThreadId(first); + + const visibleUserId = randomUUID(); + await methods.saveMessage( + { userId }, + { + messageId: visibleUserId, + conversationId: threadId, + parentMessageId: `${requireAccepted(first).task.taskId}:assistant`, + sender: 'User', + text: 'A non-canonical human edit.', + endpoint: EModelEndpoint.agents, + isCreatedByUser: true, + }, + ); + await methods.saveMessage( + { userId }, + { + messageId: randomUUID(), + conversationId: threadId, + parentMessageId: visibleUserId, + sender: 'Assistant', + text: 'A non-canonical answer.', + endpoint: EModelEndpoint.agents, + isCreatedByUser: false, + }, + ); + + let restored: BaseMessage[] = []; + const continued = store.start( + taskRequest(config.scopeId, { + threadId, + input: 'Continue the investigation.', + run: async (_runtime, initialMessages = []) => { + restored = initialMessages; + return { + content: 'Continued.', + messages: [ + ...initialMessages, + new HumanMessage('Continue the investigation.'), + new AIMessage('Continued.'), + ], + }; + }, + }), + ); + await waitForSettled(store, config.scopeId, continued); + + expect(restored.map((message) => message.content)).toEqual([ + 'Investigate the issue.', + 'Completed the investigation.', + ]); + }); + + it('reuses the original task and thread for an idempotent replay', async () => { + const userId = 'idempotent-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const request = taskRequest(config.scopeId, { idempotencyKey: 'same-attempt' }); + + const first = store.start(request); + const replay = store.start(request); + + expect(first.accepted).toBe(true); + expect(replay).toMatchObject({ + accepted: true, + isNew: false, + task: requireAccepted(first).task, + }); + await waitForSettled(store, config.scopeId, first); + expect( + (await methods.getMessages({ user: userId, conversationId: requireThreadId(first) })).length, + ).toBe(2); + }); + + it('replays a completed attempt across API workers without executing or billing twice', async () => { + const userId = 'durable-idempotency-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const firstWorker = new SubagentThreadTaskStore(methods); + const secondWorker = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); + const firstRun = jest.fn(async () => ({ + content: 'Original durable result.', + messages: [new HumanMessage('Run once.'), new AIMessage('Original durable result.')], + })); + const first = firstWorker.start( + taskRequest(config.scopeId, { + idempotencyKey: 'cross-worker-attempt', + requestFingerprint: 'same-inputs', + input: 'Run once.', + run: firstRun, + }), + ); + await waitForSettled(firstWorker, config.scopeId, first); + + const replayRun = jest.fn(taskRequest(config.scopeId).run); + const replay = secondWorker.start( + taskRequest(config.scopeId, { + idempotencyKey: 'cross-worker-attempt', + requestFingerprint: 'same-inputs', + input: 'Run once.', + run: replayRun, + }), + ); + expect(requireThreadId(replay)).toBe(requireThreadId(first)); + await waitForSettled(secondWorker, config.scopeId, replay); + + expect(firstRun).toHaveBeenCalledTimes(1); + expect(replayRun).not.toHaveBeenCalled(); + expect(secondWorker.claim(config.scopeId, requireAccepted(replay).task.taskId)).toMatchObject({ + status: 'completed', + result: 'Original durable result.', + }); + expect( + await methods.getMessages({ user: userId, conversationId: requireThreadId(first) }), + ).toHaveLength(2); + }); + + it('rejects a conflicting durable retry across API workers', async () => { + const userId = 'durable-conflict-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const firstWorker = new SubagentThreadTaskStore(methods); + const secondWorker = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); + const first = firstWorker.start( + taskRequest(config.scopeId, { + idempotencyKey: 'reused-key', + requestFingerprint: 'original-inputs', + }), + ); + await waitForSettled(firstWorker, config.scopeId, first); + + const conflictingRun = jest.fn(taskRequest(config.scopeId).run); + const conflicting = secondWorker.start( + taskRequest(config.scopeId, { + idempotencyKey: 'reused-key', + requestFingerprint: 'different-inputs', + run: conflictingRun, + }), + ); + await waitForSettled(secondWorker, config.scopeId, conflicting); + + expect(conflictingRun).not.toHaveBeenCalled(); + expect( + secondWorker.claim(config.scopeId, requireAccepted(conflicting).task.taskId), + ).toMatchObject({ status: 'error' }); + }); + + it('closes an abandoned durable attempt without re-executing it', async () => { + const userId = 'abandoned-attempt-user'; + const parentConversationId = randomUUID(); + const threadId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + await methods.saveConvo( + { userId }, + { + conversationId: threadId, + endpoint: EModelEndpoint.agents, + title: 'Abandoned child', + agent_id: 'researcher-agent', + subagentThread: { + rootConversationId: parentConversationId, + parentConversationId, + parentMessageId: 'parent-run', + parentToolCallId: 'parent-tool', + parentAgentId: 'parent-agent', + subagentType: 'researcher-agent', + subagentKind: 'agent', + depth: 1, + }, + }, + ); + await methods.saveMessage( + { userId }, + { + messageId: 'abandoned:user', + conversationId: threadId, + parentMessageId: String(Constants.NO_PARENT), + sender: 'User', + text: 'Run once.', + endpoint: EModelEndpoint.agents, + isCreatedByUser: true, + subagentTask: { + attemptKey: createSubagentAttemptKey(config.scopeId, 'abandoned-attempt'), + requestFingerprint: 'same-inputs', + status: 'running', + }, + }, + ); + const run = jest.fn(taskRequest(config.scopeId).run); + const retry = store.start( + taskRequest(config.scopeId, { + threadId, + idempotencyKey: 'abandoned-attempt', + requestFingerprint: 'same-inputs', + run, + }), + ); + await waitForSettled(store, config.scopeId, retry); + + expect(run).not.toHaveBeenCalled(); + expect(store.claim(config.scopeId, requireAccepted(retry).task.taskId)).toMatchObject({ + status: 'error', + error: + 'Subagent task failed: The prior execution ended before its result could be persisted.', + }); + const messages = await methods.getMessages( + { user: userId, conversationId: threadId }, + '+subagentTask', + ); + expect(messages.map((message) => message.subagentTask?.status)).toEqual(['running', 'error']); + }); + + it('holds one active lease per child and exposes provisional ownership safely', async () => { + const userId = 'lease-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let release = (_value: { content: string; messages: BaseMessage[] }): void => undefined; + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const result = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => { + release = resolve; + }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => { + markEntered(); + return result; + }, + }), + ); + const threadId = requireThreadId(started); + + expect(store.isThreadActiveForOwner(userId, threadId)).toBe(true); + expect(store.isThreadActiveForOwner('different-user', threadId)).toBe(false); + expect(store.isThreadActiveForOwner(userId, threadId, 'tenant-a')).toBe(false); + await entered; + expect( + store.start( + taskRequest(config.scopeId, { + threadId, + idempotencyKey: 'different-attempt', + }), + ), + ).toEqual({ accepted: false, reason: 'capacity' }); + + release({ + content: 'Lease completed.', + messages: [new HumanMessage('Investigate the issue.'), new AIMessage('Lease completed.')], + }); + await waitForSettled(store, config.scopeId, started); + expect(store.isThreadActiveForOwner(userId, threadId)).toBe(false); + }); + + it('serializes continuations across independent API worker stores', async () => { + const userId = 'cross-worker-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const firstWorker = new SubagentThreadTaskStore(methods); + const secondWorker = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); + const initial = firstWorker.start(taskRequest(config.scopeId)); + await waitForSettled(firstWorker, config.scopeId, initial); + const threadId = requireThreadId(initial); + + let releaseFirst = (_value: { content: string; messages: BaseMessage[] }): void => undefined; + let markFirstEntered = (): void => undefined; + const firstEntered = new Promise((resolve) => { + markFirstEntered = resolve; + }); + const firstResult = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => { + releaseFirst = resolve; + }); + const first = firstWorker.start( + taskRequest(config.scopeId, { + threadId, + input: 'First continuation.', + run: async () => { + markFirstEntered(); + return firstResult; + }, + }), + ); + await firstEntered; + + const secondRun = jest.fn(taskRequest(config.scopeId).run); + const second = secondWorker.start( + taskRequest(config.scopeId, { + threadId, + input: 'Overlapping continuation.', + run: secondRun, + }), + ); + await waitForSettled(secondWorker, config.scopeId, second); + expect(secondRun).not.toHaveBeenCalled(); + expect(secondWorker.claim(config.scopeId, requireAccepted(second).task.taskId)).toMatchObject({ + status: 'error', + }); + + releaseFirst({ + content: 'First continuation completed.', + messages: [new HumanMessage('First continuation.'), new AIMessage('Completed.')], + }); + await waitForSettled(firstWorker, config.scopeId, first); + }); + + it('renews the shared lease while a continuation transcript is being prepared', async () => { + const userId = 'slow-prepare-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let slowThreadId = ''; + let blockNextRead = false; + let markPreparing = (): void => undefined; + const preparing = new Promise((resolve) => { + markPreparing = resolve; + }); + let releasePreparation = (): void => undefined; + const preparationRelease = new Promise((resolve) => { + releasePreparation = resolve; + }); + const slowMethods = { + ...methods, + getMessages: jest.fn(async (...args: Parameters) => { + if (blockNextRead && args[0].conversationId === slowThreadId) { + blockNextRead = false; + markPreparing(); + await preparationRelease; + } + return methods.getMessages(...args); + }), + }; + const options = { leaseTtlMs: 60, leaseHeartbeatMs: 10 }; + const firstWorker = new SubagentThreadTaskStore(slowMethods, options); + const secondWorker = new SubagentThreadTaskStore(methods, options); + const config = buildSubagentThreadTaskConfig(firstWorker, { userId, parentConversationId }); + const initial = firstWorker.start(taskRequest(config.scopeId)); + await waitForSettled(firstWorker, config.scopeId, initial); + slowThreadId = requireThreadId(initial); + blockNextRead = true; + + const firstRun = jest.fn(taskRequest(config.scopeId).run); + const first = firstWorker.start( + taskRequest(config.scopeId, { + threadId: slowThreadId, + idempotencyKey: 'slow-preparation', + run: firstRun, + }), + ); + await preparing; + await new Promise((resolve) => setTimeout(resolve, 100)); + + const overlappingRun = jest.fn(taskRequest(config.scopeId).run); + const overlapping = secondWorker.start( + taskRequest(config.scopeId, { + threadId: slowThreadId, + idempotencyKey: 'overlapping-preparation', + run: overlappingRun, + }), + ); + await waitForSettled(secondWorker, config.scopeId, overlapping); + expect(overlappingRun).not.toHaveBeenCalled(); + + releasePreparation(); + await waitForSettled(firstWorker, config.scopeId, first); + expect(firstRun).toHaveBeenCalledTimes(1); + }); + + it('rechecks account deletion after acquiring the shared lease', async () => { + const userId = 'lease-fence-gap-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let ownerActive = true; + const fencedMethods = { + ...methods, + acquireSubagentThreadLease: jest.fn( + async (...args: Parameters) => { + const acquired = await methods.acquireSubagentThreadLease(...args); + ownerActive = false; + return acquired; + }, + ), + }; + const store = new SubagentThreadTaskStore(fencedMethods, { + isOwnerActive: async () => ownerActive, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(taskRequest(config.scopeId).run); + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).not.toHaveBeenCalled(); + expect(await methods.getConvo(userId, requireThreadId(started))).toBeNull(); + expect(await methods.countActiveSubagentThreadLeases({ user: userId, now: new Date() })).toBe( + 0, + ); + }); + + it('lets account deletion on one worker drain a child running on another', async () => { + const userId = 'owner-drain-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let ownerActive = true; + const options = { + isOwnerActive: async () => ownerActive, + leaseTtlMs: 60, + leaseHeartbeatMs: 10, + ownerDrainTimeoutMs: 1_000, + ownerDrainPollMs: 5, + }; + const workerStore = new SubagentThreadTaskStore(methods, options); + const deletingStore = new SubagentThreadTaskStore(methods, options); + const config = buildSubagentThreadTaskConfig(workerStore, { userId, parentConversationId }); + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const started = workerStore.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; + + ownerActive = false; + await deletingStore.cancelAndDrainForOwner(userId); + await waitForSettled(workerStore, config.scopeId, started); + + expect(await methods.countActiveSubagentThreadLeases({ user: userId, now: new Date() })).toBe( + 0, + ); + expect(workerStore.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'cancelled', + }); + }); + + it('fails account deletion closed while a cancelled provider still owns its lease', async () => { + const userId = 'stubborn-provider-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let ownerActive = true; + const options = { + isOwnerActive: async () => ownerActive, + leaseTtlMs: 1_000, + leaseHeartbeatMs: 100, + ownerDrainTimeoutMs: 1_500, + ownerDrainPollMs: 20, + }; + const workerStore = new SubagentThreadTaskStore(methods, options); + const deletingStore = new SubagentThreadTaskStore(methods, options); + const config = buildSubagentThreadTaskConfig(workerStore, { userId, parentConversationId }); + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + let releaseProvider = (_value: { content: string; messages: BaseMessage[] }): void => undefined; + const provider = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => { + releaseProvider = resolve; + }); + const started = workerStore.start( + taskRequest(config.scopeId, { + run: async () => { + markEntered(); + return provider; + }, + }), + ); + await entered; + + ownerActive = false; + await expect(deletingStore.cancelAndDrainForOwner(userId)).rejects.toThrow( + 'Timed out draining detached subagent tasks', + ); + expect(await methods.countActiveSubagentThreadLeases({ user: userId, now: new Date() })).toBe( + 1, + ); + + ownerActive = true; + releaseProvider({ content: 'Stopped.', messages: [new AIMessage('Stopped.')] }); + await waitForSettled(workerStore, config.scopeId, started); + }); + + it('does not overwrite a child title changed while detached execution is running', async () => { + const userId = 'title-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let release = (_value: { content: string; messages: BaseMessage[] }): void => undefined; + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const result = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => { + release = resolve; + }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => { + markEntered(); + return result; + }, + }), + ); + await entered; + await methods.saveConvo( + { userId }, + { conversationId: requireThreadId(started), title: 'Renamed while running' }, + { noUpsert: true }, + ); + + release({ + content: 'Renamed child completed.', + messages: [ + new HumanMessage('Investigate the issue.'), + new AIMessage('Renamed child completed.'), + ], + }); + await waitForSettled(store, config.scopeId, started); + + expect(await methods.getConvo(userId, requireThreadId(started))).toMatchObject({ + title: 'Renamed while running', + }); + }); + + it('lets deletion win over a concurrently settling detached result', async () => { + const userId = 'deletion-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let deletedThreadId = ''; + const deletingMethods = { + ...methods, + saveMessage: jest.fn(async (...args: Parameters) => { + const message = args[1]; + if (message.text === 'Result after deletion.') { + deletedThreadId = message.conversationId ?? ''; + await methods.deleteConvos(userId, { conversationId: deletedThreadId }); + } + return methods.saveMessage(...args); + }), + }; + const store = new SubagentThreadTaskStore(deletingMethods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ + content: 'Result after deletion.', + messages: [new AIMessage('Result after deletion.')], + }), + }), + ); + await waitForSettled(store, config.scopeId, started); + + expect(deletedThreadId).toBe(requireThreadId(started)); + expect(await methods.getConvo(userId, deletedThreadId)).toBeNull(); + expect(await methods.getMessages({ user: userId, conversationId: deletedThreadId })).toEqual( + [], + ); + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'error', + }); + }); + + it('bills detached usage independently and persists its rollup on the child result', async () => { + const userId = 'usage-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const parentUsage: UsageMetadata[] = []; + const recordDetachedUsage = jest.fn().mockResolvedValue(undefined); + const sink = createSubagentUsageSink( + parentUsage, + (usage) => { + usage.cost = 0.25; + }, + recordDetachedUsage, + ); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => { + await sink({ + usage: { input_tokens: 100, output_tokens: 20, total_tokens: 120 }, + model: 'gpt-5-mini', + provider: 'openAI', + subagentType: 'researcher-agent', + subagentRunId: 'child-run', + subagentAgentId: 'researcher-agent', + runId: 'parent-run', + }); + return { + content: 'Usage recorded.', + messages: [ + new HumanMessage('Investigate the issue.'), + new AIMessage('Usage recorded.'), + ], + }; + }, + }), + ); + await waitForSettled(store, config.scopeId, started); + + expect(parentUsage).toEqual([]); + expect(recordDetachedUsage).toHaveBeenCalledTimes(1); + const messages = await methods.getMessages({ + user: userId, + conversationId: requireThreadId(started), + }); + expect(messages[messages.length - 1]?.metadata?.usage).toEqual({ + input: 100, + output: 20, + cacheWrite: 0, + cacheRead: 0, + cost: 0.25, + }); + }); + + it('persists graph children without assigning a saved-agent identity', async () => { + const userId = 'graph-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, { + subagentKind: 'graph', + subagentType: 'research-team', + }), + ); + await waitForSettled(store, config.scopeId, started); + + const conversation = await methods.getConvo(userId, requireThreadId(started)); + expect(conversation?.agent_id).toBeUndefined(); + expect(conversation?.subagentThread).toMatchObject({ + subagentKind: 'graph', + subagentType: 'research-team', + }); + }); + + it('inherits tenant isolation and rejects a cross-tenant continuation', async () => { + const userId = 'tenant-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId, { tenantId: 'tenant-a' }); + const store = new SubagentThreadTaskStore(methods); + const tenantA = buildSubagentThreadTaskConfig(store, { + userId, + parentConversationId, + tenantId: 'tenant-a', + }); + const first = store.start(taskRequest(tenantA.scopeId)); + await waitForSettled(store, tenantA.scopeId, first); + + const tenantB = buildSubagentThreadTaskConfig(store, { + userId, + parentConversationId, + tenantId: 'tenant-b', + }); + const run = jest.fn(taskRequest(tenantB.scopeId).run); + const crossTenant = store.start( + taskRequest(tenantB.scopeId, { threadId: requireThreadId(first), run }), + ); + await waitForSettled(store, tenantB.scopeId, crossTenant); + + expect(run).not.toHaveBeenCalled(); + expect(store.claim(tenantB.scopeId, requireAccepted(crossTenant).task.taskId)).toMatchObject({ + status: 'error', + }); + }); + + it('removes a newly-created child when its first input cannot be persisted', async () => { + const userId = 'rollback-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const failingMethods = { + ...methods, + saveMessage: jest.fn(async (...args: Parameters) => { + if (args[1].isCreatedByUser === true) { + return null; + } + return methods.saveMessage(...args); + }), + }; + const store = new SubagentThreadTaskStore(failingMethods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(taskRequest(config.scopeId).run); + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).not.toHaveBeenCalled(); + expect(await methods.getConvo(userId, requireThreadId(started))).toBeNull(); + }); + + it('persists a compacted canonical replacement without replaying superseded history', async () => { + const userId = 'compaction-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const first = store.start(taskRequest(config.scopeId)); + await waitForSettled(store, config.scopeId, first); + const threadId = requireThreadId(first); + + const compactedMessages: BaseMessage[] = [ + new HumanMessage('Condensed prior work.'), + new AIMessage('Compact state.'), + new HumanMessage('Continue from the compact state.'), + new AIMessage('Compacted continuation complete.'), + ]; + const second = store.start( + taskRequest(config.scopeId, { + threadId, + input: 'Continue from the compact state.', + run: async () => ({ + content: 'Compacted continuation complete.', + messages: compactedMessages, + }), + }), + ); + await waitForSettled(store, config.scopeId, second); + + let restored: BaseMessage[] = []; + const third = store.start( + taskRequest(config.scopeId, { + threadId, + input: 'One more turn.', + run: async (_runtime, initialMessages = []) => { + restored = initialMessages; + return { content: 'Done.', messages: [...initialMessages, new AIMessage('Done.')] }; + }, + }), + ); + await waitForSettled(store, config.scopeId, third); + + expect(restored.map((message) => message.content)).toEqual( + compactedMessages.map((message) => message.content), + ); + const messages = await methods.getMessages( + { user: userId, conversationId: threadId }, + '+subagentTranscript', + ); + expect( + messages.find( + (message) => message.messageId === `${requireAccepted(second).task.taskId}:assistant`, + )?.subagentTranscript?.mode, + ).toBe('replace'); + }); + + it('holds the child lease through cancellation and discards a late success', async () => { + const userId = 'cancel-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const first = store.start(taskRequest(config.scopeId)); + await waitForSettled(store, config.scopeId, first); + const threadId = requireThreadId(first); + + let finishLate = (_value: { content: string; messages: BaseMessage[] }): void => undefined; + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const lateResult = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => { + finishLate = resolve; + }); + const cancelled = store.start( + taskRequest(config.scopeId, { + threadId, + input: 'Long child turn.', + run: async () => { + markEntered(); + return lateResult; + }, + }), + ); + await entered; + expect( + store.control(config.scopeId, requireAccepted(cancelled).task.taskId, { action: 'cancel' }) + .status, + ).toBe('cancelled'); + expect( + store.start( + taskRequest(config.scopeId, { + threadId, + input: 'Must not overlap.', + }), + ), + ).toEqual({ accepted: false, reason: 'capacity' }); + + finishLate({ + content: 'Late success must be discarded.', + messages: [new AIMessage('Late success must be discarded.')], + }); + for (let attempt = 0; attempt < 200; attempt += 1) { + if (!store.isThreadActiveForOwner(userId, threadId)) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + if (attempt === 199) { + throw new Error('Cancelled child execution did not release its durable lease.'); + } + } + + const messages = await methods.getMessages({ user: userId, conversationId: threadId }); + expect(messages.map((message) => message.text)).not.toContain( + 'Late success must be discarded.', + ); + expect(messages.map((message) => message.text)).toContain('Subagent task was cancelled.'); + }); + + it('cancels an active descendant and lets parent deletion remove its durable thread', async () => { + const userId = 'parent-delete-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finishLate = (_value: { content: string; messages: BaseMessage[] }): void => undefined; + let markEntered = (): void => undefined; + const entered = new Promise((resolve) => { + markEntered = resolve; + }); + const lateResult = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => { + finishLate = resolve; + }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => { + markEntered(); + return lateResult; + }, + }), + ); + await entered; + const threadId = requireThreadId(started); + + expect(store.cancelForConversations(userId, [parentConversationId])).toBe(1); + await methods.deleteConvos(userId, { conversationId: parentConversationId }); + finishLate({ content: 'Too late.', messages: [new AIMessage('Too late.')] }); + await waitForSettled(store, config.scopeId, started); + + expect(await methods.getConvo(userId, parentConversationId)).toBeNull(); + expect(await methods.getConvo(userId, threadId)).toBeNull(); + expect(await methods.getMessages({ user: userId, conversationId: threadId })).toEqual([]); + }); + + it('persists task timeouts as failures rather than cancellations', async () => { + const userId = 'timeout-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { taskTimeoutMs: 20 }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + input: 'Run until timeout.', + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + await waitForSettled(store, config.scopeId, started); + for (let attempt = 0; attempt < 200; attempt += 1) { + if (!store.isThreadActiveForOwner(userId, requireThreadId(started))) { + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + if (attempt === 199) { + throw new Error('Timed-out child execution did not finish durable settlement.'); + } + } + + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'error', + error: 'Detached subagent task timed out.', + }); + const messages = await methods.getMessages({ + user: userId, + conversationId: requireThreadId(started), + }); + expect(messages.map((message) => message.text)).toContain( + 'Subagent task failed: The child run could not be completed.', + ); + expect(messages.map((message) => message.text)).not.toContain('Subagent task was cancelled.'); + }); + + it('lets durable settlement win once a successful commit has started', async () => { + const userId = 'commit-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + let markCommitStarted = (): void => undefined; + let releaseCommit = (): void => undefined; + const commitStarted = new Promise((resolve) => { + markCommitStarted = resolve; + }); + const commitRelease = new Promise((resolve) => { + releaseCommit = resolve; + }); + const blockingMethods = { + ...methods, + saveMessage: jest.fn(async (...args: Parameters) => { + if (args[1].text === 'Committed result.') { + markCommitStarted(); + await commitRelease; + } + return methods.saveMessage(...args); + }), + }; + const store = new SubagentThreadTaskStore(blockingMethods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async () => ({ + content: 'Committed result.', + messages: [new AIMessage('Committed result.')], + }), + }), + ); + await commitStarted; + + expect( + store.control(config.scopeId, requireAccepted(started).task.taskId, { action: 'cancel' }), + ).toMatchObject({ status: 'not_running' }); + releaseCommit(); + await waitForSettled(store, config.scopeId, started); + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'completed', + result: 'Committed result.', + }); + }); + + it('fails closed for unknown, cross-parent, and mismatched-identity continuations', async () => { + const userId = 'lineage-user'; + const firstParentId = randomUUID(); + const secondParentId = randomUUID(); + await Promise.all([saveParent(userId, firstParentId), saveParent(userId, secondParentId)]); + const store = new SubagentThreadTaskStore(methods); + const firstConfig = buildSubagentThreadTaskConfig(store, { + userId, + parentConversationId: firstParentId, + }); + const secondConfig = buildSubagentThreadTaskConfig(store, { + userId, + parentConversationId: secondParentId, + }); + const created = store.start( + taskRequest(firstConfig.scopeId, { + parentAgentId: 'agent-a', + subagentType: 'self', + }), + ); + await waitForSettled(store, firstConfig.scopeId, created); + + const attempts = [ + { + config: firstConfig, + threadId: randomUUID(), + overrides: {}, + }, + { + config: secondConfig, + threadId: requireThreadId(created), + overrides: { parentAgentId: 'agent-a', subagentType: 'self' }, + }, + { + config: firstConfig, + threadId: requireThreadId(created), + overrides: { parentAgentId: 'agent-b', subagentType: 'self' }, + }, + { + config: firstConfig, + threadId: requireThreadId(created), + overrides: { + parentAgentId: 'agent-a', + subagentKind: 'graph' as const, + subagentType: 'self', + }, + }, + ]; + + for (const attempt of attempts) { + const run = jest.fn(taskRequest(attempt.config.scopeId).run); + const rejected = store.start( + taskRequest(attempt.config.scopeId, { + threadId: attempt.threadId, + run, + ...attempt.overrides, + }), + ); + await waitForSettled(store, attempt.config.scopeId, rejected); + expect(run).not.toHaveBeenCalled(); + expect( + store.claim(attempt.config.scopeId, requireAccepted(rejected).task.taskId), + ).toMatchObject({ status: 'error' }); + } + }); + + it('does not persist arbitrary executor details into the visible child chat', async () => { + loggerErrorSpy.mockClear(); + const userId = 'safe-error-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 () => { + throw new Error('Authorization: Bearer provider-secret'); + }, + }), + ); + await waitForSettled(store, config.scopeId, started); + + const messages = await methods.getMessages({ + user: userId, + conversationId: requireThreadId(started), + }); + expect(messages[messages.length - 1]?.text).toBe( + 'Subagent task failed: The child run could not be completed.', + ); + expect(JSON.stringify(messages)).not.toContain('provider-secret'); + expect( + JSON.stringify(store.claim(config.scopeId, requireAccepted(started).task.taskId)), + ).not.toContain('provider-secret'); + expect(JSON.stringify(loggerErrorSpy.mock.calls)).not.toContain('provider-secret'); + }); + + it('bounds durable delegation depth to one by default', async () => { + const userId = 'depth-user'; + const rootConversationId = randomUUID(); + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId, { + subagentThread: { + rootConversationId, + parentConversationId: rootConversationId, + parentMessageId: randomUUID(), + parentToolCallId: randomUUID(), + parentAgentId: 'root-agent', + subagentType: 'researcher-agent', + subagentKind: 'agent', + depth: 1, + }, + }); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const run = jest.fn(taskRequest(config.scopeId).run); + + expect(store.canCreateChildThread(0)).toBe(true); + expect(store.canCreateChildThread(1)).toBe(false); + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).not.toHaveBeenCalled(); + expect(store.claim(config.scopeId, requireAccepted(started).task.taskId)).toMatchObject({ + status: 'error', + }); + expect(await methods.getConvo(userId, requireThreadId(started))).toBeNull(); + }); +}); diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts new file mode 100644 index 0000000000..dc0f047706 --- /dev/null +++ b/packages/api/src/agents/subagentThreads.ts @@ -0,0 +1,1255 @@ +import { randomUUID } from 'node:crypto'; +import { logger } from '@librechat/data-schemas'; +import { InMemorySubagentTaskStore } from '@librechat/agents'; +import { EModelEndpoint, Constants } from 'librechat-data-provider'; +import { + mapChatMessagesToStoredMessages, + mapStoredMessagesToChatMessages, +} from '@librechat/agents/langchain/messages'; +import type { + InMemorySubagentTaskStoreOptions, + SubagentTaskConfig, + SubagentTaskControlCommand, + SubagentTaskControlResult, + SubagentTaskRuntime, + SubagentTaskStartRequest, + SubagentTaskStartResult, +} from '@librechat/agents'; +import type { + AllMethods, + IConversation, + IMessage, + MessageMethods, + ConversationMethods, +} from '@librechat/data-schemas'; +import type { BaseMessage, StoredMessage } from '@librechat/agents/langchain/messages'; +import type { UsageMetadata } from '~/stream/interfaces/IJobStore'; +import { createSubagentAttemptKey, createSubagentThreadId } from './subagentThreadIds'; +import { runWithDetachedSubagentUsage } from './subagentTaskContext'; +import { aggregateEmittedUsage } from './usage'; + +const SCOPE_VERSION = 1; +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; +const DEFAULT_OWNER_DRAIN_POLL_MS = 100; +const MAX_TRANSCRIPT_BYTES = 12 * 1024 * 1024; +const TRANSCRIPT_SELECT = + 'messageId parentMessageId text createdAt +subagentTranscript +subagentTask'; + +class SubagentThreadPublicError extends Error {} +class SubagentThreadDeletedError extends SubagentThreadPublicError {} + +type SubagentThreadMethods = Pick< + AllMethods, + | 'acquireSubagentThreadLease' + | 'countActiveSubagentThreadLeases' + | 'deleteConvos' + | 'deleteMessages' + | 'getConvo' + | 'getMessages' + | 'reserveSubagentThread' + | 'releaseSubagentThreadLease' + | 'renewSubagentThreadLease' + | 'saveConvo' + | 'saveMessage' +>; + +interface SubagentThreadScope { + version: typeof SCOPE_VERSION; + userId: string; + parentConversationId: string; + tenantId?: string; +} + +interface PreparedThread { + conversation: IConversation; + initialMessages: BaseMessage[]; + initialStoredMessages: StoredMessage[]; + attemptKey: string; + userMessageId?: string; + replay?: { + status: 'completed' | 'error' | 'cancelled'; + content: string; + }; +} + +type ThreadMessage = Pick< + IMessage, + 'messageId' | 'parentMessageId' | 'text' | 'createdAt' | 'subagentTranscript' | 'subagentTask' +>; + +interface TaskThreadLease { + idempotencyKey: string; + taskId: string; + running: boolean; + settling: boolean; + shared?: { + token: string; + lost: boolean; + heartbeat?: ReturnType; + heartbeatInFlight?: Promise; + }; +} + +export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStoreOptions { + maxThreadDepth?: number; + leaseTtlMs?: number; + leaseHeartbeatMs?: number; + ownerDrainTimeoutMs?: number; + ownerDrainPollMs?: number; + isOwnerActive?: (userId: string) => Promise; +} + +function positiveInteger(value: number | undefined, fallback: number): number { + return Number.isSafeInteger(value) && value != null && value > 0 ? value : fallback; +} + +function isNonEmptyString(value: unknown): value is string { + return typeof value === 'string' && value.trim() !== ''; +} + +function normalizedRequestFingerprint(request: SubagentTaskStartRequest): string | undefined { + const fingerprint = request.requestFingerprint?.trim(); + return fingerprint == null || fingerprint === '' ? undefined : fingerprint; +} + +function parseScope(scopeId: string): SubagentThreadScope { + let parsed: unknown; + try { + parsed = JSON.parse(scopeId); + } catch { + throw new Error('Invalid subagent thread scope.'); + } + if (parsed == null || typeof parsed !== 'object' || Array.isArray(parsed)) { + throw new Error('Invalid subagent thread scope.'); + } + const candidate = parsed as Partial; + if ( + candidate.version !== SCOPE_VERSION || + !isNonEmptyString(candidate.userId) || + !isNonEmptyString(candidate.parentConversationId) || + (candidate.tenantId != null && !isNonEmptyString(candidate.tenantId)) + ) { + throw new Error('Invalid subagent thread scope.'); + } + return { + version: SCOPE_VERSION, + userId: candidate.userId, + parentConversationId: candidate.parentConversationId, + ...(candidate.tenantId == null ? {} : { tenantId: candidate.tenantId }), + }; +} + +function matchesTenant(actual: string | undefined, expected: string | undefined): boolean { + return actual === expected; +} + +function assertParentPersistence( + value: unknown, + scope: SubagentThreadScope, +): asserts value is { message: { messageId: string; conversationId: string } } { + const message = + value != null && typeof value === 'object' + ? (value as { message?: unknown }).message + : undefined; + if ( + message == null || + typeof message !== 'object' || + !isNonEmptyString((message as { messageId?: unknown }).messageId) || + (message as { conversationId?: unknown }).conversationId !== scope.parentConversationId + ) { + throw new Error('The parent message was not persisted.'); + } +} + +function selectLatestBranch(messages: ThreadMessage[]): ThreadMessage[] { + const byId = new Map(messages.map((message) => [message.messageId, message])); + const branch: ThreadMessage[] = []; + const seen = new Set(); + let current: ThreadMessage | undefined = messages[messages.length - 1]; + while (current != null && !seen.has(current.messageId)) { + branch.push(current); + seen.add(current.messageId); + const parentId: string | undefined = current.parentMessageId ?? undefined; + current = + parentId == null || parentId === '' || parentId === Constants.NO_PARENT + ? undefined + : byId.get(parentId); + } + return branch.reverse(); +} + +function parseStoredMessages(value: string): StoredMessage[] { + const parsed = JSON.parse(value) as unknown; + if (!Array.isArray(parsed)) { + throw new SubagentThreadPublicError('Invalid persisted subagent transcript.'); + } + for (const message of parsed) { + if ( + message == null || + typeof message !== 'object' || + Array.isArray(message) || + !isNonEmptyString((message as { type?: unknown }).type) || + (message as { data?: unknown }).data == null || + typeof (message as { data?: unknown }).data !== 'object' || + Array.isArray((message as { data?: unknown }).data) + ) { + throw new SubagentThreadPublicError('Invalid persisted subagent transcript.'); + } + } + return parsed as StoredMessage[]; +} + +function restoreThreadMessages(branch: ThreadMessage[]): BaseMessage[] { + let storedMessages: StoredMessage[] = []; + for (const message of branch) { + const transcript = message.subagentTranscript; + if (transcript == null) { + continue; + } + const segment = parseStoredMessages(transcript.messagesJson); + if (transcript.mode === 'replace') { + storedMessages = segment; + } else { + storedMessages.push(...segment); + } + } + return mapStoredMessagesToChatMessages(storedMessages); +} + +function isStoredPrefix(prefix: StoredMessage[], messages: StoredMessage[]): boolean { + if (prefix.length > messages.length) { + return false; + } + for (let index = 0; index < prefix.length; index += 1) { + if (JSON.stringify(prefix[index]) !== JSON.stringify(messages[index])) { + return false; + } + } + return true; +} + +function serializeTranscript( + taskId: string, + initialMessages: StoredMessage[], + resultMessages: BaseMessage[] | undefined, +): IMessage['subagentTranscript'] { + if (resultMessages == null) { + return undefined; + } + const storedResult = mapChatMessagesToStoredMessages(resultMessages); + const storedResultJson = JSON.stringify(storedResult); + if (Buffer.byteLength(storedResultJson, 'utf8') > MAX_TRANSCRIPT_BYTES) { + throw new SubagentThreadPublicError( + 'Subagent thread transcript is too large to persist safely.', + ); + } + const append = isStoredPrefix(initialMessages, storedResult); + const messages = append ? storedResult.slice(initialMessages.length) : storedResult; + if (messages.length === 0) { + return undefined; + } + const messagesJson = append ? JSON.stringify(messages) : storedResultJson; + return { + taskId, + mode: append ? 'append' : 'replace', + messagesJson, + }; +} + +function retentionFields(conversation: IConversation): { + isTemporary?: boolean; + expiredAt?: Date; + tenantId?: string; +} { + return { + ...(conversation.isTemporary == null ? {} : { isTemporary: conversation.isTemporary }), + ...(conversation.expiredAt == null ? {} : { expiredAt: conversation.expiredAt }), + ...(conversation.tenantId == null ? {} : { tenantId: conversation.tenantId }), + }; +} + +function childAgentId(request: SubagentTaskStartRequest): string | undefined { + if (request.subagentKind === 'graph') { + return undefined; + } + return request.subagentType === 'self' ? request.parentAgentId : request.subagentType; +} + +function publicFailureDetail(error: unknown): string { + return error instanceof SubagentThreadPublicError + ? error.message.slice(0, 2_000) + : 'The child run could not be completed.'; +} + +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. */ +export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { + readonly supportsThreadContinuation = true; + private readonly activeThreads = new Map(); + 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 isOwnerActive: (userId: string) => Promise; + + constructor( + private readonly methods: SubagentThreadMethods, + options: SubagentThreadTaskStoreOptions = {}, + ) { + super(options); + this.maxThreadDepth = + Number.isSafeInteger(options.maxThreadDepth) && (options.maxThreadDepth ?? 0) > 0 + ? (options.maxThreadDepth as number) + : DEFAULT_MAX_THREAD_DEPTH; + this.leaseTtlMs = positiveInteger(options.leaseTtlMs, DEFAULT_LEASE_TTL_MS); + this.leaseHeartbeatMs = Math.min( + positiveInteger(options.leaseHeartbeatMs, DEFAULT_LEASE_HEARTBEAT_MS), + Math.max(1, Math.floor(this.leaseTtlMs / 2)), + ); + this.ownerDrainTimeoutMs = positiveInteger( + options.ownerDrainTimeoutMs, + DEFAULT_OWNER_DRAIN_TIMEOUT_MS, + ); + this.ownerDrainPollMs = positiveInteger(options.ownerDrainPollMs, DEFAULT_OWNER_DRAIN_POLL_MS); + this.isOwnerActive = options.isOwnerActive ?? (async () => true); + } + + /** Gates child creation on the ordinary parent write without retaining request state. */ + registerParentPersistence(scopeId: string, persistence: Promise): void { + const scope = parseScope(scopeId); + const gate = Promise.resolve(persistence).then((result) => { + assertParentPersistence(result, scope); + return result; + }); + this.parentPersistence.set(scopeId, gate); + void gate + .then(() => { + if (this.parentPersistence.get(scopeId) === gate) { + this.parentPersistence.delete(scopeId); + } + }) + .catch(() => undefined); + } + + override start(request: SubagentTaskStartRequest): SubagentTaskStartResult { + if (request.subagentKind !== 'agent' && request.subagentKind !== 'graph') { + throw new Error('Subagent task kind must be agent or graph.'); + } + const scope = parseScope(request.scopeId); + const parentReady = this.parentPersistence.get(request.scopeId); + const requestedThreadId = request.threadId?.trim(); + const isContinuation = requestedThreadId != null && requestedThreadId !== ''; + const idempotencyKey = request.idempotencyKey.trim(); + const threadId = isContinuation + ? requestedThreadId + : createSubagentThreadId(request.scopeId, idempotencyKey); + const lockKey = `${request.scopeId}\u0000${threadId}`; + const active = this.activeThreads.get(lockKey); + if (active != null && active.idempotencyKey !== idempotencyKey) { + return { accepted: false, reason: 'capacity' }; + } + + const lease: TaskThreadLease = active ?? { + idempotencyKey, + taskId: '', + running: false, + settling: false, + }; + const ownsLease = active == null; + if (ownsLease) { + this.activeThreads.set(lockKey, lease); + } + + let started: SubagentTaskStartResult; + try { + started = super.start({ + ...request, + threadId, + run: async (runtime: SubagentTaskRuntime) => { + lease.taskId = runtime.taskId; + lease.running = true; + const detachedUsage: UsageMetadata[] = []; + try { + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + await parentReady; + const prepared = await this.prepareThread( + request.scopeId, + scope, + threadId, + isContinuation, + request, + runtime.taskId, + lease, + ); + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + if (prepared.replay != null) { + if (prepared.replay.status === 'completed') { + return { content: prepared.replay.content }; + } + throw new SubagentThreadPublicError(prepared.replay.content); + } + if (!(await this.renewSharedLease(scope, threadId, lease))) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + const result = await runWithDetachedSubagentUsage(detachedUsage, () => + request.run(runtime, prepared.initialMessages), + ); + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + if (!(await this.renewSharedLease(scope, threadId, lease))) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + lease.settling = true; + await this.persistResult( + scope, + request, + runtime.taskId, + prepared, + result, + detachedUsage, + ); + return result; + } catch (error) { + const mayPersist = + lease.shared == null || (await this.renewSharedLease(scope, threadId, lease)); + const terminalTask = this.get(request.scopeId, runtime.taskId); + if (runtime.signal.aborted && terminalTask?.status === 'cancelled') { + if (mayPersist) { + await this.persistCancellation( + scope, + threadId, + request, + runtime.taskId, + detachedUsage, + ).catch((persistError) => { + logger.error( + '[subagentThreads] Failed to persist child-thread cancellation', + persistError, + ); + }); + } + throw error; + } + logger.error( + '[subagentThreads] Child-thread execution failed', + publicFailureDetail(error), + ); + if (mayPersist) { + await this.persistFailure( + scope, + threadId, + request, + runtime.taskId, + error, + detachedUsage, + ).catch((persistError) => { + logger.error( + '[subagentThreads] Failed to persist child-thread failure', + persistError, + ); + }); + } + throw new Error(publicFailureDetail(error)); + } finally { + await this.stopAndReleaseSharedLease(scope, threadId, lease); + if (this.activeThreads.get(lockKey) === lease) { + this.activeThreads.delete(lockKey); + } + } + }, + }); + } catch (error) { + if (ownsLease && this.activeThreads.get(lockKey) === lease) { + this.activeThreads.delete(lockKey); + } + throw error; + } + + if (started.accepted && started.isNew) { + lease.taskId = started.task.taskId; + } else if (ownsLease && this.activeThreads.get(lockKey) === lease) { + this.activeThreads.delete(lockKey); + } + return started; + } + + override control( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + ): SubagentTaskControlResult { + const snapshot = this.get(scopeId, taskId); + const lockKey = snapshot?.threadId == null ? undefined : `${scopeId}\u0000${snapshot.threadId}`; + const lease = lockKey == null ? undefined : this.activeThreads.get(lockKey); + if (lease?.taskId === taskId && lease.settling) { + return { status: 'not_running', task: snapshot }; + } + const result = super.control(scopeId, taskId, command); + if ( + command.action === 'cancel' && + result.status === 'cancelled' && + lockKey != null && + lease?.taskId === taskId && + !lease.running + ) { + this.activeThreads.delete(lockKey); + } + return result; + } + + /** Finds a provisional child lease before its conversation is durable. */ + isThreadActiveForOwner(userId: string, threadId: string, tenantId?: string): boolean { + const suffix = `\u0000${threadId}`; + for (const lockKey of this.activeThreads.keys()) { + if (!lockKey.endsWith(suffix)) { + continue; + } + const scope = parseScope(lockKey.slice(0, -suffix.length)); + if (scope.userId === userId && matchesTenant(scope.tenantId, tenantId)) { + return true; + } + } + return false; + } + + /** Cancels active descendants before their owning conversations are removed. */ + cancelForConversations( + userId: string, + conversationIds: Iterable, + tenantId?: string, + ): number { + const targets = new Set(conversationIds); + return this.cancelMatchingThreads( + (scope, threadId) => + scope.userId === userId && + matchesTenant(scope.tenantId, tenantId) && + (targets.has(scope.parentConversationId) || 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( + (scope) => scope.userId === userId && matchesTenant(scope.tenantId, tenantId), + ); + } + + /** Cancels local work and waits for every replica's durable lease to drain. */ + async cancelAndDrainForOwner(userId: string, tenantId?: string): Promise { + this.cancelForOwner(userId, tenantId); + const deadline = Date.now() + this.ownerDrainTimeoutMs; + while (true) { + const active = await this.methods.countActiveSubagentThreadLeases({ + user: userId, + now: new Date(), + ...(tenantId == null ? {} : { tenantId }), + }); + if (active === 0) { + return; + } + if (Date.now() >= deadline) { + throw new Error('Timed out draining detached subagent tasks for account deletion.'); + } + await new Promise((resolve) => setTimeout(resolve, this.ownerDrainPollMs)); + } + } + + private startSharedLeaseHeartbeat( + scopeId: string, + scope: SubagentThreadScope, + threadId: string, + lease: TaskThreadLease, + ): void { + const shared = lease.shared; + if (shared == null) { + return; + } + const heartbeat = () => { + if (shared.lost || shared.heartbeatInFlight != null) { + return; + } + const renewal = (async () => { + let ownerActive = false; + try { + ownerActive = await this.isOwnerActive(scope.userId); + } catch (error) { + logger.warn('[subagentThreads] Failed to verify the child-thread owner', error); + } + if (!ownerActive) { + const task = this.get(scopeId, lease.taskId); + if (task?.status === 'running') { + super.control(scopeId, lease.taskId, { action: 'cancel' }); + } + } + if (!(await this.renewSharedLeaseFence(scope, threadId, lease))) { + const task = this.get(scopeId, lease.taskId); + if (task?.status === 'running') { + super.control(scopeId, lease.taskId, { action: 'cancel' }); + } + } + })().finally(() => { + if (shared.heartbeatInFlight === renewal) { + shared.heartbeatInFlight = undefined; + } + }); + shared.heartbeatInFlight = renewal; + }; + shared.heartbeat = setInterval(heartbeat, this.leaseHeartbeatMs); + shared.heartbeat.unref?.(); + } + + private async renewSharedLease( + scope: SubagentThreadScope, + threadId: string, + lease: TaskThreadLease, + ): Promise { + const shared = lease.shared; + if (shared == null || shared.lost) { + return false; + } + try { + if (!(await this.isOwnerActive(scope.userId))) { + return false; + } + return this.renewSharedLeaseFence(scope, threadId, lease); + } catch (error) { + logger.warn('[subagentThreads] Failed to verify the child-thread owner', error); + return false; + } + } + + private async renewSharedLeaseFence( + scope: SubagentThreadScope, + threadId: string, + lease: TaskThreadLease, + ): Promise { + const shared = lease.shared; + if (shared == null || shared.lost) { + return false; + } + try { + const now = new Date(); + const renewed = await this.methods.renewSubagentThreadLease({ + user: scope.userId, + conversationId: threadId, + token: shared.token, + now, + expiresAt: new Date(now.getTime() + this.leaseTtlMs), + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + }); + if (!renewed) { + shared.lost = true; + } + return renewed; + } catch (error) { + shared.lost = true; + logger.warn('[subagentThreads] Lost the shared child-thread lease', error); + return false; + } + } + + private async stopAndReleaseSharedLease( + scope: SubagentThreadScope, + threadId: string, + lease: TaskThreadLease, + ): Promise { + const shared = lease.shared; + if (shared == null) { + return; + } + if (shared.heartbeat != null) { + clearInterval(shared.heartbeat); + } + await shared.heartbeatInFlight; + try { + await this.methods.releaseSubagentThreadLease({ + user: scope.userId, + conversationId: threadId, + token: shared.token, + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + }); + } catch (error) { + logger.warn('[subagentThreads] Failed to release the shared child-thread lease', error); + } + } + + /** Whether a durable child may be created below the supplied conversation depth. */ + canCreateChildThread(parentDepth: number): boolean { + return ( + Number.isSafeInteger(parentDepth) && parentDepth >= 0 && parentDepth < this.maxThreadDepth + ); + } + + private async prepareThread( + scopeId: string, + scope: SubagentThreadScope, + threadId: string, + isContinuation: boolean, + request: SubagentTaskStartRequest, + taskId: string, + lease: TaskThreadLease, + ): Promise { + if (!(await this.isOwnerActive(scope.userId))) { + throw new SubagentThreadPublicError('The thread owner is unavailable.'); + } + const [parent, existing] = await Promise.all([ + this.methods.getConvo(scope.userId, scope.parentConversationId), + this.methods.getConvo(scope.userId, threadId), + ]); + if (parent == null || !matchesTenant(parent.tenantId, scope.tenantId)) { + throw new SubagentThreadPublicError('Parent thread is unavailable.'); + } + + let createdThread = false; + let sharedLeaseAcquired = false; + try { + let conversation = existing; + if (conversation == null && isContinuation) { + throw new SubagentThreadPublicError( + 'Child thread is unavailable for this subagent and parent scope.', + ); + } + if (conversation == null) { + const parentDepth = parent.subagentThread?.depth ?? 0; + if (!this.canCreateChildThread(parentDepth)) { + throw new SubagentThreadPublicError( + `Subagent thread depth exceeds the configured limit of ${this.maxThreadDepth}.`, + ); + } + const depth = parentDepth + 1; + const agentId = childAgentId(request); + const reserved = await this.methods.reserveSubagentThread({ + user: scope.userId, + conversationId: threadId, + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + conversation: { + conversationId: threadId, + endpoint: EModelEndpoint.agents, + title: `Subagent: ${request.subagentType}`.slice(0, 120), + ...(agentId == null ? {} : { agent_id: agentId }), + ...retentionFields(parent), + subagentThread: { + rootConversationId: + parent.subagentThread?.rootConversationId ?? scope.parentConversationId, + parentConversationId: scope.parentConversationId, + parentMessageId: request.parentRunId || request.parentToolCallId, + parentToolCallId: request.parentToolCallId, + ...(request.parentAgentId == null ? {} : { parentAgentId: request.parentAgentId }), + subagentType: request.subagentType, + subagentKind: request.subagentKind as 'agent' | 'graph', + depth, + }, + }, + }); + conversation = reserved.conversation; + createdThread = reserved.created; + } + + this.assertContinuation(scope, request, conversation); + const sharedToken = randomUUID(); + const now = new Date(); + sharedLeaseAcquired = await this.methods.acquireSubagentThreadLease({ + user: scope.userId, + conversationId: threadId, + token: sharedToken, + taskId, + now, + expiresAt: new Date(now.getTime() + this.leaseTtlMs), + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + }); + if (!sharedLeaseAcquired) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + lease.shared = { token: sharedToken, lost: false }; + this.startSharedLeaseHeartbeat(scopeId, scope, threadId, lease); + /** Account deletion can fence the owner after the optimistic probe but before + * this lease exists. Once the lease is visible, revalidate so deletion either + * observes and drains us or wins before any provider work can begin. */ + if (!(await this.isOwnerActive(scope.userId))) { + throw new SubagentThreadDeletedError('The thread owner is unavailable.'); + } + const allMessages = (await this.methods.getMessages( + { conversationId: threadId, user: scope.userId }, + TRANSCRIPT_SELECT, + { sort: { createdAt: 1, _id: 1 } }, + )) as ThreadMessage[]; + const attemptKey = createSubagentAttemptKey(scopeId, request.idempotencyKey); + const requestFingerprint = normalizedRequestFingerprint(request); + const priorAttempt = allMessages.filter( + (message) => message.subagentTask?.attemptKey === attemptKey, + ); + if (priorAttempt.length > 0) { + if ( + priorAttempt.some( + (message) => message.subagentTask?.requestFingerprint !== requestFingerprint, + ) + ) { + throw new SubagentThreadPublicError( + 'The same parent tool call was already used with different subagent arguments.', + ); + } + const terminal = [...priorAttempt] + .reverse() + .find((message) => message.subagentTask?.status !== 'running'); + if (terminal?.subagentTask != null) { + return { + conversation, + initialMessages: [], + initialStoredMessages: [], + attemptKey, + replay: { + status: terminal.subagentTask.status as 'completed' | 'error' | 'cancelled', + content: + terminal.text ?? + (terminal.subagentTask.status === 'completed' + ? 'Subagent task completed.' + : 'The prior subagent task did not complete successfully.'), + }, + }; + } + /** Reaching this point while holding the thread lease proves the original + * worker no longer owns settlement. Close the abandoned attempt once rather + * than either re-billing it or leaving every retry permanently "running". */ + const abandoned = priorAttempt[priorAttempt.length - 1]; + const abandonedMessage = + 'Subagent task failed: The prior execution ended before its result could be persisted.'; + const savedAbandoned = await this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: `${taskId}:assistant`, + conversationId: threadId, + parentMessageId: abandoned.messageId, + sender: request.subagentType, + text: abandonedMessage, + endpoint: EModelEndpoint.agents, + isCreatedByUser: false, + unfinished: false, + error: true, + subagentTask: { + attemptKey, + ...(requestFingerprint == null ? {} : { requestFingerprint }), + status: 'error', + }, + ...retentionFields(conversation), + }, + { context: 'SubagentThreadTaskStore.prepareThread.abandonedAttempt' }, + ); + if (savedAbandoned == null) { + throw new Error('Unable to close the abandoned subagent attempt.'); + } + await this.touchAfterMessage(scope, threadId, taskId, 'failed'); + return { + conversation, + initialMessages: [], + initialStoredMessages: [], + attemptKey, + replay: { status: 'error', content: abandonedMessage }, + }; + } + const branch = selectLatestBranch(allMessages); + const initialMessages = restoreThreadMessages(branch); + const userMessageId = `${taskId}:user`; + /** A crashed lease can leave its input row behind. Continue from the latest + * terminal task row instead of making that incomplete input canonical. */ + let parentMessageId: string = Constants.NO_PARENT; + for (let index = branch.length - 1; index >= 0; index -= 1) { + if (branch[index].messageId.endsWith(':assistant')) { + parentMessageId = branch[index].messageId; + break; + } + } + const savedUserMessage = await this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: userMessageId, + conversationId: threadId, + parentMessageId, + sender: 'User', + text: request.input, + endpoint: EModelEndpoint.agents, + isCreatedByUser: true, + subagentTask: { + attemptKey, + ...(requestFingerprint == null ? {} : { requestFingerprint }), + status: 'running', + }, + ...retentionFields(conversation), + }, + { context: 'SubagentThreadTaskStore.prepareThread' }, + ); + if (savedUserMessage == null) { + throw new Error('Unable to persist the child-thread input.'); + } + const currentParent = await this.methods.getConvo(scope.userId, scope.parentConversationId); + if (currentParent == null || !matchesTenant(currentParent.tenantId, scope.tenantId)) { + throw new SubagentThreadPublicError('Parent thread is unavailable.'); + } + return { + conversation, + initialMessages, + initialStoredMessages: mapChatMessagesToStoredMessages(initialMessages), + attemptKey, + userMessageId, + }; + } catch (error) { + if (sharedLeaseAcquired) { + await this.rollbackPreparation(scope, threadId, taskId, createdThread); + } + throw error; + } + } + + private assertContinuation( + scope: SubagentThreadScope, + request: SubagentTaskStartRequest, + conversation: IConversation, + ): void { + if (!this.isContinuationAllowed(scope, request, conversation)) { + throw new SubagentThreadPublicError( + 'Child thread is unavailable for this subagent and parent scope.', + ); + } + } + + private isContinuationAllowed( + scope: SubagentThreadScope, + request: SubagentTaskStartRequest, + conversation: IConversation, + ): boolean { + const lineage = conversation.subagentThread; + const expectedAgentId = childAgentId(request); + const agentIdentityMatches = + request.subagentKind === 'graph' + ? conversation.agent_id == null + : conversation.agent_id === expectedAgentId; + return !( + lineage == null || + conversation.endpoint !== EModelEndpoint.agents || + !agentIdentityMatches || + lineage.parentConversationId !== scope.parentConversationId || + lineage.parentAgentId !== request.parentAgentId || + lineage.subagentType !== request.subagentType || + lineage.subagentKind !== request.subagentKind || + lineage.depth > this.maxThreadDepth || + !matchesTenant(conversation.tenantId, scope.tenantId) + ); + } + + private async persistResult( + scope: SubagentThreadScope, + request: SubagentTaskStartRequest, + taskId: string, + prepared: PreparedThread, + result: { content: string; messages?: BaseMessage[] }, + detachedUsage: UsageMetadata[], + ): Promise { + if (prepared.userMessageId == null) { + throw new Error('The child-thread input was not prepared.'); + } + const subagentTranscript = serializeTranscript( + taskId, + prepared.initialStoredMessages, + result.messages, + ); + const conversation = await this.requireCurrentConversation( + scope, + request, + prepared.conversation.conversationId, + ); + const usage = this.aggregateDetachedUsage(detachedUsage); + const savedAssistantMessage = await this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: `${taskId}:assistant`, + conversationId: conversation.conversationId, + parentMessageId: prepared.userMessageId, + sender: request.subagentType, + text: result.content, + endpoint: EModelEndpoint.agents, + isCreatedByUser: false, + unfinished: false, + ...(subagentTranscript == null ? {} : { subagentTranscript }), + subagentTask: { + attemptKey: prepared.attemptKey, + ...(normalizedRequestFingerprint(request) == null + ? {} + : { requestFingerprint: normalizedRequestFingerprint(request) }), + status: 'completed', + }, + ...(usage == null ? {} : { metadata: { usage } }), + ...retentionFields(conversation), + }, + { context: 'SubagentThreadTaskStore.persistResult' }, + ); + if (savedAssistantMessage == null) { + throw new Error('Unable to persist the child-thread result.'); + } + await this.touchAfterMessage(scope, conversation.conversationId, taskId, 'completed'); + } + + private async persistFailure( + scope: SubagentThreadScope, + threadId: string, + request: SubagentTaskStartRequest, + taskId: string, + error: unknown, + detachedUsage: UsageMetadata[], + ): Promise { + const conversation = await this.currentConversation(scope, request, threadId); + if (conversation == null || !(await this.taskInputExists(scope, threadId, taskId))) { + return; + } + const usage = this.aggregateDetachedUsage(detachedUsage); + const savedFailure = await this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: `${taskId}:assistant`, + conversationId: threadId, + parentMessageId: `${taskId}:user`, + sender: request.subagentType, + text: safeErrorMessage(error), + endpoint: EModelEndpoint.agents, + isCreatedByUser: false, + unfinished: false, + error: true, + subagentTask: { + attemptKey: createSubagentAttemptKey(request.scopeId, request.idempotencyKey), + ...(normalizedRequestFingerprint(request) == null + ? {} + : { requestFingerprint: normalizedRequestFingerprint(request) }), + status: 'error', + }, + ...(usage == null ? {} : { metadata: { usage } }), + ...retentionFields(conversation), + }, + { context: 'SubagentThreadTaskStore.persistFailure' }, + ); + if (savedFailure == null) { + throw new Error('Unable to persist the child-thread failure.'); + } + await this.touchAfterMessage(scope, threadId, taskId, 'failed'); + } + + private async persistCancellation( + scope: SubagentThreadScope, + threadId: string, + request: SubagentTaskStartRequest, + taskId: string, + detachedUsage: UsageMetadata[], + ): Promise { + const conversation = await this.currentConversation(scope, request, threadId); + if (conversation == null || !(await this.taskInputExists(scope, threadId, taskId))) { + return; + } + const usage = this.aggregateDetachedUsage(detachedUsage); + const savedCancellation = await this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: `${taskId}:assistant`, + conversationId: threadId, + parentMessageId: `${taskId}:user`, + sender: request.subagentType, + text: 'Subagent task was cancelled.', + endpoint: EModelEndpoint.agents, + isCreatedByUser: false, + unfinished: false, + subagentTask: { + attemptKey: createSubagentAttemptKey(request.scopeId, request.idempotencyKey), + ...(normalizedRequestFingerprint(request) == null + ? {} + : { requestFingerprint: normalizedRequestFingerprint(request) }), + status: 'cancelled', + }, + ...(usage == null ? {} : { metadata: { usage } }), + ...retentionFields(conversation), + }, + { context: 'SubagentThreadTaskStore.persistCancellation' }, + ); + if (savedCancellation == null) { + throw new Error('Unable to persist the child-thread cancellation.'); + } + await this.touchAfterMessage(scope, threadId, taskId, 'cancelled'); + } + + private async currentConversation( + scope: SubagentThreadScope, + request: SubagentTaskStartRequest, + threadId: string, + ): Promise { + const [ownerActive, parent, conversation] = await Promise.all([ + this.isOwnerActive(scope.userId), + this.methods.getConvo(scope.userId, scope.parentConversationId), + this.methods.getConvo(scope.userId, threadId), + ]); + if (conversation == null || !this.isContinuationAllowed(scope, request, conversation)) { + return null; + } + if (!ownerActive) { + return null; + } + if (parent != null && matchesTenant(parent.tenantId, scope.tenantId)) { + return conversation; + } + await this.methods.deleteConvos(scope.userId, { conversationId: threadId }).catch((error) => { + logger.warn('[subagentThreads] Failed to remove an orphaned child thread', error); + }); + return null; + } + + private cancelMatchingThreads( + matches: (scope: SubagentThreadScope, threadId: string) => boolean, + ): number { + let cancelled = 0; + for (const [lockKey, lease] of this.activeThreads) { + const separator = lockKey.lastIndexOf('\u0000'); + if (separator < 0 || lease.taskId === '') { + continue; + } + const scopeId = lockKey.slice(0, separator); + const threadId = lockKey.slice(separator + 1); + const scope = parseScope(scopeId); + if (!matches(scope, threadId)) { + continue; + } + const result = this.control(scopeId, lease.taskId, { action: 'cancel' }); + if (result.status === 'cancelled') { + cancelled += 1; + } + } + return cancelled; + } + + private async requireCurrentConversation( + scope: SubagentThreadScope, + request: SubagentTaskStartRequest, + threadId: string, + ): Promise { + const conversation = await this.currentConversation(scope, request, threadId); + if (conversation == null) { + throw new SubagentThreadDeletedError('Child thread was deleted before settlement.'); + } + return conversation; + } + + private async taskInputExists( + scope: SubagentThreadScope, + threadId: string, + taskId: string, + ): Promise { + const messages = await this.methods.getMessages( + { conversationId: threadId, user: scope.userId, messageId: `${taskId}:user` }, + 'messageId', + { limit: 1 }, + ); + return messages.length > 0; + } + + private async rollbackPreparation( + scope: SubagentThreadScope, + threadId: string, + taskId: string, + createdThread: boolean, + ): Promise { + try { + if (createdThread) { + await this.methods.deleteConvos(scope.userId, { conversationId: threadId }); + return; + } + await this.deleteTaskMessages(scope, threadId, taskId); + } catch (cleanupError) { + logger.error('[subagentThreads] Failed to roll back child-thread setup', cleanupError); + } + } + + private async deleteTaskMessages( + scope: SubagentThreadScope, + threadId: string, + taskId: string, + ): Promise { + await this.methods.deleteMessages({ + user: scope.userId, + conversationId: threadId, + messageId: { $in: [`${taskId}:user`, `${taskId}:assistant`] }, + }); + } + + private async touchAfterMessage( + scope: SubagentThreadScope, + threadId: string, + taskId: string, + outcome: 'cancelled' | 'completed' | 'failed', + ): Promise { + try { + const saved = await this.methods.saveConvo( + { userId: scope.userId }, + { conversationId: threadId }, + { context: 'SubagentThreadTaskStore.touchAfterMessage', noUpsert: true }, + ); + if (saved == null) { + throw new SubagentThreadDeletedError('Child thread was deleted before settlement.'); + } + if ('message' in saved) { + throw new Error('Unable to refresh the child thread.'); + } + } catch (error) { + if (error instanceof SubagentThreadDeletedError) { + await this.deleteTaskMessages(scope, threadId, taskId); + throw error; + } + logger.error(`[subagentThreads] Failed to refresh ${outcome} child thread`, error); + } + } + + private aggregateDetachedUsage(detachedUsage: UsageMetadata[]) { + return aggregateEmittedUsage( + detachedUsage.map((entry) => ({ ...entry, usage_type: 'subagent' as const })), + ); + } +} + +export function createSubagentThreadTaskStore( + methods: Pick< + ConversationMethods, + | 'acquireSubagentThreadLease' + | 'countActiveSubagentThreadLeases' + | 'deleteConvos' + | 'getConvo' + | 'releaseSubagentThreadLease' + | 'reserveSubagentThread' + | 'renewSubagentThreadLease' + | 'saveConvo' + > & + Pick, + options?: SubagentThreadTaskStoreOptions, +): SubagentThreadTaskStore { + return new SubagentThreadTaskStore(methods, options); +} + +export function buildSubagentThreadTaskConfig( + store: SubagentThreadTaskStore, + scope: Omit, +): SubagentTaskConfig { + return { + store, + scopeId: JSON.stringify({ version: SCOPE_VERSION, ...scope }), + }; +} diff --git a/packages/api/src/agents/usage.spec.ts b/packages/api/src/agents/usage.spec.ts index 4c0992e27e..a9a5d070cd 100644 --- a/packages/api/src/agents/usage.spec.ts +++ b/packages/api/src/agents/usage.spec.ts @@ -5,6 +5,7 @@ import type { BulkWriteDeps, PricingFns } from './transactions'; import { computeUsageCostUSD, aggregateEmittedUsage, + createDetachedSubagentUsageRecorder, createSubagentUsageSink, recordCollectedUsage, resolveAgentTokenConfig, @@ -13,6 +14,7 @@ import { computeSummaryUsedTokens, priorRunOutputTokens, } from './usage'; +import { runWithDetachedSubagentUsage } from './subagentTaskContext'; describe('recordCollectedUsage', () => { let mockSpendTokens: jest.Mock; @@ -1472,7 +1474,9 @@ describe('createSubagentUsageSink', () => { it('tags the child agent id so the host can price with the subagent endpoint config', () => { const collectedUsage: UsageMetadata[] = []; const emitted: UsageMetadata[] = []; - const sink = createSubagentUsageSink(collectedUsage, (u) => emitted.push(u)); + const sink = createSubagentUsageSink(collectedUsage, (u) => { + emitted.push(u); + }); sink(makeEvent({ subagentAgentId: 'agent_xyz' })); @@ -1548,6 +1552,55 @@ describe('createSubagentUsageSink', () => { expect(collectedUsage).toEqual([]); }); + it('routes detached usage to its awaited billing and durable child collectors', async () => { + const collectedUsage: UsageMetadata[] = []; + const detachedUsage: UsageMetadata[] = []; + const emitted: UsageMetadata[] = []; + const recordDetachedUsage = jest.fn().mockResolvedValue(undefined); + const sink = createSubagentUsageSink( + collectedUsage, + (usage) => { + emitted.push(usage); + }, + recordDetachedUsage, + ); + + await runWithDetachedSubagentUsage(detachedUsage, async () => { + await sink(makeEvent()); + }); + + expect(collectedUsage).toEqual([]); + expect(detachedUsage).toHaveLength(1); + expect(emitted[0]).toBe(detachedUsage[0]); + expect(recordDetachedUsage).toHaveBeenCalledWith(detachedUsage[0]); + + /** The same sink still batches ordinary foreground subagents with the parent. */ + await sink(makeEvent({ subagentRunId: 'foreground-child' })); + expect(collectedUsage).toHaveLength(1); + expect(recordDetachedUsage).toHaveBeenCalledTimes(1); + }); + + it('still records detached usage when the auxiliary emitter throws', async () => { + const collectedUsage: UsageMetadata[] = []; + const detachedUsage: UsageMetadata[] = []; + const recordDetachedUsage = jest.fn().mockResolvedValue(undefined); + const sink = createSubagentUsageSink( + collectedUsage, + () => { + throw new Error('parent transport was disposed'); + }, + recordDetachedUsage, + ); + + await runWithDetachedSubagentUsage(detachedUsage, async () => { + await sink(makeEvent()); + }); + + expect(collectedUsage).toEqual([]); + expect(detachedUsage).toHaveLength(1); + expect(recordDetachedUsage).toHaveBeenCalledWith(detachedUsage[0]); + }); + it('round-trips into recordCollectedUsage as billed subagent transactions', async () => { const collectedUsage: UsageMetadata[] = []; const sink = createSubagentUsageSink(collectedUsage); @@ -2160,3 +2213,78 @@ describe('resolveAgentTokenConfig', () => { expect(resolveAgentTokenConfig({ agentId: 'primary', fallback: primary })).toBe(primary); }); }); + +describe('createDetachedSubagentUsageRecorder', () => { + it('snapshots per-agent pricing and records each call as subagent usage', async () => { + const spendTokens = jest.fn().mockResolvedValue(undefined); + const childConfig = { 'child-model': { prompt: 0.01, completion: 0.02, context: 4096 } }; + const configs = new Map([['child-agent', childConfig]]); + const recorder = createDetachedSubagentUsageRecorder( + { + spendTokens, + spendStructuredTokens: jest.fn().mockResolvedValue(undefined), + }, + { + user: 'user-1', + conversationId: 'parent-1', + messageId: 'response-1', + model: 'parent-model', + endpointTokenConfigByAgentId: configs, + }, + ); + configs.set('child-agent', { + 'child-model': { prompt: 99, completion: 99, context: 4096 }, + }); + + await recorder({ + usage_type: 'subagent', + input_tokens: 12, + output_tokens: 4, + model: 'child-model', + agentId: 'child-agent', + }); + + expect(spendTokens).toHaveBeenCalledWith( + expect.objectContaining({ + user: 'user-1', + conversationId: 'parent-1', + messageId: 'response-1', + context: 'subagent', + model: 'child-model', + endpointTokenConfig: childConfig, + }), + { promptTokens: 12, completionTokens: 4 }, + ); + }); + + it('does not recreate billing records after the owning principal is fenced or deleted', async () => { + const spendTokens = jest.fn().mockResolvedValue(undefined); + const updateBalance = jest.fn().mockResolvedValue(undefined); + const insertMany = jest.fn().mockResolvedValue(undefined); + const recorder = createDetachedSubagentUsageRecorder( + { + spendTokens, + spendStructuredTokens: jest.fn().mockResolvedValue(undefined), + bulkWriteOps: { updateBalance, insertMany }, + isPrincipalActive: jest.fn().mockResolvedValue(false), + }, + { + user: 'deleted-user', + conversationId: 'parent-1', + messageId: 'response-1', + model: 'child-model', + }, + ); + + await recorder({ + usage_type: 'subagent', + input_tokens: 12, + output_tokens: 4, + model: 'child-model', + }); + + expect(spendTokens).not.toHaveBeenCalled(); + expect(updateBalance).not.toHaveBeenCalled(); + expect(insertMany).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/api/src/agents/usage.ts b/packages/api/src/agents/usage.ts index 935c356a0a..941244d1f7 100644 --- a/packages/api/src/agents/usage.ts +++ b/packages/api/src/agents/usage.ts @@ -27,6 +27,7 @@ import { bulkWriteTransactions, prepareTokenSpend, } from './transactions'; +import { collectDetachedSubagentUsage } from './subagentTaskContext'; type SpendTokensFn = (txData: TxMetadata, tokenUsage: TokenUsage) => Promise; type SpendStructuredTokensFn = ( @@ -136,6 +137,7 @@ export interface RecordUsageDeps { spendStructuredTokens: SpendStructuredTokensFn; pricing?: PricingFns; bulkWriteOps?: BulkWriteDeps; + isPrincipalActive?: (userId: string) => Promise; } /** @@ -514,6 +516,17 @@ export interface RecordUsageResult { output_tokens: number; } +export interface DetachedSubagentUsageRecorderParams { + user: string; + conversationId: string; + model?: string; + messageId?: string; + balance?: Partial | null; + transactions?: Partial; + endpointTokenConfig?: EndpointTokenConfig; + endpointTokenConfigByAgentId?: Map; +} + /** * Records token usage for collected LLM calls and spends tokens against balance. * This handles both sequential execution (tool calls) and parallel execution (multiple agents). @@ -690,6 +703,50 @@ export async function recordCollectedUsage( }; } +/** + * Creates an immutable, request-independent billing adapter for detached child + * calls. The recorder owns pricing selection and failure isolation so legacy + * controllers only provide their database dependencies and request snapshot. + */ +export function createDetachedSubagentUsageRecorder( + deps: RecordUsageDeps, + params: DetachedSubagentUsageRecorderParams, +): (usage: UsageMetadata) => Promise { + const billing = { + ...params, + endpointTokenConfigByAgentId: + params.endpointTokenConfigByAgentId == null + ? undefined + : new Map(params.endpointTokenConfigByAgentId), + }; + return async (usage) => { + try { + if (deps.isPrincipalActive != null && !(await deps.isPrincipalActive(billing.user))) { + return; + } + await recordCollectedUsage(deps, { + user: billing.user, + conversationId: billing.conversationId, + collectedUsage: [usage], + model: billing.model, + context: 'subagent', + messageId: billing.messageId, + balance: billing.balance, + transactions: billing.transactions, + endpointTokenConfig: billing.endpointTokenConfig, + resolveEndpointTokenConfig: (entry) => + resolveAgentTokenConfig({ + agentId: entry.agentId, + byAgentId: billing.endpointTokenConfigByAgentId, + fallback: billing.endpointTokenConfig, + }), + }); + } catch (error) { + logger.error('[agents/usage] Failed to record detached subagent usage', error); + } + }; +} + /** SDK-owned usage envelope re-exported for host billing consumers. */ export type SubagentUsageEvent = AgentsSubagentUsageEvent; @@ -698,15 +755,17 @@ export type SubagentUsageEvent = AgentsSubagentUsageEvent; * graphs execute outside the run's `streamEvents` loop, so their model calls * never reach the `CHAT_MODEL_END` handler (`ModelEndHandler`) — the SDK * reports them through this sink instead. Each event is tagged - * `usage_type: 'subagent'` with the child's model/provider and pushed onto - * the same `collectedUsage` array the handler fills, so - * {@link recordCollectedUsage} bills child calls (transactions + balance) - * alongside the parent's. + * `usage_type: 'subagent'` with the child's model/provider. Foreground child + * calls join the parent `collectedUsage` batch. Detached calls are recognized + * through their task-local context, billed immediately through + * `recordDetachedUsage`, and persisted with the durable child result instead + * of depending on a parent turn that may already have closed. */ export function createSubagentUsageSink( collectedUsage: UsageMetadata[], - onUsage?: (usage: UsageMetadata) => void, -): (event: SubagentUsageEvent) => void { + onUsage?: (usage: UsageMetadata) => void | Promise, + recordDetachedUsage?: (usage: UsageMetadata) => void | Promise, +): (event: SubagentUsageEvent) => void | Promise { return (event) => { if (event?.usage == null) { return; @@ -728,10 +787,38 @@ export function createSubagentUsageSink( if (billingAgentId != null && billingAgentId !== '') { usage.agentId = billingAgentId; } + /** Usage emission is observability/UI plumbing. It must never prevent the + * authoritative billing path from running when a detached child outlives + * its parent transport. The host emitter normally contains its own error + * handling; this boundary also protects custom hosts and synchronous + * lifecycle failures. */ + const emitUsage = () => { + try { + const emitted = onUsage?.(usage); + if (emitted != null) { + void Promise.resolve(emitted).catch((err) => { + logger.warn('[createSubagentUsageSink] Failed to emit subagent usage', err); + }); + } + } catch (err) { + logger.warn('[createSubagentUsageSink] Failed to emit subagent usage', err); + } + }; + /** A detached task can finish after its parent turn's one-time billing + * flush. Its AsyncLocalStorage context therefore owns the usage: persist + * it with the child transcript and bill it immediately. Foreground child + * calls retain the existing parent-turn batch path. */ + if (recordDetachedUsage != null && collectDetachedSubagentUsage(usage)) { + /** Emission is already retained/flushed by the host and must not add + * transport latency to the child model loop. Billing is the durable + * side effect the SDK needs to await. */ + emitUsage(); + return Promise.resolve(recordDetachedUsage(usage)).then(() => undefined); + } collectedUsage.push(usage); /** Lets the host stream the billed child usage to the client (tagged * `subagent`, so it folds into session cost/totals but not the live * gauge) — child runs never reach ModelEndHandler's emit path. */ - onUsage?.(usage); + emitUsage(); }; } diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 1f1a5590f8..de13becfd2 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -527,6 +527,8 @@ export interface UsageMetadata { /** Agent that produced this usage (graph agent id / subagent agent id). Lets * multi-endpoint graphs price each call with its own endpoint token config. */ agentId?: string; + /** Authoritative display cost attached by the host before durable child persistence. */ + cost?: number; /** * OpenAI-style cache token details. * Present for OpenAI models (GPT-4, o1, etc.) diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index a1bcb6489a..0720263d27 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -55,6 +55,7 @@ export const defaultRetrievalModels = [ export const excludedKeys = new Set([ 'conversationId', + 'subagentThread', 'title', 'iconURL', 'greeting', diff --git a/packages/data-provider/src/schemas.spec.ts b/packages/data-provider/src/schemas.spec.ts index 01c235ccd1..48b0f83397 100644 --- a/packages/data-provider/src/schemas.spec.ts +++ b/packages/data-provider/src/schemas.spec.ts @@ -10,6 +10,7 @@ import { eReasoningEffortSchema, eReasoningModeSchema, eReasoningContextSchema, + subagentThreadLineageSchema, } from './schemas'; describe('anthropicSettings', () => { @@ -621,3 +622,30 @@ describe('ReasoningContext', () => { expect(() => eReasoningContextSchema.parse('next_turn')).toThrow(); }); }); + +describe('subagentThreadLineageSchema', () => { + const lineage = { + rootConversationId: 'root-conversation', + parentConversationId: 'parent-conversation', + parentMessageId: 'parent-message', + parentToolCallId: 'tool-call', + parentAgentId: 'parent-agent', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }; + + it('accepts durable child-thread lineage', () => { + expect(subagentThreadLineageSchema.parse(lineage)).toEqual(lineage); + }); + + it('rejects non-positive depth and unknown execution shapes', () => { + expect(() => subagentThreadLineageSchema.parse({ ...lineage, depth: 0 })).toThrow(); + expect(() => + subagentThreadLineageSchema.parse({ ...lineage, subagentKind: 'workflow' }), + ).toThrow(); + expect(() => + subagentThreadLineageSchema.parse({ ...lineage, parentConversationId: '' }), + ).toThrow(); + }); +}); diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index a8ff1f4dae..97fbd700d2 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -951,6 +951,19 @@ const DocumentType: z.ZodType = z.lazy(() => ]), ); +export const subagentThreadLineageSchema = z.object({ + rootConversationId: z.string().min(1), + parentConversationId: z.string().min(1), + parentMessageId: z.string().min(1), + parentToolCallId: z.string().min(1), + parentAgentId: z.string().min(1).optional(), + subagentType: z.string().min(1), + subagentKind: z.enum(['agent', 'graph']), + depth: z.number().int().positive(), +}); + +export type TSubagentThreadLineage = z.infer; + export const tConversationSchema = z.object({ conversationId: z.string().nullable(), endpoint: eModelEndpointSchema.nullable(), @@ -1026,6 +1039,8 @@ export const tConversationSchema = z.object({ assistant_id: z.string().optional(), /* agents */ agent_id: z.string().optional(), + /** Durable parent/child navigation for a subagent thread. */ + subagentThread: subagentThreadLineageSchema.optional(), /* AWS Bedrock */ region: z.string().optional(), maxTokens: coerceNumber.optional(), diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index 963e8778fe..1c69d79d6d 100644 --- a/packages/data-schemas/src/methods/conversation.spec.ts +++ b/packages/data-schemas/src/methods/conversation.spec.ts @@ -1583,6 +1583,86 @@ describe('Conversation Operations', () => { expect(deletedConvo).toBeNull(); }); + it('cascades parent deletion through owner-scoped child-thread lineage', async () => { + const parentId = uuidv4(); + const childId = uuidv4(); + const grandchildId = uuidv4(); + const otherUsersChildId = uuidv4(); + const lineage = ( + parentConversationId: string, + rootConversationId: string, + depth: number, + ) => ({ + rootConversationId, + parentConversationId, + parentMessageId: `message-${depth}`, + parentToolCallId: `tool-${depth}`, + subagentType: 'agent-child', + subagentKind: 'agent', + depth, + }); + await Conversation.create([ + { conversationId: parentId, user: 'user123', endpoint: EModelEndpoint.agents }, + { + conversationId: childId, + user: 'user123', + endpoint: EModelEndpoint.agents, + subagentThread: lineage(parentId, parentId, 1), + }, + { + conversationId: grandchildId, + user: 'user123', + endpoint: EModelEndpoint.agents, + subagentThread: lineage(childId, parentId, 2), + }, + { + conversationId: otherUsersChildId, + user: 'other-user', + endpoint: EModelEndpoint.agents, + subagentThread: lineage(parentId, parentId, 1), + }, + ]); + deleteMessages.mockResolvedValue({ acknowledged: true, deletedCount: 3 }); + + const result = await deleteConvos('user123', { conversationId: parentId }); + + expect(result.deletedCount).toBe(3); + expect(result.conversationIds).toEqual([parentId, childId, grandchildId]); + expect(deleteMessages).toHaveBeenCalledWith({ + conversationId: { $in: [parentId, childId, grandchildId] }, + user: 'user123', + }); + expect(await Conversation.find({ user: 'user123' })).toHaveLength(0); + expect(await Conversation.findOne({ conversationId: otherUsersChildId })).not.toBeNull(); + }); + + it('does not delete a parent when deleting one child thread', async () => { + const parentId = uuidv4(); + const childId = uuidv4(); + await Conversation.create([ + { conversationId: parentId, user: 'user123', endpoint: EModelEndpoint.agents }, + { + conversationId: childId, + user: 'user123', + endpoint: EModelEndpoint.agents, + subagentThread: { + rootConversationId: parentId, + parentConversationId: parentId, + parentMessageId: 'message-1', + parentToolCallId: 'tool-1', + subagentType: 'agent-child', + subagentKind: 'agent', + depth: 1, + }, + }, + ]); + + const result = await deleteConvos('user123', { conversationId: childId }); + + expect(result.conversationIds).toEqual([childId]); + expect(await Conversation.findOne({ conversationId: parentId })).not.toBeNull(); + }); + it('should throw error if no conversations found', async () => { await expect(deleteConvos('user123', { conversationId: 'non-existent' })).rejects.toThrow( 'Conversation not found or already deleted.', @@ -3190,6 +3270,153 @@ describe('Conversation Operations', () => { }); }); + describe('subagent thread leases', () => { + it('reserves child lineage once without overwriting a concurrent winner', async () => { + await Conversation.init(); + const conversationId = uuidv4(); + const lineage = { + rootConversationId: 'root', + parentConversationId: 'parent', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool', + parentAgentId: 'parent-agent', + subagentType: 'researcher', + subagentKind: 'agent' as const, + depth: 1, + }; + const reservations = await Promise.all( + ['First title', 'Second title'].map((title) => + methods.reserveSubagentThread({ + user: 'reservation-user', + conversationId, + conversation: { endpoint: EModelEndpoint.agents, title, subagentThread: lineage }, + }), + ), + ); + + expect(reservations.filter((reservation) => reservation.created)).toHaveLength(1); + const saved = await methods.getConvo('reservation-user', conversationId); + expect(saved?.title).toBe( + reservations.find((reservation) => reservation.created)?.conversation.title, + ); + expect(saved?.subagentThread).toEqual(lineage); + }); + + it('admits one cross-replica owner and fences renewal and release by token', async () => { + const conversationId = uuidv4(); + await Conversation.create({ + conversationId, + user: 'lease-user', + endpoint: EModelEndpoint.agents, + subagentThread: { + rootConversationId: 'root', + parentConversationId: 'parent', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool', + parentAgentId: 'parent-agent', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + }); + const now = new Date('2026-08-18T00:00:00.000Z'); + const expiresAt = new Date(now.getTime() + 30_000); + + const claims = await Promise.all( + ['token-a', 'token-b'].map((token) => + methods.acquireSubagentThreadLease({ + user: 'lease-user', + conversationId, + token, + taskId: `task-${token}`, + now, + expiresAt, + }), + ), + ); + expect(claims.filter(Boolean)).toHaveLength(1); + 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); + expect(await methods.getConvo('lease-user', conversationId)).not.toHaveProperty( + 'subagentThreadLease', + ); + await expect( + methods.renewSubagentThreadLease({ + user: 'lease-user', + conversationId, + token: loser, + now, + expiresAt: new Date(expiresAt.getTime() + 30_000), + }), + ).resolves.toBe(false); + await expect( + methods.releaseSubagentThreadLease({ + user: 'lease-user', + conversationId, + token: loser, + }), + ).resolves.toBe(false); + await expect( + methods.releaseSubagentThreadLease({ + user: 'lease-user', + conversationId, + token: winner, + }), + ).resolves.toBe(true); + expect(await methods.countActiveSubagentThreadLeases({ user: 'lease-user', now })).toBe(0); + }); + + it('allows takeover only after the prior lease expires', async () => { + const conversationId = uuidv4(); + await Conversation.create({ + conversationId, + user: 'takeover-user', + endpoint: EModelEndpoint.agents, + subagentThread: { + rootConversationId: 'root', + parentConversationId: 'parent', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool', + subagentType: 'graph', + subagentKind: 'graph', + depth: 1, + }, + }); + const startedAt = new Date('2026-08-18T00:00:00.000Z'); + const firstExpiry = new Date(startedAt.getTime() + 1_000); + await methods.acquireSubagentThreadLease({ + user: 'takeover-user', + conversationId, + token: 'first', + taskId: 'task-first', + now: startedAt, + expiresAt: firstExpiry, + }); + + await expect( + methods.acquireSubagentThreadLease({ + user: 'takeover-user', + conversationId, + token: 'second', + taskId: 'task-second', + now: new Date(firstExpiry.getTime() - 1), + expiresAt: new Date(firstExpiry.getTime() + 1_000), + }), + ).resolves.toBe(false); + await expect( + methods.acquireSubagentThreadLease({ + user: 'takeover-user', + conversationId, + token: 'second', + taskId: 'task-second', + now: firstExpiry, + expiresAt: new Date(firstExpiry.getTime() + 1_000), + }), + ).resolves.toBe(true); + }); + }); + describe('tenantId stripping', () => { it('saveConvo should not write caller-supplied tenantId to the document', async () => { const conversationId = uuidv4(); diff --git a/packages/data-schemas/src/methods/conversation.ts b/packages/data-schemas/src/methods/conversation.ts index ff3c21ab5e..79b9d84ad9 100644 --- a/packages/data-schemas/src/methods/conversation.ts +++ b/packages/data-schemas/src/methods/conversation.ts @@ -1,7 +1,13 @@ import { RetentionMode } from 'librechat-data-provider'; import type { FilterQuery, Model, SortOrder, Types } from 'mongoose'; import type { DeleteResult } from 'mongoose'; -import type { AppConfig, IChatProjectDocument, IConversation, ISharedLink } from '~/types'; +import type { + AppConfig, + IChatProjectDocument, + IConversation, + ISharedLink, + ISubagentThreadReservation, +} from '~/types'; import type { MessageMethods } from './message'; import { activeExpirationFilter, @@ -143,6 +149,40 @@ export interface ConversationMethods { convoMap: Record; }>; getConvo(user: string, conversationId: string): Promise; + reserveSubagentThread(input: { + user: string; + conversationId: string; + conversation: Partial; + tenantId?: string; + }): Promise; + acquireSubagentThreadLease(input: { + user: string; + conversationId: string; + token: string; + taskId: string; + now: Date; + expiresAt: Date; + tenantId?: string; + }): Promise; + renewSubagentThreadLease(input: { + user: string; + conversationId: string; + token: string; + now: Date; + expiresAt: Date; + tenantId?: string; + }): Promise; + releaseSubagentThreadLease(input: { + user: string; + conversationId: string; + token: string; + tenantId?: string; + }): Promise; + countActiveSubagentThreadLeases(input: { + user: string; + now: Date; + tenantId?: string; + }): Promise; getConvoOwnership( user: string, conversationId: string, @@ -203,6 +243,152 @@ export function createConversationMethods( } } + /** Creates immutable child lineage exactly once without overwriting a concurrent winner. */ + async function reserveSubagentThread(input: { + user: string; + conversationId: string; + conversation: Partial; + tenantId?: string; + }): Promise { + const Conversation = mongoose.models.Conversation as Model; + const filter = { + user: input.user, + conversationId: input.conversationId, + ...(input.tenantId == null ? { tenantId: { $exists: false } } : { tenantId: input.tenantId }), + }; + try { + const result = (await Conversation.findOneAndUpdate( + filter, + { + $setOnInsert: { + ...input.conversation, + user: input.user, + conversationId: input.conversationId, + messages: [], + }, + }, + { new: true, upsert: true, includeResultMetadata: true, setDefaultsOnInsert: true }, + )) as unknown as ConversationUpdateResult; + if (result.value == null) { + throw new Error('Unable to reserve the subagent thread.'); + } + return { + conversation: result.value.toObject(), + created: result.lastErrorObject?.updatedExisting !== true, + }; + } catch (error) { + /** Concurrent upserts can race at the unique index. The document that won is + * the reservation; callers still validate its immutable lineage before use. */ + if ((error as { code?: number }).code === 11000) { + const existing = await Conversation.findOne(filter).lean(); + if (existing != null) { + return { conversation: existing, created: false }; + } + } + throw error; + } + } + + function subagentLeaseTenantFilter(tenantId?: string): FilterQuery { + return tenantId == null ? { tenantId: { $exists: false } } : { tenantId }; + } + + /** Atomically claims one durable child thread across all API replicas. */ + async function acquireSubagentThreadLease(input: { + user: string; + conversationId: string; + token: string; + taskId: string; + now: Date; + expiresAt: Date; + tenantId?: string; + }): Promise { + const Conversation = mongoose.models.Conversation as Model; + const result = await Conversation.updateOne( + { + user: input.user, + conversationId: input.conversationId, + subagentThread: { $exists: true }, + ...subagentLeaseTenantFilter(input.tenantId), + $or: [ + { subagentThreadLease: { $exists: false } }, + { 'subagentThreadLease.expiresAt': { $lte: input.now } }, + { 'subagentThreadLease.token': input.token }, + ], + }, + { + $set: { + subagentThreadLease: { + token: input.token, + taskId: input.taskId, + expiresAt: input.expiresAt, + }, + }, + }, + { timestamps: false }, + ); + return result.matchedCount === 1; + } + + /** Renews only the unexpired lease owned by this exact task token. */ + async function renewSubagentThreadLease(input: { + user: string; + conversationId: string; + token: string; + now: Date; + expiresAt: Date; + tenantId?: string; + }): Promise { + const Conversation = mongoose.models.Conversation as Model; + const result = await Conversation.updateOne( + { + user: input.user, + conversationId: input.conversationId, + ...subagentLeaseTenantFilter(input.tenantId), + 'subagentThreadLease.token': input.token, + 'subagentThreadLease.expiresAt': { $gt: input.now }, + }, + { $set: { 'subagentThreadLease.expiresAt': input.expiresAt } }, + { timestamps: false }, + ); + return result.matchedCount === 1; + } + + /** Releases only the lease owned by this exact task token. */ + async function releaseSubagentThreadLease(input: { + user: string; + conversationId: string; + token: string; + tenantId?: string; + }): Promise { + const Conversation = mongoose.models.Conversation as Model; + const result = await Conversation.updateOne( + { + user: input.user, + conversationId: input.conversationId, + ...subagentLeaseTenantFilter(input.tenantId), + 'subagentThreadLease.token': input.token, + }, + { $unset: { subagentThreadLease: 1 } }, + { timestamps: false }, + ); + return result.modifiedCount === 1; + } + + /** Counts live child execution fences while account deletion drains. */ + async function countActiveSubagentThreadLeases(input: { + user: string; + now: Date; + tenantId?: string; + }): Promise { + const Conversation = mongoose.models.Conversation as Model; + return Conversation.countDocuments({ + user: input.user, + ...subagentLeaseTenantFilter(input.tenantId), + 'subagentThreadLease.expiresAt': { $gt: input.now }, + }); + } + /** * Ownership probe for request validation: resolves only the owning user id * instead of materializing the full conversation document (preset spread + @@ -1071,12 +1257,58 @@ export function createConversationMethods( const Conversation = mongoose.models.Conversation as Model; const { deleteMessages } = getMessageMethods(); const userFilter = { ...filter, user }; - const conversations = await Conversation.find(userFilter).select( - 'conversationId chatProjectId tags', - ); - const conversationIds = conversations.map((c) => c.conversationId); + type DeletionConversation = Pick; + const conversations = await Conversation.find(userFilter) + .select('conversationId chatProjectId tags') + .lean(); + if (!conversations.length) { + throw new Error('Conversation not found or already deleted.'); + } + + /** + * Delete roots first, then walk their owner-scoped lineage. Apart from + * making child threads share the parent's lifecycle, deleting each wave + * before discovering the next closes the child-creation race: a creator + * that started concurrently sees its parent disappear and rolls back, + * while a child already committed is found by the next query. + */ + const deletedConversations: DeletionConversation[] = []; + const seen = new Set(); + let pending = conversations; + let acknowledged = true; + let deletedCount = 0; + while (pending.length > 0) { + const wave = pending.filter((conversation) => !seen.has(conversation.conversationId)); + if (wave.length === 0) { + break; + } + const waveIds = wave.map((conversation) => conversation.conversationId); + try { + const result = await Conversation.deleteMany({ user, conversationId: { $in: waveIds } }); + acknowledged &&= result.acknowledged; + deletedCount += result.deletedCount; + for (const conversation of wave) { + seen.add(conversation.conversationId); + deletedConversations.push(conversation); + } + pending = await Conversation.find({ + user, + 'subagentThread.parentConversationId': { $in: waveIds }, + }) + .select('conversationId chatProjectId tags') + .lean(); + } catch (error) { + if (deletedConversations.length === 0) { + throw error; + } + logger.error('[deleteConvos] Root deleted but child-thread cascade failed', error); + break; + } + } + + const conversationIds = deletedConversations.map((c) => c.conversationId); const projectIds = new Set( - conversations + deletedConversations .map((conversation) => conversation.chatProjectId) .filter((projectId): projectId is string => Boolean(projectId)), ); @@ -1087,7 +1319,7 @@ export function createConversationMethods( * the bookmark count once. */ const tagDecrements: string[] = []; - for (const conversation of conversations) { + for (const conversation of deletedConversations) { if (!conversation.tags?.length) { continue; } @@ -1096,11 +1328,7 @@ export function createConversationMethods( } } - if (!conversationIds.length) { - throw new Error('Conversation not found or already deleted.'); - } - - const deleteConvoResult = await Conversation.deleteMany(userFilter); + const deleteConvoResult: DeleteResult = { acknowledged, deletedCount }; const deleted = deleteConvoResult.deletedCount > 0; /** @@ -1292,6 +1520,11 @@ export function createConversationMethods( getConvosByCursor, getConvosQueried, getConvo, + reserveSubagentThread, + acquireSubagentThreadLease, + renewSubagentThreadLease, + releaseSubagentThreadLease, + countActiveSubagentThreadLeases, getConvoOwnership, getConvoRetention, getConvoTitle, diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index dd0922d886..097fe99810 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -511,6 +511,11 @@ describe('Message Operations', () => { usage: { input: 10, output: 20 }, thoughtSignatures: { tool_1: 'opaque' }, }, + subagentTask: { + attemptKey: 'private-attempt', + requestFingerprint: 'private-fingerprint', + status: 'running', + }, attachments: [ { type: 'web_search', @@ -566,6 +571,7 @@ describe('Message Operations', () => { 'contextMeta', 'langfuseSampled', 'langfuseDestinationIds', + 'subagentTask', ]) { expect(hidden[field]).toBeUndefined(); } diff --git a/packages/data-schemas/src/schema/convo.ts b/packages/data-schemas/src/schema/convo.ts index f5db5e1184..5b389e466e 100644 --- a/packages/data-schemas/src/schema/convo.ts +++ b/packages/data-schemas/src/schema/convo.ts @@ -29,6 +29,31 @@ const convoSchema: Schema = new Schema( agent_id: { type: String, }, + subagentThread: { + type: { + rootConversationId: { type: String, required: true }, + parentConversationId: { type: String, required: true }, + parentMessageId: { type: String, required: true }, + parentToolCallId: { type: String, required: true }, + parentAgentId: { type: String }, + subagentType: { type: String, required: true }, + subagentKind: { type: String, enum: ['agent', 'graph'], required: true }, + depth: { type: Number, min: 1, required: true }, + }, + _id: false, + default: undefined, + }, + /** Mongo-backed continuation lease shared by every API replica. */ + subagentThreadLease: { + type: { + token: { type: String, required: true }, + taskId: { type: String, required: true }, + expiresAt: { type: Date, required: true }, + }, + _id: false, + default: undefined, + select: false, + }, tags: { type: [String], default: [], @@ -77,6 +102,9 @@ convoSchema.index({ user: 1, isArchived: 1, archivedAt: -1, createdAt: -1, _id: convoSchema.index({ user: 1, pinned: 1, updatedAt: -1, _id: -1 }); convoSchema.index({ user: 1, isTemporary: 1, expiredAt: 1 }); +/** Owner-scoped child-thread cascade lookup used when a parent is deleted. */ +convoSchema.index({ user: 1, 'subagentThread.parentConversationId': 1 }); +convoSchema.index({ user: 1, 'subagentThreadLease.expiresAt': 1 }); // index for MeiliSearch sync operations convoSchema.index({ _meiliIndex: 1, isTemporary: 1, expiredAt: 1 }); diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index 2d74d6b71c..d8279ca05e 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -124,6 +124,31 @@ const messageSchema: Schema = new Schema( type: String, }, metadata: { type: mongoose.Schema.Types.Mixed }, + subagentTranscript: { + type: { + taskId: { type: String, required: true }, + mode: { type: String, enum: ['append', 'replace'], required: true }, + messagesJson: { type: String, required: true }, + }, + _id: false, + select: false, + default: undefined, + }, + /** Durable, server-only marker used to make detached retries at-most-once. */ + subagentTask: { + type: { + attemptKey: { type: String, required: true }, + requestFingerprint: { type: String }, + status: { + type: String, + enum: ['running', 'completed', 'error', 'cancelled'], + required: true, + }, + }, + _id: false, + select: false, + default: undefined, + }, contextMeta: { type: { calibrationRatio: { type: Number }, diff --git a/packages/data-schemas/src/types/convo.ts b/packages/data-schemas/src/types/convo.ts index 9652b85819..5fa5fbf0a9 100644 --- a/packages/data-schemas/src/types/convo.ts +++ b/packages/data-schemas/src/types/convo.ts @@ -1,5 +1,17 @@ +import type { TSubagentThreadLineage } from 'librechat-data-provider'; import type { Document, Types } from 'mongoose'; +export interface ISubagentThreadLease { + token: string; + taskId: string; + expiresAt: Date; +} + +export interface ISubagentThreadReservation { + conversation: IConversation; + created: boolean; +} + // @ts-ignore export interface IConversation extends Document { conversationId: string; @@ -35,6 +47,9 @@ export interface IConversation extends Document { resendFiles?: boolean; imageDetail?: string; agent_id?: string; + subagentThread?: TSubagentThreadLineage; + /** Internal execution fence. Excluded from ordinary conversation reads. */ + subagentThreadLease?: ISubagentThreadLease; assistant_id?: string; instructions?: string; stop?: string[]; diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index 60e354a870..9b1d4492f4 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -42,6 +42,18 @@ export interface IMessage extends Document { iconURL?: string; addedConvo?: boolean; metadata?: Record; + /** Server-private canonical message delta for durable subagent-thread continuation. */ + subagentTranscript?: { + taskId: string; + mode: 'append' | 'replace'; + messagesJson: string; + }; + /** Server-private durable idempotency marker for one detached subagent turn. */ + subagentTask?: { + attemptKey: string; + requestFingerprint?: string; + status: 'running' | 'completed' | 'error' | 'cancelled'; + }; contextMeta?: { calibrationRatio?: number; encoding?: string;