diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index 649bb8a010..e513f12e2a 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -2,8 +2,15 @@ import mongoose from 'mongoose'; import { randomUUID } from 'node:crypto'; import { MongoMemoryServer } from 'mongodb-memory-server'; import { Constants, EModelEndpoint } from 'librechat-data-provider'; -import { createMethods, createModels, logger } from '@librechat/data-schemas'; import { AIMessage, HumanMessage } from '@librechat/agents/langchain/messages'; +import { + createMethods, + createModels, + getTenantId, + getUserId, + logger, + tenantStorage, +} from '@librechat/data-schemas'; import type { SubagentTaskClaim, SubagentTaskControlCommand, @@ -376,6 +383,50 @@ describe('SubagentThreadTaskStore', () => { expect(await methods.getConvo(userId, requireThreadId(started))).not.toBeNull(); }); + it('reconstructs trusted owner context after the admitting request has ended', async () => { + const userId = 'detached-context-user'; + const tenantId = 'detached-context-tenant'; + const parentConversationId = randomUUID(); + await tenantStorage.run({ tenantId, userId }, async () => + saveParent(userId, parentConversationId, { tenantId }), + ); + + const observedContexts: Array<{ tenantId?: string; userId?: string }> = []; + const observeContext = () => { + observedContexts.push({ tenantId: getTenantId(), userId: getUserId() }); + }; + const store = new SubagentThreadTaskStore(methods, { + isOwnerActive: async () => { + observeContext(); + return true; + }, + }); + const config = buildSubagentThreadTaskConfig(store, { + userId, + tenantId, + parentConversationId, + }); + const defaultRun = taskRequest(config.scopeId).run; + const run = jest.fn(async (...args: Parameters) => { + observeContext(); + return defaultRun(...args); + }); + + /** `start` deliberately runs outside `tenantStorage.run`: the detached task + * owns only its serialized host scope once the HTTP request has returned. */ + const started = store.start(taskRequest(config.scopeId, { run })); + await waitForSettled(store, config.scopeId, started); + + expect(run).toHaveBeenCalledTimes(1); + expect(observedContexts.length).toBeGreaterThan(0); + expect(observedContexts).toEqual(observedContexts.map(() => ({ tenantId, userId }))); + const messages = await tenantStorage.run({ tenantId, userId }, async () => + methods.getMessages({ user: userId, conversationId: requireThreadId(started) }), + ); + expect(messages).toHaveLength(2); + expect(messages.every((message) => message.tenantId === tenantId)).toBe(true); + }); + it('fails without leaving an orphan when parent persistence rejects', async () => { const userId = 'parent-gate-failure-user'; const parentConversationId = randomUUID(); diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index becac41104..eb3f21fcac 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -1,6 +1,6 @@ import { randomUUID } from 'node:crypto'; -import { logger } from '@librechat/data-schemas'; import { InMemorySubagentTaskStore } from '@librechat/agents'; +import { logger, tenantStorage } from '@librechat/data-schemas'; import { EModelEndpoint, Constants } from 'librechat-data-provider'; import { mapChatMessagesToStoredMessages, @@ -63,6 +63,7 @@ export interface SubagentCancellationPlan { } /** Three missed 10-second transport heartbeats retire a crashed owner. */ const DEFAULT_TASK_ROUTING_TTL_MS = 30_000; +const SLOW_PREPARATION_WARN_MS = 5_000; const MAX_TRANSCRIPT_BYTES = 12 * 1024 * 1024; const TRANSCRIPT_SELECT = 'messageId parentMessageId text createdAt +subagentTranscript +subagentTask'; @@ -380,6 +381,21 @@ function safeErrorMessage(error: unknown): string { return `Subagent task failed: ${publicFailureDetail(error).slice(0, 2_000)}`; } +async function observeSlowPreparation( + operation: Promise, + context: { stage: string; taskId: string; threadId: string }, +): Promise { + const warning = setTimeout(() => { + logger.warn('[subagentThreads] Child-thread preparation is still waiting', context); + }, SLOW_PREPARATION_WARN_MS); + warning.unref?.(); + try { + return await operation; + } finally { + clearTimeout(warning); + } +} + /** Persists view-only logical child threads with owner-routed controls and a shared execution fence. */ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { readonly supportsThreadContinuation = true; @@ -519,128 +535,129 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { started = super.start({ ...request, threadId, - run: async (runtime: SubagentTaskRuntime) => { - lease.taskId = runtime.taskId; - lease.running = true; - const detachedUsage: UsageMetadata[] = []; - let prepared: PreparedThread | undefined; - try { - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - /** Publish the owner address before any provider work: a child running - * while unaddressable cannot be polled, controlled, or cancelled, and its - * side effects would already have happened by the time a heartbeat - * republished it. A failed registration fails the task closed instead. */ - await this.taskControlTransport?.registerTask( - request.scopeId, - runtime.taskId, - this.taskRoutingTtlMs, - ); - await parentReady; - prepared = await this.prepareThread( - request.scopeId, - scope, - threadId, - isContinuation, - request, - runtime.taskId, - lease, - ); - await this.registerTaskWakeup(scope, prepared.conversation.conversationId, request, { - taskId: prepared.replay?.taskId ?? runtime.taskId, - parentRunId: prepared.replay?.parentRunId ?? request.parentRunId, - createdAt: prepared.taskCreatedAt, - }); - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - if (prepared.replay != null) { - if (prepared.replay.status === 'completed') { - return { content: prepared.replay.content }; + run: (runtime: SubagentTaskRuntime) => + this.runWithOwnerContext(scope, async () => { + lease.taskId = runtime.taskId; + lease.running = true; + const detachedUsage: UsageMetadata[] = []; + let prepared: PreparedThread | undefined; + try { + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); } - throw new SubagentThreadPublicError(prepared.replay.content); - } - if (!(await this.renewSharedLease(scope, threadId, lease))) { - throw new SubagentThreadPublicError( - 'This child thread is already being continued by another run.', + /** Publish the owner address before any provider work: a child running + * while unaddressable cannot be polled, controlled, or cancelled, and its + * side effects would already have happened by the time a heartbeat + * republished it. A failed registration fails the task closed instead. */ + await this.taskControlTransport?.registerTask( + request.scopeId, + runtime.taskId, + this.taskRoutingTtlMs, ); - } - const preparedThread = prepared; - const result = await runWithDetachedSubagentUsage(detachedUsage, () => - request.run(runtime, preparedThread.initialMessages), - ); - if (runtime.signal.aborted) { - throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); - } - if (!(await this.renewSharedLease(scope, threadId, lease))) { - throw new SubagentThreadPublicError( - 'This child thread is already being continued by another run.', + await parentReady; + prepared = await this.prepareThread( + request.scopeId, + scope, + threadId, + isContinuation, + request, + runtime.taskId, + lease, + ); + await this.registerTaskWakeup(scope, prepared.conversation.conversationId, request, { + taskId: prepared.replay?.taskId ?? runtime.taskId, + parentRunId: prepared.replay?.parentRunId ?? request.parentRunId, + createdAt: prepared.taskCreatedAt, + }); + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + if (prepared.replay != null) { + if (prepared.replay.status === 'completed') { + return { content: prepared.replay.content }; + } + throw new SubagentThreadPublicError(prepared.replay.content); + } + if (!(await this.renewSharedLease(scope, threadId, lease))) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + const preparedThread = prepared; + const result = await runWithDetachedSubagentUsage(detachedUsage, () => + request.run(runtime, preparedThread.initialMessages), + ); + if (runtime.signal.aborted) { + throw runtime.signal.reason ?? new Error('Subagent task was cancelled.'); + } + if (!(await this.renewSharedLease(scope, threadId, lease))) { + throw new SubagentThreadPublicError( + 'This child thread is already being continued by another run.', + ); + } + lease.settling = true; + await this.persistResult( + scope, + request, + runtime.taskId, + prepared, + result, + detachedUsage, + ); + return result; + } catch (error) { + /** A replay is already terminal in Mongo. A temporary wakeup-queue + * outage must not overwrite that canonical result with a new error. */ + if (prepared?.replay != null) { + throw error; + } + const mayPersist = + lease.shared == null || (await this.renewSharedLease(scope, threadId, lease)); + const terminalTask = this.get(request.scopeId, runtime.taskId); + if (runtime.signal.aborted && terminalTask?.status === 'cancelled') { + if (mayPersist) { + await this.persistCancellation( + scope, + threadId, + request, + runtime.taskId, + detachedUsage, + ).catch((persistError) => { + logger.error( + '[subagentThreads] Failed to persist child-thread cancellation', + persistError, + ); + }); + } + throw error; + } + logger.error( + '[subagentThreads] Child-thread execution failed', + publicFailureDetail(error), ); - } - lease.settling = true; - await this.persistResult( - scope, - request, - runtime.taskId, - prepared, - result, - detachedUsage, - ); - return result; - } catch (error) { - /** A replay is already terminal in Mongo. A temporary wakeup-queue - * outage must not overwrite that canonical result with a new error. */ - if (prepared?.replay != null) { - throw error; - } - const mayPersist = - lease.shared == null || (await this.renewSharedLease(scope, threadId, lease)); - const terminalTask = this.get(request.scopeId, runtime.taskId); - if (runtime.signal.aborted && terminalTask?.status === 'cancelled') { if (mayPersist) { - await this.persistCancellation( + await this.persistFailure( scope, threadId, request, runtime.taskId, + error, detachedUsage, ).catch((persistError) => { logger.error( - '[subagentThreads] Failed to persist child-thread cancellation', + '[subagentThreads] Failed to persist child-thread failure', persistError, ); }); } - throw error; + throw new Error(publicFailureDetail(error)); + } finally { + await this.stopAndReleaseSharedLease(scope, threadId, lease); + if (this.activeThreads.get(lockKey) === lease) { + this.activeThreads.delete(lockKey); + } } - 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) { @@ -657,6 +674,20 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { return started; } + /** Detached tasks intentionally outlive the HTTP request that admitted them. + * Reconstruct only the trusted owner identity carried by the opaque task scope + * so tenant-isolated database reads and lazy child initialization do not depend + * on request AsyncLocalStorage remaining alive after the parent turn returns. */ + private runWithOwnerContext(scope: SubagentThreadScope, run: () => Promise): Promise { + return tenantStorage.run( + { + userId: scope.userId, + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + }, + run, + ); + } + /** * Claims locally when possible, otherwise asks the registered owning replica. * @@ -1566,13 +1597,22 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { /** Account deletion can fence the owner after the optimistic probe but before * this lease exists. Once the lease is visible, revalidate so deletion either * observes and drains us or wins before any provider work can begin. */ - if (!(await this.isOwnerActive(scope.userId))) { + if ( + !(await observeSlowPreparation(this.isOwnerActive(scope.userId), { + stage: 'owner_recheck', + taskId, + threadId, + })) + ) { throw new SubagentThreadDeletedError('The thread owner is unavailable.'); } - const allMessages = (await this.methods.getMessages( - { conversationId: threadId, user: scope.userId }, - TRANSCRIPT_SELECT, - { sort: { createdAt: 1, _id: 1 } }, + const allMessages = (await observeSlowPreparation( + this.methods.getMessages( + { conversationId: threadId, user: scope.userId }, + TRANSCRIPT_SELECT, + { sort: { createdAt: 1, _id: 1 } }, + ), + { stage: 'transcript_read', taskId, threadId }, )) as ThreadMessage[]; const attemptKey = createSubagentAttemptKey(scopeId, request.idempotencyKey); const requestFingerprint = normalizedRequestFingerprint(request); @@ -1685,25 +1725,28 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { break; } } - const savedUserMessage = await this.methods.saveMessage( - { userId: scope.userId }, - { - messageId: userMessageId, - conversationId: threadId, - parentMessageId, - sender: 'User', - text: request.input, - endpoint: EModelEndpoint.agents, - isCreatedByUser: true, - subagentTask: { - attemptKey, - parentRunId: request.parentRunId, - ...(requestFingerprint == null ? {} : { requestFingerprint }), - status: 'running', + const savedUserMessage = await observeSlowPreparation( + this.methods.saveMessage( + { userId: scope.userId }, + { + messageId: userMessageId, + conversationId: threadId, + parentMessageId, + sender: 'User', + text: request.input, + endpoint: EModelEndpoint.agents, + isCreatedByUser: true, + subagentTask: { + attemptKey, + parentRunId: request.parentRunId, + ...(requestFingerprint == null ? {} : { requestFingerprint }), + status: 'running', + }, + ...retentionFields(conversation), }, - ...retentionFields(conversation), - }, - { context: 'SubagentThreadTaskStore.prepareThread' }, + { context: 'SubagentThreadTaskStore.prepareThread' }, + ), + { stage: 'seed_write', taskId, threadId }, ); if (savedUserMessage == null) { throw new Error('Unable to persist the child-thread input.');