From d641c398d56989793f350e2edda3f13d898bfab0 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 24 Aug 2026 20:05:39 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B3=20fix:=20Port=20Subagent=20Control?= =?UTF-8?q?=20Receipt=20Writes=20to=20DocumentDB-Safe=20Operators=20(#1517?= =?UTF-8?q?1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: harden subagent control receipt persistence * fix: harden durable subagent control replay * fix: await terminal control receipts on shutdown * fix: close subagent control replay races * test: type stale-owner transport fixture * fix: quiesce durable subagent controls * test: await subagent shutdown durability boundary * fix: serialize durable subagent controls * fix: fail shutdown on cleanup errors * fix: report cancellable result availability accurately * fix: fence subagent control receipt ownership * fix: close distributed control receipt races * test: type control receipt race fixture * chore: require authoritative control receipts * fix: close subagent control lifecycle races * style: separate control reservation member * test: harden subagent settlement wait * fix: preserve authoritative control replay state --- .../Endpoints/agents/subagentThreadStore.js | 31 +- .../agents/subagentThreadStore.spec.js | 23 +- packages/api/src/agents/guard.spec.ts | 1 + .../src/agents/subagentTaskRouting.spec.ts | 44 + .../api/src/agents/subagentTaskRouting.ts | 75 +- .../api/src/agents/subagentThreads.spec.ts | 992 +++++++++++++++++- packages/api/src/agents/subagentThreads.ts | 826 +++++++++++++-- packages/api/src/agents/view.spec.ts | 9 + packages/api/src/agents/view.ts | 15 +- packages/api/src/app/shutdown.spec.ts | 4 +- packages/api/src/app/shutdown.ts | 9 +- .../data-schemas/src/methods/message.spec.ts | 456 +++++++- packages/data-schemas/src/methods/message.ts | 558 ++++++---- packages/data-schemas/src/schema/message.ts | 2 +- packages/data-schemas/src/types/message.ts | 7 +- 15 files changed, 2649 insertions(+), 403 deletions(-) diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.js b/api/server/services/Endpoints/agents/subagentThreadStore.js index c094c2850c..6ab0398f61 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.js @@ -62,6 +62,7 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore( deleteConvos: db.deleteConvos, deleteMessages: db.deleteMessages, getConvo: db.getConvo, + getSubagentTaskControlReplay: db.getSubagentTaskControlReplay, getMessages: db.getMessages, listActiveSubagentThreadLeases: db.listActiveSubagentThreadLeases, recordSubagentTaskControlReceipt: db.recordSubagentTaskControlReceipt, @@ -93,6 +94,20 @@ registerShutdownTask( ); let taskRoutingConfigured = false; +let disconnectTaskRouting = () => {}; + +/** Store quiescence is required even without Redis. Optional transport cleanup + * is attached after configuration, but local child cancellation and the final + * durable receipt flush always participate in graceful shutdown. */ +registerShutdownTask( + 'subagent task store', + async () => { + await subagentThreadTaskStore.destroyTaskControlTransport(); + subagentThreadTaskStore.destroyActivityStream(); + disconnectTaskRouting(); + }, + { priority: 90 }, +); /** Starts the optional Redis owner directory before HTTP admission opens. */ async function configureSubagentTaskRouting() { @@ -126,17 +141,11 @@ async function configureSubagentTaskRouting() { throw error; } taskRoutingConfigured = true; - registerShutdownTask( - 'subagent task control transport', - async () => { - await subagentThreadTaskStore.destroyTaskControlTransport(); - subagentThreadTaskStore.destroyActivityStream(); - publisher.disconnect(); - activitySubscriber.disconnect(); - activityPublisher.disconnect(); - }, - { priority: 90 }, - ); + disconnectTaskRouting = () => { + publisher.disconnect(); + activitySubscriber.disconnect(); + activityPublisher.disconnect(); + }; } module.exports = subagentThreadTaskStore; diff --git a/api/server/services/Endpoints/agents/subagentThreadStore.spec.js b/api/server/services/Endpoints/agents/subagentThreadStore.spec.js index 513e765df0..f829c2092c 100644 --- a/api/server/services/Endpoints/agents/subagentThreadStore.spec.js +++ b/api/server/services/Endpoints/agents/subagentThreadStore.spec.js @@ -28,6 +28,7 @@ jest.mock('~/models', () => ({ deleteConvos: jest.fn(), deleteMessages: jest.fn(), getConvo: jest.fn(), + getSubagentTaskControlReplay: jest.fn(), getMessages: jest.fn(), listActiveSubagentThreadLeases: jest.fn(), recordSubagentTaskControlReceipt: jest.fn(), @@ -61,12 +62,16 @@ const db = require('~/models'); const activityPrepareRegistration = registerShutdownTask.mock.calls.find( ([name]) => name === 'subagent activity streams prepare', ); +const taskStoreShutdownRegistration = registerShutdownTask.mock.calls.find( + ([name]) => name === 'subagent task store', +); describe('subagent thread Redis lifecycle', () => { it('wires durable control receipt persistence into the host store', () => { expect(taskStoreMethods.recordSubagentTaskControlReceipt).toBe( db.recordSubagentTaskControlReceipt, ); + expect(taskStoreMethods.getSubagentTaskControlReplay).toBe(db.getSubagentTaskControlReplay); }); it('reads completion wakeup rollout state at task preparation time', async () => { @@ -83,6 +88,14 @@ describe('subagent thread Redis lifecycle', () => { expect(subagentThreadTaskStore.completionWakeupsEnabled).toBe(false); }); + it('registers local task-store quiescence independently of optional Redis setup', () => { + expect(taskStoreShutdownRegistration).toEqual([ + 'subagent task store', + expect.any(Function), + { priority: 90 }, + ]); + }); + it('closes activity SSE before drain and disconnects its subscriber after drain', async () => { const taskSubscriber = { disconnect: jest.fn() }; const activitySubscriber = { disconnect: jest.fn() }; @@ -102,18 +115,16 @@ describe('subagent thread Redis lifecycle', () => { expect.any(Function), { phase: 'pre-drain', priority: 100 }, ]); - expect(registerShutdownTask).toHaveBeenCalledWith( - 'subagent task control transport', + expect(taskStoreShutdownRegistration).toEqual([ + 'subagent task store', expect.any(Function), { priority: 90 }, - ); + ]); const prepare = activityPrepareRegistration[1]; prepare(); expect(mockTaskStore.prepareActivityForShutdown).toHaveBeenCalledTimes(1); - const shutdown = registerShutdownTask.mock.calls.find( - ([name]) => name === 'subagent task control transport', - )[1]; + const shutdown = taskStoreShutdownRegistration[1]; await shutdown(); expect(mockTaskStore.destroyTaskControlTransport).toHaveBeenCalledTimes(1); diff --git a/packages/api/src/agents/guard.spec.ts b/packages/api/src/agents/guard.spec.ts index e329acef46..c59366b396 100644 --- a/packages/api/src/agents/guard.spec.ts +++ b/packages/api/src/agents/guard.spec.ts @@ -37,6 +37,7 @@ function makeStore(): SubagentThreadTaskStore { deleteConvos: unused as AllMethods['deleteConvos'], deleteMessages: unused as AllMethods['deleteMessages'], getConvo: unused as AllMethods['getConvo'], + getSubagentTaskControlReplay: unused as AllMethods['getSubagentTaskControlReplay'], getMessages: unused as AllMethods['getMessages'], listActiveSubagentThreadLeases: unused as AllMethods['listActiveSubagentThreadLeases'], recordSubagentTaskControlReceipt: unused as AllMethods['recordSubagentTaskControlReceipt'], diff --git a/packages/api/src/agents/subagentTaskRouting.spec.ts b/packages/api/src/agents/subagentTaskRouting.spec.ts index 044dec0e23..21981e21e4 100644 --- a/packages/api/src/agents/subagentTaskRouting.spec.ts +++ b/packages/api/src/agents/subagentTaskRouting.spec.ts @@ -177,6 +177,7 @@ function taskHandler( claim: () => ({ status: 'not_found' }), control: () => ({ status: 'not_found' }), list: () => [], + retainsTaskOwnership: () => false, cancelScope: () => 0, ...overrides, }; @@ -475,6 +476,49 @@ describe('RedisSubagentTaskControlTransport', () => { await Promise.all([owner.destroy(), requester.destroy()]); }); + it('keeps a receipt-only owner registered until deletion cleanup can reach it', async () => { + const bus = new FakeRedisBus(); + const owner = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'owner', registrationHeartbeatMs: 5 }, + ); + const requester = new RedisSubagentTaskControlTransport( + asRedis(bus.createClient()), + asRedis(bus.createClient()), + { namespace: 'test', instanceId: 'requester', requestTimeoutMs: 100, retryDelayMs: 5 }, + ); + let receiptPending = true; + const cancelScope = jest.fn(() => 0); + await owner.bind( + taskHandler({ + retainsTaskOwnership: (_scopeId, taskId) => receiptPending && taskId === 'task-1', + cancelScope, + }), + ); + await requester.bind(taskHandler()); + await owner.registerTask('scope-1', 'task-1', 60_000); + + /** Model the SDK task/result buckets dropping the task and Redis losing the + * directory entry before the next owner heartbeat. Pending receipt work is + * the only remaining reason this process can still handle deletion cleanup. */ + bus.hashes.clear(); + for (let attempt = 0; attempt < 100 && !(await requester.hasTasks('scope-1')); attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await expect(requester.hasTasks('scope-1')).resolves.toBe(true); + await expect(requester.cancelScope('scope-1', null, ['deleted-child-thread'])).resolves.toBe(0); + expect(cancelScope).toHaveBeenCalledWith('scope-1', null, ['deleted-child-thread']); + + receiptPending = false; + bus.hashes.clear(); + for (let attempt = 0; attempt < 100 && (await requester.hasTasks('scope-1')); attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await expect(requester.hasTasks('scope-1')).resolves.toBe(false); + await Promise.all([owner.destroy(), requester.destroy()]); + }); + it('expires a dead owner independently while another owner keeps the scope active', async () => { const bus = new FakeRedisBus(); const deadOwner = new RedisSubagentTaskControlTransport( diff --git a/packages/api/src/agents/subagentTaskRouting.ts b/packages/api/src/agents/subagentTaskRouting.ts index f7b8eb502e..e1147db53a 100644 --- a/packages/api/src/agents/subagentTaskRouting.ts +++ b/packages/api/src/agents/subagentTaskRouting.ts @@ -36,6 +36,7 @@ const MAX_PROGRESS_LABEL_CHARS = 1_024; /** Bounds the model-facing task list, per owner reply and across the merged result. */ export const MAX_TASK_SNAPSHOTS = 200; const MAX_CANCEL_THREAD_IDS = 200; +const MAX_REMOVED_CONVERSATION_IDS = MAX_CANCEL_THREAD_IDS + 1; /** Matches the deletion drain so bounded fan-out stays well inside the lease TTL. */ const ROUTING_FANOUT_CONCURRENCY = 32; /** Contains every bounded response even when JSON escapes each retained character. */ @@ -97,7 +98,13 @@ type RoutedRequest = RoutedRequestBase & invocationId: string; } | { operation: 'list' } - | { operation: 'cancel'; threadIds: string[] | null } + | { + operation: 'cancel'; + threadIds: string[] | null; + /** Rows already committed as deleted by the requester. Owners must drop + * receipt retry work for these exact conversations after cancellation. */ + removedConversationIds?: string[]; + } ); type RoutedRequestPayload = @@ -110,7 +117,12 @@ type RoutedRequestPayload = invocationId: string; } | { operation: 'list'; scopeId: string } - | { operation: 'cancel'; scopeId: string; threadIds: string[] | null }; + | { + operation: 'cancel'; + scopeId: string; + threadIds: string[] | null; + removedConversationIds?: string[]; + }; interface RoutedResponse { version: typeof PROTOCOL_VERSION; @@ -174,7 +186,14 @@ export interface SubagentTaskControlHandler { invocationId: string, ): Promise | SubagentTaskControlResult; list(scopeId: string): SubagentTaskSnapshot[]; - cancelScope(scopeId: string, threadIds: string[] | null): number; + /** Receipt retry work can outlive the SDK task/result buckets. Keep its owner + * addressable so deletion can revoke work whose durable target was removed. */ + retainsTaskOwnership(scopeId: string, taskId: string): boolean; + cancelScope( + scopeId: string, + threadIds: string[] | null, + removedConversationIds?: string[], + ): number; } /** Optional host transport for reaching the process that owns a live child task. */ @@ -190,7 +209,11 @@ export interface SubagentTaskControlTransport { invocationId: string, ): Promise; list(scopeId: string): Promise; - cancelScope(scopeId: string, threadIds: string[] | null): Promise; + cancelScope( + scopeId: string, + threadIds: string[] | null, + removedConversationIds?: string[], + ): Promise; destroy(): Promise; } @@ -547,6 +570,7 @@ function parseRequest(value: unknown): RoutedRequest | undefined { taskId?: unknown; command?: unknown; threadIds?: unknown; + removedConversationIds?: unknown; invocationId?: unknown; expiresAt?: unknown; }; @@ -579,6 +603,14 @@ function parseRequest(value: unknown): RoutedRequest | undefined { if (candidate.threadIds !== null && !isCancelThreadIds(candidate.threadIds)) { return undefined; } + if ( + candidate.removedConversationIds !== undefined && + (!Array.isArray(candidate.removedConversationIds) || + candidate.removedConversationIds.length > MAX_REMOVED_CONVERSATION_IDS || + !candidate.removedConversationIds.every((id) => isBoundedString(id, MAX_THREAD_ID_CHARS))) + ) { + return undefined; + } return { version: PROTOCOL_VERSION, kind: 'request', @@ -588,6 +620,9 @@ function parseRequest(value: unknown): RoutedRequest | undefined { operation: 'cancel', scopeId: candidate.scopeId, threadIds: candidate.threadIds, + ...(candidate.removedConversationIds === undefined + ? {} + : { removedConversationIds: candidate.removedConversationIds }), }; } if (!isBoundedString(candidate.taskId, MAX_TASK_ID_CHARS)) { @@ -852,7 +887,11 @@ export class RedisSubagentTaskControlTransport implements SubagentTaskControlTra * predicate to its complete local task set, so deletion never depends on the * bounded model-facing list and cannot miss a task beyond that cap. */ - async cancelScope(scopeId: string, threadIds: string[] | null): Promise { + async cancelScope( + scopeId: string, + threadIds: string[] | null, + removedConversationIds: string[] = [], + ): Promise { this.assertScope(scopeId); if (threadIds != null && threadIds.length === 0) { return 0; @@ -880,11 +919,24 @@ export class RedisSubagentTaskControlTransport implements SubagentTaskControlTra } const cancelSlot = createConcurrencyLimiter(ROUTING_FANOUT_CONCURRENCY); const requests: Array> = []; + const allTargetThreadIds = threadIds == null ? null : new Set(threadIds); for (const ownerId of owners) { for (const batch of batches) { + const batchThreadIds = batch == null ? null : new Set(batch); + const removedForBatch = removedConversationIds.filter( + (conversationId) => + allTargetThreadIds == null || + !allTargetThreadIds.has(conversationId) || + batchThreadIds?.has(conversationId) === true, + ); requests.push( cancelSlot(() => - this.sendRequest(ownerId, { operation: 'cancel', scopeId, threadIds: batch }), + this.sendRequest(ownerId, { + operation: 'cancel', + scopeId, + threadIds: batch, + ...(removedForBatch.length === 0 ? {} : { removedConversationIds: removedForBatch }), + }), ), ); } @@ -1014,7 +1066,13 @@ export class RedisSubagentTaskControlTransport implements SubagentTaskControlTra truncated: tasks.length > bounded.length, }; } else if (request.operation === 'cancel') { - result = { cancelled: handler.cancelScope(request.scopeId, request.threadIds) }; + result = { + cancelled: handler.cancelScope( + request.scopeId, + request.threadIds, + request.removedConversationIds, + ), + }; } else if (request.operation === 'claim') { const claim = boundedClaim(handler.claim(request.scopeId, request.taskId)); replayable = consumesResult(claim); @@ -1286,7 +1344,8 @@ export class RedisSubagentTaskControlTransport implements SubagentTaskControlTra * address outlives the task itself until the result is acknowledged. */ if ( localTaskIds.has(taskId) || - this.claimReplays.entries.has(this.claimReplayKey(scopeId, taskId)) + this.claimReplays.entries.has(this.claimReplayKey(scopeId, taskId)) || + handler.retainsTaskOwnership(scopeId, taskId) ) { retained.push(registration); continue; diff --git a/packages/api/src/agents/subagentThreads.spec.ts b/packages/api/src/agents/subagentThreads.spec.ts index c56c78e4fe..f1f8c383e0 100644 --- a/packages/api/src/agents/subagentThreads.spec.ts +++ b/packages/api/src/agents/subagentThreads.spec.ts @@ -35,7 +35,7 @@ import { createSubagentThreadTaskStore, SubagentThreadTaskStore, } from './subagentThreads'; -import { SubagentTaskOwnerUnavailableError } from './subagentTaskRouting'; +import { controlFingerprint, SubagentTaskOwnerUnavailableError } from './subagentTaskRouting'; import { SUBAGENT_COMPLETION_DELIVERY } from './subagentDelivery'; import { createSubagentAttemptKey } from './subagentThreadIds'; import { SubagentActivityStream } from './subagentActivity'; @@ -110,10 +110,14 @@ class TestTaskControlTransport implements SubagentTaskControlTransport { return [...this.remoteOwners(scopeId)].flatMap((owner) => owner.handler?.list(scopeId) ?? []); } - async cancelScope(scopeId: string, threadIds: string[] | null): Promise { + async cancelScope( + scopeId: string, + threadIds: string[] | null, + removedConversationIds?: string[], + ): Promise { let cancelled = 0; for (const owner of this.remoteOwners(scopeId)) { - cancelled += owner.handler?.cancelScope(scopeId, threadIds) ?? 0; + cancelled += owner.handler?.cancelScope(scopeId, threadIds, removedConversationIds) ?? 0; } return cancelled; } @@ -195,7 +199,10 @@ async function waitForSettled( started: SubagentTaskStartResult, ): Promise { const accepted = requireAccepted(started); - for (let attempt = 0; attempt < 200; attempt += 1) { + /** Coverage shards can briefly starve this polling loop while Mongo-backed + * suites build indexes in parallel. Keep the assertion bounded without + * treating two seconds of runner contention as a task-lifecycle failure. */ + for (let attempt = 0; attempt < 1000; attempt += 1) { const task = store.get(scopeId, accepted.task.taskId); if (task != null && task.status !== 'running') { return; @@ -2026,7 +2033,7 @@ describe('SubagentThreadTaskStore', () => { const parentConversationId = randomUUID(); await saveParent(userId, parentConversationId); const hub = new TestTaskRoutingHub(); - const ownerStore = new ReceiptTestSubagentThreadTaskStore(methods); + const ownerStore = new SubagentThreadTaskStore(methods); const requesterStore = new SubagentThreadTaskStore(methods); await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); @@ -2035,9 +2042,13 @@ describe('SubagentThreadTaskStore', () => { const result = new Promise<{ content: string }>((resolve) => { finish = resolve; }); + let runtime: SubagentTaskRuntime | undefined; const started = ownerStore.start( taskRequest(config.scopeId, { - run: async () => result, + run: async (taskRuntime) => { + runtime = taskRuntime; + return result; + }, }), ); const taskId = requireAccepted(started).task.taskId; @@ -2052,6 +2063,7 @@ describe('SubagentThreadTaskStore', () => { await new Promise((resolve) => setTimeout(resolve, 10)); } expect(durableInput).toBeDefined(); + await waitUntil(() => runtime != null, 'the controlled child runtime'); const steer = { action: 'queue' as const, message: 'Verify the primary source too.' }; const routed = await requesterStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'); @@ -2080,15 +2092,14 @@ describe('SubagentThreadTaskStore', () => { const acceptedControlId = routed.status === 'accepted' && routed.controlId != null ? routed.controlId : undefined; expect(acceptedControlId).toBeDefined(); - const transitionTime = Date.now(); - ownerStore.emitControlReceiptForTest(config.scopeId, taskId, { - controlId: acceptedControlId as string, - action: 'queue', - status: 'applied', - createdAt: transitionTime - 1, - updatedAt: transitionTime, - boundary: 'turn', - }); + const receiptIndex = ( + ownerStore as unknown as { controlInvocationByReceipt: Map } + ).controlInvocationByReceipt; + expect(receiptIndex.size).toBe(1); + expect(runtime?.drain('turn')).toEqual([ + expect.objectContaining({ content: steer.message, source: 'steer' }), + ]); + expect(receiptIndex.size).toBe(0); await waitUntil( () => ownerStore.get(config.scopeId, taskId) != null, 'the owner task to remain available', @@ -2109,11 +2120,14 @@ describe('SubagentThreadTaskStore', () => { }), ]); - /** A delayed retry can replay accepted in memory but cannot downgrade the - * already-applied durable receipt. */ + /** A delayed retry reflects the applied authoritative transition and cannot + * downgrade the durable receipt to the original accepted snapshot. */ await expect( ownerStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'), - ).resolves.toEqual(routed); + ).resolves.toMatchObject({ + status: 'accepted', + task: { pendingControls: 0 }, + }); [durableInput] = await methods.getMessages( { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, '+subagentTask', @@ -2129,7 +2143,7 @@ describe('SubagentThreadTaskStore', () => { 'invocation-1', ), ).resolves.toMatchObject({ status: 'invalid' }); - expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(1); + expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(0); finish({ content: 'Cross-replica result.' }); await waitForSettled(ownerStore, config.scopeId, started); @@ -2139,11 +2153,444 @@ describe('SubagentThreadTaskStore', () => { ]); }); + it('replays a durable control after owner loss and rejects fingerprint reuse', async () => { + const userId = 'durable-control-replay-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const ownerStore = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(ownerStore, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = ownerStore.start(taskRequest(config.scopeId, { run: async () => result })); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil( + async () => + ( + await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ) + ).length === 1, + 'the durable task input', + ); + const command = { action: 'queue' as const, message: 'Keep the citation.' }; + await expect( + ownerStore.controlTask(config.scopeId, taskId, command, 'durable-invocation'), + ).resolves.toMatchObject({ status: 'accepted' }); + finish({ content: 'Done.' }); + await waitForSettled(ownerStore, config.scopeId, started); + /** Owner shutdown is the production durability boundary. Awaiting it is both + * stronger and less load-sensitive than polling Mongo while the async receipt + * tail is still settling under Jest coverage instrumentation. */ + await ownerStore.destroyTaskControlTransport(); + const [settledInput] = await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ); + expect(settledInput?.subagentTask?.controlReceipts).toContainEqual( + expect.objectContaining({ + invocationId: 'durable-invocation', + status: 'rejected', + reason: 'task_completed', + }), + ); + + const restartedStore = new SubagentThreadTaskStore(methods); + ( + restartedStore as unknown as { + taskControlTransport: { + control: () => Promise; + destroy: () => Promise; + }; + } + ).taskControlTransport = { + control: async () => { + throw new SubagentTaskOwnerUnavailableError(); + }, + destroy: async () => undefined, + }; + await expect( + restartedStore.controlTask(config.scopeId, taskId, command, 'durable-invocation'), + ).resolves.toMatchObject({ status: 'not_running' }); + await expect( + restartedStore.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Different command.' }, + 'durable-invocation', + ), + ).resolves.toMatchObject({ status: 'invalid' }); + await restartedStore.destroyTaskControlTransport(); + }); + + it('replays the task-wide durable pending-control count after owner loss', async () => { + const store = new SubagentThreadTaskStore(methods); + const { scopeId } = buildSubagentThreadTaskConfig(store, { + userId: 'durable-pending-count-user', + parentConversationId: randomUUID(), + }); + const command = { action: 'queue' as const, message: 'Keep both queued instructions.' }; + const now = new Date('2026-08-24T12:00:00.000Z'); + const replay = jest.spyOn(methods, 'getSubagentTaskControlReplay').mockResolvedValue({ + receipt: { + invocationId: 'pending-count-invocation', + fingerprint: controlFingerprint(command), + controlId: 'pending-count-control', + action: 'queue', + status: 'accepted', + createdAt: now, + updatedAt: now, + }, + task: { + taskId: 'pending-count-task', + threadId: randomUUID(), + subagentType: 'researcher', + status: 'running', + resultAvailable: false, + resultClaimed: false, + pendingControls: 2, + createdAt: now, + updatedAt: now, + }, + }); + + await expect( + store.controlTask(scopeId, 'pending-count-task', command, 'pending-count-invocation'), + ).resolves.toMatchObject({ + status: 'accepted', + task: { pendingControls: 2 }, + }); + + replay.mockRestore(); + await store.destroyTaskControlTransport(); + }); + + it('waits for a raced authoritative receipt generation before acknowledging control', async () => { + const userId = 'receipt-generation-race-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => (finish = resolve)); + let runtime: SubagentTaskRuntime | undefined; + const started = store.start( + taskRequest(config.scopeId, { + run: async (taskRuntime) => { + runtime = taskRuntime; + return result; + }, + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil( + async () => + ( + await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ) + ).length === 1, + 'the generation-race task seed', + ); + await waitUntil(() => runtime != null, 'the generation-race child runtime'); + + const originalRecord = methods.recordSubagentTaskControlReceipt.bind(methods); + let releaseAccepted!: () => void; + let releaseApplied!: () => void; + let acceptedEntered!: () => void; + let appliedEntered!: () => void; + const acceptedGate = new Promise((resolve) => (releaseAccepted = resolve)); + const appliedGate = new Promise((resolve) => (releaseApplied = resolve)); + const sawAccepted = new Promise((resolve) => (acceptedEntered = resolve)); + const sawApplied = new Promise((resolve) => (appliedEntered = resolve)); + const persistence = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockImplementation(async (args) => { + if (args.receipt.status === 'accepted') { + acceptedEntered(); + await acceptedGate; + } else if (args.receipt.status === 'applied') { + appliedEntered(); + await appliedGate; + } + return originalRecord(args); + }); + + const control = store.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Apply at the next boundary.' }, + 'generation-race-invocation', + ); + await sawAccepted; + expect(runtime?.drain('turn')).toHaveLength(1); + releaseAccepted(); + await sawApplied; + let acknowledged = false; + void control.then(() => (acknowledged = true)); + await Promise.resolve(); + expect(acknowledged).toBe(false); + + releaseApplied(); + await expect(control).resolves.toMatchObject({ + status: 'accepted', + task: { pendingControls: 0 }, + }); + + persistence.mockRestore(); + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + await store.destroyTaskControlTransport(); + }); + + it('refreshes retained cancellation result flags after durable collection', async () => { + const userId = 'retained-control-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { + run: async (runtime) => + new Promise((_resolve, reject) => { + runtime.signal.addEventListener('abort', () => reject(runtime.signal.reason), { + once: true, + }); + }), + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil( + async () => + ( + await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ) + ).length === 1, + 'the cancellation replay task seed', + ); + + await expect( + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, 'cancel-invocation'), + ).resolves.toMatchObject({ status: 'cancelled' }); + await waitForSettled(store, config.scopeId, started); + await waitUntil( + async () => + ( + await methods.getMessages({ + user: userId, + conversationId: threadId, + messageId: `${taskId}:assistant`, + }) + ).length === 1, + 'the durable cancelled result', + ); + await expect(store.claimTask(config.scopeId, taskId, 'poll-invocation')).resolves.toMatchObject( + { + status: 'cancelled', + task: { resultClaimed: true }, + }, + ); + + await expect( + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, 'cancel-invocation'), + ).resolves.toMatchObject({ + status: 'cancelled', + task: { resultAvailable: false, resultClaimed: true }, + }); + + await store.destroyTaskControlTransport(); + }); + + it('normalizes terminal replay storage failures at the owner boundary', async () => { + const userId = 'terminal-replay-storage-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)); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + const replay = jest + .spyOn(methods, 'getSubagentTaskControlReplay') + .mockRejectedValue(new Error('database unavailable')); + try { + await expect( + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, 'storage-invocation'), + ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); + } finally { + replay.mockRestore(); + await store.destroyTaskControlTransport(); + } + }); + + it('reserves one durable fingerprint before concurrent owners can apply controls', async () => { + const userId = 'concurrent-control-reservation-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let runtime: SubagentTaskRuntime | undefined; + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start( + taskRequest(config.scopeId, { + run: async (taskRuntime) => { + runtime = taskRuntime; + return result; + }, + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitUntil(() => runtime != null, 'the concurrent-control child runtime'); + + const results = await Promise.all([ + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, 'shared-invocation'), + store.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Only apply if this fingerprint wins.' }, + 'shared-invocation', + ), + ]); + expect(results.filter((candidate) => candidate.status === 'invalid')).toHaveLength(1); + const winner = results.find((candidate) => candidate.status !== 'invalid'); + expect(winner?.status === 'cancelled' || winner?.status === 'accepted').toBe(true); + if (winner?.status === 'accepted') { + expect(runtime?.signal.aborted).toBe(false); + expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(1); + } else { + expect(runtime?.signal.aborted).toBe(true); + } + + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + await store.destroyTaskControlTransport(); + }); + + it('bounds concurrent durable reservation writers below the storage CAS retry limit', async () => { + const userId = 'bounded-control-reservation-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { maxControlsPerTask: 100 }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start(taskRequest(config.scopeId, { run: async () => result })); + const taskId = requireAccepted(started).task.taskId; + await waitUntil(() => store.get(config.scopeId, taskId)?.status === 'running', 'running task'); + + let activeReservations = 0; + let maxActiveReservations = 0; + let releaseReservations = (): void => undefined; + const reservationGate = new Promise((resolve) => { + releaseReservations = resolve; + }); + const persist = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockImplementation(async ({ receipt }) => { + if (receipt.status !== 'reserved') return true; + activeReservations += 1; + maxActiveReservations = Math.max(maxActiveReservations, activeReservations); + await reservationGate; + activeReservations -= 1; + return true; + }); + try { + const controls = Array.from({ length: 65 }, (_, index) => + store.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: `Control ${index}` }, + `bounded-reservation-${index}`, + ), + ); + await waitUntil(() => activeReservations === 32, 'the reservation writer bound'); + expect(maxActiveReservations).toBe(32); + releaseReservations(); + await expect(Promise.all(controls)).resolves.toHaveLength(65); + expect(maxActiveReservations).toBe(32); + } finally { + releaseReservations(); + persist.mockRestore(); + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + await store.destroyTaskControlTransport(); + } + }); + + it('bounds retained terminal control invocations while persisting their receipts', async () => { + const userId = 'terminal-control-window-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { run: async () => ({ content: 'Done.' }) }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + for (let index = 0; index < 65; index += 1) { + await expect( + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, `terminal-${index}`), + ).resolves.toMatchObject({ status: 'not_running' }); + } + const retained = (store as unknown as { terminalControlInvocations: Map }) + .terminalControlInvocations; + expect(retained.size).toBe(64); + const replayLookup = jest + .spyOn(methods, 'getSubagentTaskControlReplay') + .mockRejectedValue(new Error('database unavailable')); + await expect( + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, 'terminal-64'), + ).resolves.toMatchObject({ status: 'not_running' }); + expect(replayLookup).not.toHaveBeenCalled(); + replayLookup.mockRestore(); + await store.destroyTaskControlTransport(); + }); + + it('refuses to evict an unpersisted terminal invocation', async () => { + const userId = 'terminal-control-admission-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + const started = store.start( + taskRequest(config.scopeId, { run: async () => ({ content: 'Done.' }) }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitForSettled(store, config.scopeId, started); + + for (let index = 0; index < 64; index += 1) { + expect( + store.controlInvocation(config.scopeId, taskId, { action: 'cancel' }, `pending-${index}`), + ).toMatchObject({ status: 'not_running' }); + } + expect( + store.controlInvocation(config.scopeId, taskId, { action: 'cancel' }, 'pending-64'), + ).toMatchObject({ status: 'invalid' }); + const retained = (store as unknown as { terminalControlInvocations: Map }) + .terminalControlInvocations; + expect(retained.size).toBe(64); + await store.destroyTaskControlTransport(); + }); + it('fails acceptance closed when the durable receipt target is not ready', async () => { const userId = 'receipt-target-user'; const parentConversationId = randomUUID(); await saveParent(userId, parentConversationId); - const store = new SubagentThreadTaskStore(methods, { controlReceiptRetryMs: 10 }); + const store = new SubagentThreadTaskStore(methods, { controlReceiptRetryMs: 60_000 }); const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); let finish = (_value: { content: string }): void => undefined; const result = new Promise<{ content: string }>((resolve) => { @@ -2163,22 +2610,63 @@ describe('SubagentThreadTaskStore', () => { 'the durable task input', ); - const persistReceipt = methods.recordSubagentTaskControlReceipt; const persistence = jest .spyOn(methods, 'recordSubagentTaskControlReceipt') - .mockResolvedValueOnce(false) - .mockImplementation(persistReceipt); + .mockRejectedValue(new Error('database unavailable')); const command = { action: 'queue' as const, message: 'Check the source.' }; try { await expect( store.controlTask(config.scopeId, taskId, command, 'not-ready-invocation'), ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); - expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(1); + expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(0); + + await expect( + store.controlTask(config.scopeId, taskId, command, 'not-ready-invocation'), + ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); + expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(0); + persistence.mockRestore(); await expect( store.controlTask(config.scopeId, taskId, command, 'not-ready-invocation'), ).resolves.toMatchObject({ status: 'accepted' }); expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(1); + } finally { + if (jest.isMockFunction(methods.recordSubagentTaskControlReceipt)) persistence.mockRestore(); + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + await store.destroyTaskControlTransport(); + } + }); + + it('drops a permanent receipt conflict and rolls back its queued control', async () => { + const userId = 'receipt-conflict-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start(taskRequest(config.scopeId, { run: async () => result })); + const taskId = requireAccepted(started).task.taskId; + await waitUntil( + () => store.get(config.scopeId, taskId)?.threadId != null, + 'thread preparation', + ); + const persistence = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockResolvedValue('conflict'); + const command = { action: 'queue' as const, message: 'Check the source.' }; + try { + await expect( + store.controlTask(config.scopeId, taskId, command, 'conflicting-invocation'), + ).resolves.toMatchObject({ status: 'invalid' }); + expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(0); + await expect( + store.controlTask(config.scopeId, taskId, command, 'conflicting-invocation'), + ).resolves.toMatchObject({ status: 'invalid' }); + expect(persistence).toHaveBeenCalledTimes(2); } finally { persistence.mockRestore(); finish({ content: 'Done.' }); @@ -2187,6 +2675,128 @@ describe('SubagentThreadTaskStore', () => { } }); + it('rejects a durable fingerprint conflict before applying task cancellation', async () => { + const userId = 'receipt-preflight-conflict-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let runtime: SubagentTaskRuntime | undefined; + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start( + taskRequest(config.scopeId, { + run: async (taskRuntime) => { + runtime = taskRuntime; + return result; + }, + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil(() => runtime != null, 'the child runtime to start'); + await waitUntil( + async () => + ( + await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ) + ).length === 1, + 'the durable task input', + ); + await expect( + methods.recordSubagentTaskControlReceipt({ + userId, + conversationId: threadId, + taskId, + receipt: { + invocationId: 'conflicting-cancel', + fingerprint: controlFingerprint({ action: 'queue', message: 'Original.' }), + action: 'queue', + status: 'rejected', + reason: 'withdrawn', + createdAt: new Date(), + updatedAt: new Date(), + }, + }), + ).resolves.toBe(true); + + await expect( + store.controlTask(config.scopeId, taskId, { action: 'cancel' }, 'conflicting-cancel'), + ).resolves.toMatchObject({ status: 'invalid' }); + expect(runtime?.signal.aborted).toBe(false); + expect(store.get(config.scopeId, taskId)?.status).toBe('running'); + + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + await store.destroyTaskControlTransport(); + }); + + it('does not replay a control whose prior owner only reserved its invocation', async () => { + const userId = 'abandoned-control-reservation-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let runtime: SubagentTaskRuntime | undefined; + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start( + taskRequest(config.scopeId, { + run: async (taskRuntime) => { + runtime = taskRuntime; + return result; + }, + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil(() => runtime != null, 'the child runtime to start'); + await waitUntil( + async () => + ( + await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ) + ).length === 1, + 'the durable task input', + ); + const command = { action: 'queue' as const, message: 'Apply this once.' }; + await expect( + methods.recordSubagentTaskControlReceipt({ + userId, + conversationId: threadId, + taskId, + receipt: { + invocationId: 'abandoned-reservation', + fingerprint: controlFingerprint(command), + action: 'queue', + status: 'reserved', + message: command.message, + createdAt: new Date(), + updatedAt: new Date(), + }, + }), + ).resolves.toBe(true); + + /** The reservation proves a prior owner may have crossed the side-effect + * boundary. A retry must neither apply it again nor report false acceptance. */ + await expect( + store.controlTask(config.scopeId, taskId, command, 'abandoned-reservation'), + ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); + expect(store.get(config.scopeId, taskId)?.pendingControls).toBe(0); + + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + await store.destroyTaskControlTransport(); + }); + it('retains only the bounded public projection of a control payload', async () => { const userId = 'bounded-control-payload-user'; const parentConversationId = randomUUID(); @@ -2228,6 +2838,13 @@ describe('SubagentThreadTaskStore', () => { action: 'queue', message: 'x'.repeat(4 * 1024), }); + await waitUntil(async () => { + const [input] = await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ); + return input?.subagentTask?.controlReceipts?.[0]?.messageTruncated === true; + }, 'the bounded receipt truncation marker'); await expect( store.controlTask( config.scopeId, @@ -2298,7 +2915,10 @@ describe('SubagentThreadTaskStore', () => { 'tenant-invocation', ), ).resolves.toMatchObject({ status: 'accepted' }); - expect(contexts).toEqual([{ tenantId, userId }]); + expect(contexts).toEqual([ + { tenantId, userId }, + { tenantId, userId }, + ]); } finally { persistence.mockRestore(); finish({ content: 'Done.' }); @@ -2475,24 +3095,92 @@ describe('SubagentThreadTaskStore', () => { } }); - it('flushes a pending control receipt once during graceful shutdown', async () => { + it('retries a terminal receipt after result collection removes the local task', async () => { + const userId = 'receipt-after-claim-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods, { controlReceiptRetryMs: 5 }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let finish = (_value: { content: string }): void => undefined; + const result = new Promise<{ content: string }>((resolve) => { + finish = resolve; + }); + const started = store.start(taskRequest(config.scopeId, { run: async () => result })); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil( + async () => + ( + await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ) + ).length === 1, + 'the claim-race durable task input', + ); + await store.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Persist the terminal transition after collection.' }, + 'claim-race-invocation', + ); + + const persistReceipt = methods.recordSubagentTaskControlReceipt; + const persistence = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockRejectedValueOnce(new Error('database temporarily unavailable')) + .mockImplementation(persistReceipt); + try { + finish({ content: 'Done.' }); + await waitForSettled(store, config.scopeId, started); + expect(store.claim(config.scopeId, taskId)).toMatchObject({ status: 'completed' }); + await waitUntil( + () => persistence.mock.calls.length >= 1, + 'the terminal receipt persistence failure', + ); + /** Model the SDK retention timer expiring the claimed result before the + * storage retry fires. Pending durability must not depend on this bucket. */ + const bucket = ( + store as unknown as { + buckets: Map }>; + } + ).buckets.get(config.scopeId); + bucket?.tasks.delete(taskId); + expect(store.get(config.scopeId, taskId)).toBeUndefined(); + await waitUntil(async () => { + const [input] = await methods.getMessages( + { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, + '+subagentTask', + ); + return input?.subagentTask?.controlReceipts?.[0]?.reason === 'task_completed'; + }, 'the post-collection terminal receipt retry'); + } finally { + persistence.mockRestore(); + await store.destroyTaskControlTransport(); + } + }); + + it('quiesces receipt producers and flushes their final transition during shutdown', async () => { const userId = 'receipt-shutdown-user'; const parentConversationId = randomUUID(); await saveParent(userId, parentConversationId); const store = new ReceiptTestSubagentThreadTaskStore(methods, { controlReceiptRetryMs: 60_000, + shutdownControlReceiptBackoffMs: 1, }); const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); let runtime: SubagentTaskRuntime | undefined; - let finish = (_value: { content: string }): void => undefined; - const result = new Promise<{ content: string }>((resolve) => { - finish = resolve; - }); const started = store.start( taskRequest(config.scopeId, { run: async (taskRuntime) => { runtime = taskRuntime; - return result; + return await new Promise<{ content: string }>((_resolve, reject) => { + taskRuntime.signal.addEventListener( + 'abort', + () => reject(new Error('provider stopped after task cancellation')), + { once: true }, + ); + }); }, }), ); @@ -2512,33 +3200,163 @@ describe('SubagentThreadTaskStore', () => { .spyOn(methods, 'recordSubagentTaskControlReceipt') .mockRejectedValueOnce(new Error('database temporarily unavailable')) .mockImplementation(persistReceipt); - const controlId = accepted.status === 'accepted' ? accepted.controlId : undefined; - expect(controlId).toBeDefined(); - store.emitControlReceiptForTest(config.scopeId, taskId, { - controlId: controlId as string, - action: 'queue', - status: 'applied', - createdAt: Date.now() - 1, - updatedAt: Date.now(), - boundary: 'turn', + const saveMessage = methods.saveMessage; + let releaseCancellation = (): void => undefined; + const cancellationGate = new Promise((resolve) => { + releaseCancellation = resolve; + }); + let cancellationSaveStarted = false; + const settlement = jest.spyOn(methods, 'saveMessage').mockImplementation(async (...args) => { + const message = args[1] as IMessage; + if ( + message.messageId === `${taskId}:assistant` && + message.subagentTask?.status === 'cancelled' + ) { + cancellationSaveStarted = true; + await cancellationGate; + } + return saveMessage(...args); }); - await waitUntil(() => persistence.mock.calls.length === 1, 'the failed receipt write'); - await store.destroyTaskControlTransport(); + const shutdown = store.destroyTaskControlTransport(); + await waitUntil(() => cancellationSaveStarted, 'the cancellation settlement to start'); + let shutdownResolved = false; + void shutdown.then(() => { + shutdownResolved = true; + }); + await Promise.resolve(); + expect(shutdownResolved).toBe(false); + releaseCancellation(); + await shutdown; + expect(runtime?.signal.aborted).toBe(true); const [input] = await methods.getMessages( { user: userId, conversationId: threadId, messageId: `${taskId}:user` }, '+subagentTask', ); expect(input?.subagentTask?.controlReceipts).toContainEqual( - expect.objectContaining({ invocationId: 'shutdown-invocation', status: 'applied' }), + expect.objectContaining({ + invocationId: 'shutdown-invocation', + status: 'rejected', + reason: 'task_cancelled', + }), ); expect(persistence).toHaveBeenCalledTimes(2); persistence.mockRestore(); - finish({ content: 'Done.' }); + settlement.mockRestore(); + expect(store.get(config.scopeId, taskId)?.status).toBe('cancelled'); + }); + + it('fails graceful shutdown when terminal receipts remain unavailable', async () => { + const userId = 'receipt-shutdown-failure-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new ReceiptTestSubagentThreadTaskStore(methods, { + shutdownControlReceiptBackoffMs: 1, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let runtime: SubagentTaskRuntime | undefined; + const started = store.start( + taskRequest(config.scopeId, { + run: async (taskRuntime) => { + runtime = taskRuntime; + return await new Promise<{ content: string }>((_resolve, reject) => { + taskRuntime.signal.addEventListener('abort', () => reject(new Error('stopped')), { + once: true, + }); + }); + }, + }), + ); + const taskId = requireAccepted(started).task.taskId; + await waitUntil(() => runtime != null, 'the child runtime to start'); + await store.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Persist me.' }, + 'shutdown-failure-invocation', + ); + const persistence = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockRejectedValue(new Error('database unavailable')); + + await expect(store.destroyTaskControlTransport()).rejects.toBeInstanceOf( + SubagentTaskOwnerUnavailableError, + ); + expect(persistence.mock.calls.length).toBeGreaterThanOrEqual(4); + persistence.mockRestore(); + await store.destroyTaskControlTransport(); await waitForSettled(store, config.scopeId, started); }); + it('drops permanently unwritable control receipt work after its conversations are deleted', async () => { + const userId = 'deleted-control-receipt-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new ReceiptTestSubagentThreadTaskStore(methods, { + controlReceiptRetryMs: 60_000, + shutdownControlReceiptBackoffMs: 1, + }); + const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + let runtime: SubagentTaskRuntime | undefined; + const started = store.start( + taskRequest(config.scopeId, { + run: async (taskRuntime) => { + runtime = taskRuntime; + return await new Promise<{ content: string }>((_resolve, reject) => { + taskRuntime.signal.addEventListener('abort', () => reject(new Error('deleted')), { + once: true, + }); + }); + }, + }), + ); + const taskId = requireAccepted(started).task.taskId; + const threadId = requireThreadId(started); + await waitUntil(() => runtime != null, 'the child runtime to start'); + const accepted = await store.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'This receipt target will be deleted.' }, + 'deleted-receipt-invocation', + ); + expect(accepted).toMatchObject({ status: 'accepted' }); + const controlId = accepted.status === 'accepted' ? accepted.controlId : undefined; + expect(controlId).toBeDefined(); + + const persistence = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockRejectedValue(new Error('receipt target deleted')); + try { + const appliedAt = Date.now(); + store.emitControlReceiptForTest(config.scopeId, taskId, { + controlId: controlId as string, + action: 'queue', + status: 'applied', + createdAt: appliedAt - 1, + updatedAt: appliedAt, + boundary: 'turn', + }); + await waitUntil( + () => persistence.mock.calls.length >= 1, + 'the permanently unwritable receipt', + ); + + const plan = await store.planCancellationForConversations(userId, [parentConversationId]); + await methods.deleteConvos(userId, { conversationId: parentConversationId }); + await expect( + store.cancelPlan(plan, [parentConversationId, threadId]), + ).resolves.toBeGreaterThanOrEqual(1); + await waitForSettled(store, config.scopeId, started); + } finally { + persistence.mockRestore(); + } + + /** Deletion is a terminal storage outcome, so shutdown must not keep retrying + * receipt rows whose authorized parent or child no longer exists. */ + await expect(store.destroyTaskControlTransport()).resolves.toBeUndefined(); + }); + it('fails a child closed when its owner address cannot be published', async () => { const userId = 'unregistered-user'; const parentConversationId = randomUUID(); @@ -2900,7 +3718,7 @@ describe('SubagentThreadTaskStore', () => { }); }); - it('keeps a live task’s control invocation when settled tasks fill the window', async () => { + it('does not spend live replay slots on controls rejected by settled tasks', async () => { const userId = 'invocation-eviction-user'; const parentConversationId = randomUUID(); await saveParent(userId, parentConversationId); @@ -2924,7 +3742,16 @@ describe('SubagentThreadTaskStore', () => { const steer = { action: 'queue' as const, message: 'Verify the primary source too.' }; const applied = store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-live'); expect(applied).toMatchObject({ status: 'accepted' }); - store.controlInvocation(config.scopeId, settledTaskId, steer, 'invocation-settled'); + expect( + store.controlInvocation(config.scopeId, settledTaskId, steer, 'invocation-settled'), + ).toMatchObject({ status: 'not_running' }); + expect( + ( + store as unknown as { + controlInvocations: Map; + } + ).controlInvocations.size, + ).toBe(1); for (let attempt = 0; attempt < 100; attempt += 1) { if (store.get(config.scopeId, settledTaskId) == null) { @@ -2934,9 +3761,8 @@ describe('SubagentThreadTaskStore', () => { } expect(store.get(config.scopeId, settledTaskId)).toBeUndefined(); - /** The window is full, so admitting another invocation sweeps the records of tasks - * this store no longer holds. The live task's record survives, so a caller - * retrying it replays instead of steering that child a second time. */ + /** A second live invocation fills the window. The first live record survives, so + * a caller retrying it replays instead of steering that child a second time. */ store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-later'); expect(store.controlInvocation(config.scopeId, liveTaskId, steer, 'invocation-live')).toEqual( applied, @@ -2953,6 +3779,7 @@ describe('SubagentThreadTaskStore', () => { finish({ content: 'done' }); await waitForSettled(store, config.scopeId, live); + await store.destroyTaskControlTransport(); }); it('caps the merged local and remote task list the poll tool reads', async () => { @@ -3014,9 +3841,11 @@ describe('SubagentThreadTaskStore', () => { expect(store.controlInvocation(config.scopeId, liveTaskId, steer, 'local-1')).toMatchObject({ status: 'accepted', }); + const replayLookup = jest.spyOn(methods, 'getSubagentTaskControlReplay'); /** The window holds a live task's record and cannot be swept, but a task this - * replica never owned is the remote owner's to refuse or apply. */ + * replica never owned is the remote owner's to refuse or apply. Only that owner + * performs the durable preflight; the requester does not repeat the Mongo read. */ await expect( store.controlTask(config.scopeId, 'remote-task', { action: 'cancel' }, 'remote-1'), ).resolves.toEqual(remoteResult); @@ -3026,12 +3855,39 @@ describe('SubagentThreadTaskStore', () => { { action: 'cancel' }, 'remote-1', ); + expect(replayLookup).not.toHaveBeenCalled(); + replayLookup.mockRestore(); finish({ content: 'done' }); await waitForSettled(store, config.scopeId, live); await store.destroyTaskControlTransport(); }); + it('normalizes a storage outage after routed owner loss', async () => { + const userId = 'remote-fallback-outage-user'; + const parentConversationId = randomUUID(); + await saveParent(userId, parentConversationId); + const store = new SubagentThreadTaskStore(methods); + const { scopeId } = buildSubagentThreadTaskConfig(store, { userId, parentConversationId }); + await store.configureTaskControlTransport({ + ...replayTransport({ status: 'not_found' }), + control: async () => { + throw new SubagentTaskOwnerUnavailableError(); + }, + }); + const replayLookup = jest + .spyOn(methods, 'getSubagentTaskControlReplay') + .mockRejectedValue(new Error('database unavailable')); + + await expect( + store.controlTask(scopeId, 'remote-task', { action: 'cancel' }, 'remote-outage'), + ).rejects.toBeInstanceOf(SubagentTaskOwnerUnavailableError); + expect(replayLookup).toHaveBeenCalledTimes(1); + + replayLookup.mockRestore(); + await store.destroyTaskControlTransport(); + }); + it('fails a deletion closed when the admission fence cannot be held', async () => { const userId = 'fence-lapse-user'; const parentConversationId = randomUUID(); @@ -3299,7 +4155,10 @@ describe('SubagentThreadTaskStore', () => { const parentConversationId = randomUUID(); await saveParent(userId, parentConversationId); const hub = new TestTaskRoutingHub(); - const ownerStore = new SubagentThreadTaskStore(methods); + const ownerStore = new ReceiptTestSubagentThreadTaskStore(methods, { + controlReceiptRetryMs: 60_000, + shutdownControlReceiptBackoffMs: 1, + }); const deletingStore = new SubagentThreadTaskStore(methods); await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); await deletingStore.configureTaskControlTransport(new TestTaskControlTransport(hub)); @@ -3324,16 +4183,37 @@ describe('SubagentThreadTaskStore', () => { } expect(await methods.getConvo(userId, threadId)).not.toBeNull(); + const accepted = await ownerStore.controlTask( + config.scopeId, + taskId, + { action: 'queue', message: 'Persist this before the child is deleted.' }, + 'remote-deleted-receipt', + ); + const controlId = accepted.status === 'accepted' ? accepted.controlId : undefined; + expect(controlId).toBeDefined(); + const persistence = jest + .spyOn(methods, 'recordSubagentTaskControlReceipt') + .mockRejectedValue(new Error('receipt target deleted')); + ownerStore.emitControlReceiptForTest(config.scopeId, taskId, { + controlId: controlId as string, + action: 'queue', + status: 'applied', + createdAt: Date.now() - 1, + updatedAt: Date.now(), + boundary: 'turn', + }); + await waitUntil(() => persistence.mock.calls.length > 0, 'the remote receipt write to fail'); + /** The parent survives this deletion, so the child's own thread is the only target. */ const plan = await deletingStore.planCancellationForConversations(userId, [threadId]); - await expect(deletingStore.cancelPlan(plan)).resolves.toBe(1); + await methods.deleteConvos(userId, { conversationId: threadId }); + await expect(deletingStore.cancelPlan(plan, [threadId])).resolves.toBe(1); await waitForSettled(ownerStore, config.scopeId, started); expect(ownerStore.get(config.scopeId, taskId)).toMatchObject({ status: 'cancelled' }); + persistence.mockRestore(); - await Promise.all([ - ownerStore.destroyTaskControlTransport(), - deletingStore.destroyTaskControlTransport(), - ]); + await expect(ownerStore.destroyTaskControlTransport()).resolves.toBeUndefined(); + await deletingStore.destroyTaskControlTransport(); }); it('cancels a child admitted after the deletion snapshot from its durable lease', async () => { diff --git a/packages/api/src/agents/subagentThreads.ts b/packages/api/src/agents/subagentThreads.ts index d4fdb30eae..0474cbf4ba 100644 --- a/packages/api/src/agents/subagentThreads.ts +++ b/packages/api/src/agents/subagentThreads.ts @@ -63,8 +63,16 @@ const DEFAULT_OWNER_DRAIN_POLL_MS = 100; const DELETION_CANCEL_CONCURRENCY = 32; /** Bounds retained control invocations; one entry per applied command. */ const MAX_CONTROL_INVOCATIONS = 4_096; +/** Terminal controls are side-effect free, but retaining one bounded window + * prevents duplicate storage writers while preserving recent retry replay. */ +const MAX_TERMINAL_CONTROL_INVOCATIONS = 64; +/** Keep same-task durable reservations below the storage CAS retry bound. Receipt + * finalization is serialized per task separately, leaving ample collision headroom. */ +const CONTROL_RESERVATION_CONCURRENCY = 32; const MAX_DURABLE_CONTROL_MESSAGE_CHARS = 4 * 1024; const DEFAULT_CONTROL_RECEIPT_RETRY_MS = 5_000; +const SHUTDOWN_CONTROL_RECEIPT_FLUSH_ATTEMPTS = 4; +const DEFAULT_SHUTDOWN_CONTROL_RECEIPT_BACKOFF_MS = 1_000; /** Bounds retained live-only updates while an event transport is unavailable. */ const MAX_PENDING_ACTIVITY_EVENTS = 32; /** Live activity must never delay terminal notification indefinitely. */ @@ -89,6 +97,7 @@ const DURABLE_RESULT_SELECT = class SubagentThreadPublicError extends Error {} class SubagentThreadDeletedError extends SubagentThreadPublicError {} +class SubagentControlReceiptConflictError extends Error {} type SubagentThreadMethods = Pick< AllMethods, @@ -98,6 +107,7 @@ type SubagentThreadMethods = Pick< | 'deleteConvos' | 'deleteMessages' | 'getConvo' + | 'getSubagentTaskControlReplay' | 'getMessages' | 'listActiveSubagentThreadLeases' | 'reserveSubagentThread' @@ -156,18 +166,31 @@ type ControlInvocationRecord = { invocationId: string; fingerprint: string; command: SubagentTaskControlCommand; + commandMessageTruncated: boolean; result: SubagentTaskControlResult; createdAt: number; /** Last authoritative SDK transition, retained for idempotent retries even * after the bounded SDK snapshot evicts older receipt history. */ receipt?: ISubagentTaskControlReceipt; + /** True only after this invocation's current receipt is durable and therefore + * safe to evict from the bounded process-local replay window. */ + receiptPersisted?: boolean; + /** The current durable write, shared by same-invocation retries so a caller + * cannot observe success before the authoritative receipt is committed. */ + receiptPersistence?: Promise; }; +const hasDurableControlReceipt = (invocation: ControlInvocationRecord): boolean => + invocation.receiptPersisted === true; + interface TaskThreadLease { + scopeId: string; idempotencyKey: string; taskId: string; running: boolean; settling: boolean; + /** Resolves only after child persistence and lease cleanup finish. */ + execution?: Promise; /** Ordered observational tail; canonical child settlement never awaits it. */ activityTail?: Promise; activityPending?: number; @@ -220,6 +243,7 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor isOwnerActive?: (userId: string) => Promise; maxControlInvocations?: number; controlReceiptRetryMs?: number; + shutdownControlReceiptBackoffMs?: number; ownerFenceGraceMs?: number; fenceOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; renewOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise; @@ -477,13 +501,19 @@ function controlReceiptKey(scopeId: string, taskId: string, controlId: string): return `${scopeId}\u0000${taskId}\u0000${controlId}`; } -function boundedControlMessage(command: SubagentTaskControlCommand): { +function boundedControlMessage( + command: SubagentTaskControlCommand, + alreadyTruncated = false, +): { message?: string; messageTruncated?: boolean; } { if (!('message' in command)) return {}; if (command.message.length <= MAX_DURABLE_CONTROL_MESSAGE_CHARS) { - return { message: command.message }; + return { + message: command.message, + ...(alreadyTruncated ? { messageTruncated: true } : {}), + }; } return { message: command.message.slice(0, MAX_DURABLE_CONTROL_MESSAGE_CHARS), @@ -522,6 +552,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { readonly supportsThreadContinuation = true; private readonly activeThreads = new Map(); private readonly controlInvocations = new Map(); + private readonly terminalControlInvocations = new Map(); private readonly controlInvocationByReceipt = new Map(); private readonly pendingControlReceipts = new Map< @@ -531,7 +562,12 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { private readonly controlPersistenceTails = new Map>(); private readonly controlPersistenceRetryTimers = new Map>(); + private readonly controlReservationSlot = createConcurrencyLimiter( + CONTROL_RESERVATION_CONCURRENCY, + ); + private controlPersistenceStopping = false; + private controlCommandAdmissionClosed = false; private readonly parentPersistence = new Map>(); private readonly maxThreadDepth: number; @@ -542,6 +578,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { private readonly taskRoutingTtlMs: number; private readonly maxControlInvocations: number; private readonly controlReceiptRetryMs: number; + private readonly shutdownControlReceiptBackoffMs: number; private readonly ownerFenceGraceMs: number; private readonly isOwnerActive: (userId: string) => Promise; private readonly fenceOwnerAdmission?: ( @@ -590,6 +627,10 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { options.controlReceiptRetryMs, DEFAULT_CONTROL_RECEIPT_RETRY_MS, ); + this.shutdownControlReceiptBackoffMs = positiveInteger( + options.shutdownControlReceiptBackoffMs, + DEFAULT_SHUTDOWN_CONTROL_RECEIPT_BACKOFF_MS, + ); this.ownerFenceGraceMs = positiveInteger(options.ownerFenceGraceMs, OWNER_FENCE_GRACE_MS); this.isOwnerActive = options.isOwnerActive ?? (async () => true); this.fenceOwnerAdmission = options.fenceOwnerAdmission; @@ -601,16 +642,107 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { /** Receives payload-free authoritative transitions from the SDK task store. */ protected onControlReceipt(scopeId: string, taskId: string, receipt: SdkControlReceipt): void { + const persistence = this.queueAuthoritativeControlReceipt(scopeId, taskId, receipt); + void persistence?.catch((error) => { + logger.warn('[subagentThreads] Failed to persist a child control transition', error); + }); + } + + private queueAuthoritativeControlReceipt( + scopeId: string, + taskId: string, + receipt: SdkControlReceipt, + ): Promise | undefined { const invocation = this.controlInvocationByReceipt.get( controlReceiptKey(scopeId, taskId, receipt.controlId), ); const threadId = this.get(scopeId, taskId)?.threadId; - if (invocation == null || threadId == null) return; + if (invocation == null || threadId == null) return undefined; const durable = this.durableReceipt(invocation, receipt); invocation.receipt = durable; - void this.queueControlReceipt(scopeId, taskId, threadId, durable).catch((error) => { - logger.warn('[subagentThreads] Failed to persist a child control transition', error); + invocation.result = this.controlResultFromReceipt( + invocation, + durable, + this.get(scopeId, taskId)?.pendingControls, + ); + if (receipt.status !== 'accepted') { + this.controlInvocationByReceipt.delete(controlReceiptKey(scopeId, taskId, receipt.controlId)); + } + invocation.receiptPersisted = false; + const persistence = this.queueControlReceipt(scopeId, taskId, threadId, durable).then(() => { + if (invocation.receipt === durable) invocation.receiptPersisted = true; }); + const tracked = persistence.finally(() => { + if (invocation.receiptPersistence === tracked) invocation.receiptPersistence = undefined; + }); + invocation.receiptPersistence = tracked; + return tracked; + } + + /** Keeps same-process retries aligned with the durable receipt ledger. The SDK + * can replace an accepted control with a terminal transition after the child + * settles, so the originally returned result is no longer authoritative. */ + private controlResultFromReceipt( + invocation: ControlInvocationRecord, + receipt: ISubagentTaskControlReceipt, + pendingControls?: number, + ): SubagentTaskControlResult { + const current = invocation.result; + if (!('task' in current)) return current; + let terminalStatus = current.task.status; + if ( + (receipt.action === 'cancel' && receipt.status === 'applied') || + receipt.reason === 'task_cancelled' + ) { + terminalStatus = 'cancelled'; + } else if (receipt.reason === 'task_completed') { + terminalStatus = 'completed'; + } else if (receipt.reason === 'task_failed') { + terminalStatus = 'error'; + } + const task: SubagentTaskSnapshot = { + ...current.task, + status: terminalStatus, + updatedAt: receipt.updatedAt.getTime(), + /** A receipt can make cancellation authoritative before the assistant row + * exists. Preserve actual result materialization rather than inferring it. */ + resultAvailable: current.task.resultAvailable, + pendingControls: pendingControls ?? current.task.pendingControls, + }; + if (receipt.status === 'accepted') { + return { + status: 'accepted', + task, + ...(receipt.controlId == null ? {} : { controlId: receipt.controlId }), + }; + } + if (receipt.status === 'applied') { + return receipt.action === 'cancel' + ? { status: 'cancelled', task } + : { + status: 'accepted', + task, + ...(receipt.controlId == null ? {} : { controlId: receipt.controlId }), + }; + } + if ( + receipt.reason === 'task_not_running' || + receipt.reason === 'task_completed' || + receipt.reason === 'task_cancelled' || + receipt.reason === 'task_failed' + ) { + return { status: 'not_running', task }; + } + if (receipt.reason === 'control_not_found' || receipt.reason === 'withdrawn') { + return { status: 'control_not_found', task }; + } + return { + status: 'invalid', + message: + receipt.status === 'failed' + ? 'The prior control invocation failed.' + : 'The prior control invocation was rejected.', + }; } private durableReceipt( @@ -627,7 +759,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { updatedAt: new Date(receipt.updatedAt), ...(receipt.boundary == null ? {} : { boundary: receipt.boundary }), ...(receipt.reason == null ? {} : { reason: receipt.reason }), - ...boundedControlMessage(invocation.command), + ...boundedControlMessage(invocation.command, invocation.commandMessageTruncated), }; } @@ -655,7 +787,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { status: 'accepted', createdAt: new Date(invocation.createdAt), updatedAt: new Date(invocation.createdAt), - ...boundedControlMessage(command), + ...boundedControlMessage(command, invocation.commandMessageTruncated), }; } const now = new Date(); @@ -681,7 +813,104 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { createdAt: new Date(invocation.createdAt), updatedAt: now, ...(reason == null ? {} : { reason }), - ...boundedControlMessage(command), + ...boundedControlMessage(command, invocation.commandMessageTruncated), + }; + } + + private async replayDurableControl( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise { + const scope = parseScope(scopeId); + const replay = await this.runWithOwnerContext(scope, () => + this.methods.getSubagentTaskControlReplay({ + userId: scope.userId, + parentConversationId: scope.parentConversationId, + taskId, + invocationId, + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + }), + ); + if (replay == null) return undefined; + if (replay.receipt.fingerprint !== controlFingerprint(command)) { + return { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + } + const { receipt, task: durableTask } = replay; + if (receipt.status === 'reserved') { + /** The prior owner fenced this invocation but did not durably prove the + * side effect. Reapplying could duplicate it; reporting acceptance would lie. */ + throw new SubagentTaskOwnerUnavailableError(); + } + const task: SubagentTaskSnapshot = { + taskId, + threadId: durableTask.threadId, + subagentType: durableTask.subagentType, + status: durableTask.status, + createdAt: durableTask.createdAt.getTime(), + updatedAt: durableTask.updatedAt.getTime(), + resultAvailable: durableTask.resultAvailable, + resultClaimed: durableTask.resultClaimed, + pendingControls: durableTask.pendingControls, + ...(receipt.controlId != null && + (receipt.action === 'steer' || receipt.action === 'queue' || receipt.action === 'interrupt') + ? { + controlReceipts: [ + { + controlId: receipt.controlId, + action: receipt.action, + status: receipt.status, + createdAt: receipt.createdAt.getTime(), + updatedAt: receipt.updatedAt.getTime(), + ...(receipt.boundary == null ? {} : { boundary: receipt.boundary }), + ...(receipt.reason === 'withdrawn' || + receipt.reason === 'task_completed' || + receipt.reason === 'task_cancelled' || + receipt.reason === 'task_failed' + ? { reason: receipt.reason } + : {}), + }, + ], + } + : {}), + }; + if (receipt.status === 'accepted') { + return { + status: 'accepted', + task, + ...(receipt.controlId == null ? {} : { controlId: receipt.controlId }), + }; + } + if (receipt.status === 'applied') { + return command.action === 'cancel' + ? { status: 'cancelled', task } + : { + status: 'accepted', + task, + ...(receipt.controlId == null ? {} : { controlId: receipt.controlId }), + }; + } + if ( + receipt.reason === 'task_not_running' || + receipt.reason === 'task_completed' || + receipt.reason === 'task_cancelled' || + receipt.reason === 'task_failed' + ) { + return { status: 'not_running', task }; + } + if (receipt.reason === 'control_not_found' || receipt.reason === 'withdrawn') { + return { status: 'control_not_found', task }; + } + return { + status: 'invalid', + message: + receipt.status === 'failed' + ? 'The prior control invocation failed.' + : 'The prior control invocation was rejected.', }; } @@ -719,9 +948,18 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { receipt: candidate.receipt, }), ); + if (persisted === 'conflict') { + if (pending.get(invocationId) === candidate) pending.delete(invocationId); + if (pending.size === 0) this.pendingControlReceipts.delete(key); + throw new SubagentControlReceiptConflictError(); + } if (!persisted) { throw new Error('The child control receipt target is not ready.'); } + const invocation = this.retainedControlInvocation(scopeId, taskId, invocationId); + if (invocation?.receipt === candidate.receipt) { + invocation.receiptPersisted = true; + } if (pending.get(invocationId) === candidate) { pending.delete(invocationId); } @@ -753,13 +991,10 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { /** A terminal child may have no later caller to retrigger persistence. Keep a * single bounded retry timer per task so transient storage failures converge - * while this process still owns the task; restart durability remains AI-1737. */ + * even after result collection expires its in-memory task; restart durability + * remains AI-1737. */ private scheduleControlReceiptRetry(scopeId: string, taskId: string): void { const key = controlTaskKey(scopeId, taskId); - if (this.get(scopeId, taskId) == null) { - this.pendingControlReceipts.delete(key); - return; - } if ( this.controlPersistenceStopping || this.controlPersistenceRetryTimers.has(key) || @@ -794,25 +1029,105 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { control: (scopeId, taskId, command, invocationId) => this.controlInvocationAndPersist(scopeId, taskId, command, invocationId), list: (scopeId) => super.list(scopeId), - cancelScope: (scopeId, threadIds) => this.cancelForScope(scopeId, threadIds), + cancelScope: (scopeId, threadIds, removedConversationIds = []) => { + const cancelled = this.cancelForScope(scopeId, threadIds); + if (removedConversationIds.length > 0) { + const scope = parseScope(scopeId); + this.dropDeletedControlReceiptWork( + scope.userId, + new Set(removedConversationIds), + scope.tenantId, + ); + } + return cancelled; + }, + retainsTaskOwnership: (scopeId, taskId) => + this.pendingControlReceipts.has(controlTaskKey(scopeId, taskId)), }); this.taskControlTransport = transport; } async destroyTaskControlTransport(): Promise { + /** Close command admission and synchronously cancel every locally-owned child + * before the first await. The SDK emits all pending-control transitions while + * cancelling, so no receipt producer can race the final persistence snapshot. */ + this.controlCommandAdmissionClosed = true; + const cancellationFlushes: Promise[] = []; + for (const lease of this.activeThreads.values()) { + if (lease.taskId !== '' && this.get(lease.scopeId, lease.taskId)?.status === 'running') { + const cancellation = super.control(lease.scopeId, lease.taskId, { action: 'cancel' }); + if (cancellation.status === 'cancelled') { + /** The SDK hook above is synchronous, but retain direct promises for the + * authoritative terminal snapshot as well. This makes shutdown await the + * transition even when its first storage attempt fails under load. */ + const snapshot = cancellation.task as SnapshotWithControlReceipts; + for (const receipt of snapshot.controlReceipts ?? []) { + const persistence = this.queueAuthoritativeControlReceipt( + lease.scopeId, + lease.taskId, + receipt, + ); + if (persistence != null) cancellationFlushes.push(persistence); + } + } + } + } + const childSettlements = [...this.activeThreads.values()] + .map((lease) => lease.execution) + .filter((execution): execution is Promise => execution != null); + let childSettlementTimedOut = false; + if (childSettlements.length > 0) { + let timeout: ReturnType | undefined; + try { + await Promise.race([ + Promise.allSettled(childSettlements), + new Promise((_, reject) => { + timeout = setTimeout( + () => reject(new SubagentTaskOwnerUnavailableError()), + this.ownerDrainTimeoutMs, + ); + }), + ]); + } catch { + childSettlementTimedOut = true; + } finally { + if (timeout != null) clearTimeout(timeout); + } + } this.controlPersistenceStopping = true; for (const timer of this.controlPersistenceRetryTimers.values()) clearTimeout(timer); this.controlPersistenceRetryTimers.clear(); - const pendingTasks = [...this.pendingControlReceipts.keys()] - .map(parseControlTaskKey) - .filter((task): task is { scopeId: string; taskId: string } => task != null); - await Promise.allSettled( - pendingTasks.map(({ scopeId, taskId }) => this.flushControlReceipts(scopeId, taskId)), - ); - await Promise.allSettled(this.controlPersistenceTails.values()); + await Promise.allSettled(cancellationFlushes); + /** Cancellation can enqueue its terminal transition behind an already-failing + * acceptance write. Re-snapshot both maps after each round so work admitted + * synchronously before shutdown cannot appear just after the final snapshot. */ + for (let attempt = 0; attempt < SHUTDOWN_CONTROL_RECEIPT_FLUSH_ATTEMPTS; attempt += 1) { + const pendingTasks = [...this.pendingControlReceipts.keys()] + .map(parseControlTaskKey) + .filter((task): task is { scopeId: string; taskId: string } => task != null); + await Promise.allSettled( + pendingTasks.map(({ scopeId, taskId }) => this.flushControlReceipts(scopeId, taskId)), + ); + await Promise.allSettled([...this.controlPersistenceTails.values()]); + if (this.pendingControlReceipts.size === 0 && this.controlPersistenceTails.size === 0) { + break; + } + if (attempt + 1 < SHUTDOWN_CONTROL_RECEIPT_FLUSH_ATTEMPTS) { + await new Promise((resolve) => { + setTimeout(resolve, this.shutdownControlReceiptBackoffMs * 2 ** attempt); + }); + } + } const transport = this.taskControlTransport; this.taskControlTransport = undefined; await transport?.destroy(); + if ( + childSettlementTimedOut || + this.pendingControlReceipts.size > 0 || + this.controlPersistenceTails.size > 0 + ) { + throw new SubagentTaskOwnerUnavailableError(); + } } /** Replaces the process-local activity bus after the host's Redis service is ready. */ @@ -934,6 +1249,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { } const lease: TaskThreadLease = active ?? { + scopeId: request.scopeId, idempotencyKey, taskId: '', running: false, @@ -949,8 +1265,8 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { started = super.start({ ...request, threadId, - run: (runtime: SubagentTaskRuntime) => - this.runWithOwnerContext(scope, async () => { + run: (runtime: SubagentTaskRuntime) => { + const execution = this.runWithOwnerContext(scope, async () => { lease.taskId = runtime.taskId; lease.running = true; const detachedUsage: UsageMetadata[] = []; @@ -1102,7 +1418,14 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { this.activeThreads.delete(lockKey); } } - }), + }); + const settlement = execution.then( + () => undefined, + () => undefined, + ); + lease.execution = settlement; + return execution; + }, }); } catch (error) { if (ownsLease && this.activeThreads.get(lockKey) === lease) { @@ -1314,15 +1637,157 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { command: SubagentTaskControlCommand, invocationId: string = randomUUID(), ): Promise { - const local = await this.controlInvocationAndPersist(scopeId, taskId, command, invocationId); + /** A requester that has neither the task nor a retained invocation cannot be + * authoritative. Route first so a remote control pays only the owner's durable + * preflight instead of repeating the same Mongo read on both replicas. */ + const hasLocalAuthority = + this.get(scopeId, taskId) != null || + this.retainedControlInvocation(scopeId, taskId, invocationId) != null; + const local = hasLocalAuthority + ? await this.controlInvocationAndPersist(scopeId, taskId, command, invocationId) + : ({ status: 'not_found' } as const); if (local.status !== 'not_found') { return local; } + let routed: SubagentTaskControlResult | undefined; + try { + routed = await this.taskControlTransport?.control(scopeId, taskId, command, invocationId); + } catch (error) { + if (error instanceof SubagentTaskOwnerUnavailableError) { + const replay = await this.replayDurableControlAtBoundary( + scopeId, + taskId, + command, + invocationId, + ); + if (replay != null) return replay; + } + throw error; + } + if (routed != null && routed.status !== 'not_found') return routed; return ( - (await this.taskControlTransport?.control(scopeId, taskId, command, invocationId)) ?? local + (await this.replayDurableControlAtBoundary(scopeId, taskId, command, invocationId)) ?? + routed ?? + local ); } + /** Durable receipt reads are part of the owner boundary. Storage ambiguity must + * remain retryable instead of escaping as an unrelated tool execution failure. */ + private async replayDurableControlAtBoundary( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise { + try { + return await this.replayDurableControl(scopeId, taskId, command, invocationId); + } catch (error) { + if (error instanceof SubagentTaskOwnerUnavailableError) throw error; + throw new SubagentTaskOwnerUnavailableError(); + } + } + + private retainedControlInvocation( + scopeId: string, + taskId: string, + invocationId: string, + ): ControlInvocationRecord | undefined { + const key = `${scopeId}\u0000${taskId}\u0000${invocationId}`; + return this.controlInvocations.get(key) ?? this.terminalControlInvocations.get(key); + } + + private async replayRetainedControl( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): Promise { + const retained = this.retainedControlInvocation(scopeId, taskId, invocationId); + if (retained == null) return undefined; + if (retained.fingerprint !== controlFingerprint(command)) { + return { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + } + if (hasDurableControlReceipt(retained)) { + /** Terminal materialization and one-shot collection can change after the + * receipt becomes durable. Refresh those flags from the exact durable row so + * same-owner replay agrees with replay after owner loss. */ + if ('task' in retained.result && retained.result.task.status !== 'running') { + const current = this.get(scopeId, taskId); + if (current != null) { + retained.result = { + ...retained.result, + task: { + ...retained.result.task, + status: current.status, + updatedAt: current.updatedAt, + resultAvailable: current.resultAvailable, + resultClaimed: current.resultClaimed, + }, + }; + } else { + const durable = await this.replayDurableControlAtBoundary( + scopeId, + taskId, + command, + invocationId, + ); + if (durable != null) retained.result = durable; + } + } + return retained.result; + } + let persistenceFailed = false; + try { + await (retained.receiptPersistence ?? this.flushControlReceipts(scopeId, taskId)); + } catch { + persistenceFailed = true; + // The durable replay below distinguishes a committed result or conflict + // from a genuinely retryable storage failure. + } + if (hasDurableControlReceipt(retained)) return retained.result; + if (persistenceFailed) { + try { + const retry = this.flushControlReceipts(scopeId, taskId); + retained.receiptPersistence = retry; + await retry; + } catch { + // Durable replay below remains the authoritative discriminator. + } + } + try { + const durable = await this.replayDurableControl(scopeId, taskId, command, invocationId); + if (durable != null) { + retained.result = durable; + retained.receiptPersisted = true; + retained.receiptPersistence = undefined; + return durable; + } + } catch { + // Normalize storage outages at the owner boundary. + } + throw new SubagentTaskOwnerUnavailableError(); + } + + private retainedControlResult( + scopeId: string, + taskId: string, + command: SubagentTaskControlCommand, + invocationId: string, + ): SubagentTaskControlResult | undefined { + const retained = this.retainedControlInvocation(scopeId, taskId, invocationId); + if (retained == null) return undefined; + return retained.fingerprint === controlFingerprint(command) + ? retained.result + : { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + } + /** * Applies one logical control exactly once for its owning task. Idempotency lives * here rather than in the transport so a local and a routed caller of the same @@ -1335,25 +1800,49 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { command: SubagentTaskControlCommand, invocationId: string, ): SubagentTaskControlResult { + if (this.controlCommandAdmissionClosed) { + return { status: 'invalid', message: 'Subagent task controls are shutting down.' }; + } const key = `${scopeId}\u0000${taskId}\u0000${invocationId}`; const fingerprint = controlFingerprint(command); - const applied = this.controlInvocations.get(key); - if (applied != null) { - /** One invocation is one command; reusing its id for different content is a - * caller error rather than a retry, so it is refused instead of applied. */ - return applied.fingerprint === fingerprint - ? applied.result - : { - status: 'invalid', - message: 'This control invocation id was already used for a different command.', - }; - } - if (this.get(scopeId, taskId) == null) { + const applied = this.retainedControlResult(scopeId, taskId, command, invocationId); + if (applied != null) return applied; + const localTask = this.get(scopeId, taskId); + if (localTask == null) { /** Not this replica's task. Refusing here would keep the command from ever * reaching its owner, so local load cannot veto a remote cancellation: the * owner applies its own window to the routed request. */ return this.control(scopeId, taskId, command); } + if (localTask.status !== 'running') { + const result = this.control(scopeId, taskId, command); + if (result.status === 'not_found' || result.status === 'invalid') return result; + if (this.terminalControlInvocations.size >= MAX_TERMINAL_CONTROL_INVOCATIONS) { + const oldestPersisted = [...this.terminalControlInvocations].find( + ([, invocation]) => invocation.receiptPersisted === true, + )?.[0]; + if (oldestPersisted == null) { + return { + status: 'invalid', + message: + 'Too many terminal control invocations are awaiting persistence; retry shortly.', + }; + } + this.terminalControlInvocations.delete(oldestPersisted); + } + this.terminalControlInvocations.set(key, { + scopeId, + taskId, + invocationId, + fingerprint, + command: boundedControlCommand(command), + commandMessageTruncated: + 'message' in command && command.message.length > MAX_DURABLE_CONTROL_MESSAGE_CHARS, + result, + createdAt: Date.now(), + }); + return result; + } if (!this.makeRoomForInvocation()) { /** Every tracked invocation belongs to a task this store still holds. Applying * this command without room to record it would let a caller retry apply it a @@ -1365,7 +1854,10 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { }; } const result = this.control(scopeId, taskId, command); - if (result.status === 'not_found') { + if (result.status === 'not_found' || result.status === 'invalid') { + return result; + } + if (result.status !== 'accepted' && result.status !== 'cancelled') { return result; } const invocation: ControlInvocationRecord = { @@ -1374,6 +1866,8 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { invocationId, fingerprint, command: boundedControlCommand(command), + commandMessageTruncated: + 'message' in command && command.message.length > MAX_DURABLE_CONTROL_MESSAGE_CHARS, result, createdAt: Date.now(), }; @@ -1397,22 +1891,157 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { command: SubagentTaskControlCommand, invocationId: string, ): Promise { + const retained = await this.replayRetainedControl(scopeId, taskId, command, invocationId); + if (retained != null) return retained; + try { + const durable = await this.replayDurableControl(scopeId, taskId, command, invocationId); + if (durable != null) return durable; + } catch (error) { + logger.warn('[subagentThreads] Failed to preflight a child control receipt', error); + throw new SubagentTaskOwnerUnavailableError(); + } + const localTask = this.get(scopeId, taskId); + if (localTask?.status === 'running') { + if (localTask.threadId == null || localTask.threadId === '') { + throw new SubagentTaskOwnerUnavailableError(); + } + const scope = parseScope(scopeId); + const now = new Date(); + const reservation: ISubagentTaskControlReceipt = { + invocationId, + fingerprint: controlFingerprint(command), + action: command.action, + status: 'reserved', + createdAt: now, + updatedAt: now, + ...boundedControlMessage( + boundedControlCommand(command), + 'message' in command && command.message.length > MAX_DURABLE_CONTROL_MESSAGE_CHARS, + ), + }; + let reserved: boolean | 'unchanged' | 'conflict'; + try { + reserved = (await this.controlReservationSlot(() => + this.runWithOwnerContext(scope, () => + this.methods.recordSubagentTaskControlReceipt({ + userId: scope.userId, + conversationId: localTask.threadId as string, + taskId, + ...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }), + receipt: reservation, + }), + ), + )) as boolean | 'unchanged' | 'conflict'; + } catch (error) { + logger.warn('[subagentThreads] Failed to reserve a child control invocation', error); + throw new SubagentTaskOwnerUnavailableError(); + } + if (reserved === 'conflict') { + return { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + } + if (reserved === 'unchanged') { + try { + const replay = await this.replayDurableControl(scopeId, taskId, command, invocationId); + if (replay != null) return replay; + } catch { + // Normalize storage ambiguity at the owner boundary below. + } + throw new SubagentTaskOwnerUnavailableError(); + } + if (!reserved) throw new SubagentTaskOwnerUnavailableError(); + } + const invocationKey = `${scopeId}\u0000${taskId}\u0000${invocationId}`; const result = this.controlInvocation(scopeId, taskId, command, invocationId); - const invocation = this.controlInvocations.get( - `${scopeId}\u0000${taskId}\u0000${invocationId}`, - ); + const retainedInvocation = + this.controlInvocations.get(invocationKey) ?? + this.terminalControlInvocations.get(invocationKey); + const invocation: ControlInvocationRecord | undefined = + retainedInvocation ?? + (result.status === 'not_found' || result.status === 'invalid' + ? undefined + : { + scopeId, + taskId, + invocationId, + fingerprint: controlFingerprint(command), + command: boundedControlCommand(command), + commandMessageTruncated: + 'message' in command && command.message.length > MAX_DURABLE_CONTROL_MESSAGE_CHARS, + result, + createdAt: Date.now(), + }); const threadId = 'task' in result ? result.task.threadId : undefined; if (invocation == null || threadId == null) return result; const receipt = this.controlResultReceipt(invocation); - if (receipt != null) { - try { - await this.queueControlReceipt(scopeId, taskId, threadId, receipt); - } catch (error) { - logger.warn('[subagentThreads] Failed to durably accept a child control', error); - throw new SubagentTaskOwnerUnavailableError(); + let persistedReceipt: ISubagentTaskControlReceipt | undefined; + let persistence: Promise | undefined; + try { + if (receipt != null) { + persistedReceipt = receipt; + invocation.receipt = persistedReceipt; + invocation.receiptPersisted = false; + persistence = this.queueControlReceipt(scopeId, taskId, threadId, persistedReceipt); + invocation.receiptPersistence = persistence; + await persistence; + /** A terminal SDK transition can replace the accepted projection while its + * older write is awaiting Mongo. Mark only the exact generation awaited. */ + if (invocation.receipt === persistedReceipt) invocation.receiptPersisted = true; + /** Do not acknowledge an older generation while a newer authoritative SDK + * transition is still queued. There is no async gap after this loop, so the + * generation proven durable is the one returned to the caller. */ + await this.awaitCurrentControlReceipt(scopeId, taskId, invocation); + } + } catch (error) { + if (error instanceof SubagentControlReceiptConflictError) { + const invalid: SubagentTaskControlResult = { + status: 'invalid', + message: 'This control invocation id was already used for a different command.', + }; + invocation.result = invalid; + invocation.receiptPersisted = true; + if (result.status === 'accepted' && result.controlId != null) { + this.controlInvocationByReceipt.delete( + controlReceiptKey(scopeId, taskId, result.controlId), + ); + super.control(scopeId, taskId, { + action: 'cancel_message', + controlId: result.controlId, + }); + } + return invalid; + } + logger.warn('[subagentThreads] Failed to durably accept a child control', error); + throw new SubagentTaskOwnerUnavailableError(); + } finally { + if ( + persistence != null && + persistedReceipt != null && + invocation.receiptPersistence === persistence && + invocation.receipt === persistedReceipt && + invocation.receiptPersisted === true + ) { + invocation.receiptPersistence = undefined; } } - return result; + return invocation.result; + } + + private async awaitCurrentControlReceipt( + scopeId: string, + taskId: string, + invocation: ControlInvocationRecord, + ): Promise { + while (invocation.receipt != null && invocation.receiptPersisted !== true) { + const receipt = invocation.receipt; + const persistence = + invocation.receiptPersistence ?? this.flushControlReceipts(scopeId, taskId); + invocation.receiptPersistence = persistence; + await persistence; + if (invocation.receipt === receipt && hasDurableControlReceipt(invocation)) return; + } } /** @@ -1435,7 +2064,6 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { controlReceiptKey(invocation.scopeId, invocation.taskId, result.controlId), ); } - this.pendingControlReceipts.delete(controlTaskKey(invocation.scopeId, invocation.taskId)); } } return this.controlInvocations.size < this.maxControlInvocations; @@ -1602,6 +2230,9 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { let cancelled = this.cancelForConversations(userId, targets, tenantId); const transport = this.taskControlTransport; if (transport == null) { + if (removed.size > 0) { + this.dropDeletedControlReceiptWork(userId, removed, tenantId); + } return cancelled; } const cancelSlot = createConcurrencyLimiter(DELETION_CANCEL_CONCURRENCY); @@ -1615,9 +2246,16 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { }), threadIds: null, })); - const scopeCancellations = [...plan.scopes, ...cascadeScopes].map((scope) => - cancelSlot(() => transport.cancelScope(scope.scopeId, scope.threadIds)), - ); + const scopeCancellations = [...plan.scopes, ...cascadeScopes].map((scope) => { + const parsed = parseScope(scope.scopeId); + const removedForScope = [ + ...(removed.has(parsed.parentConversationId) ? [parsed.parentConversationId] : []), + ...(scope.threadIds ?? []).filter((threadId) => removed.has(threadId)), + ]; + return cancelSlot(() => + transport.cancelScope(scope.scopeId, scope.threadIds, removedForScope), + ); + }); const leaseCancellations = plan.leases .filter( (lease) => removed.has(lease.parentConversationId) || removed.has(lease.conversationId), @@ -1629,7 +2267,14 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { parentConversationId: lease.parentConversationId, ...(tenantId ? { tenantId } : {}), }); - const stopped = await transport.cancelScope(scopeId, [lease.conversationId]); + const removedForLease = [lease.parentConversationId, lease.conversationId].filter((id) => + removed.has(id), + ); + const stopped = await transport.cancelScope( + scopeId, + [lease.conversationId], + removedForLease, + ); if (stopped > 0 || this.cancelUnroutedTask == null) return stopped; return (await this.cancelUnroutedTask({ userId, @@ -1641,13 +2286,74 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore { : 0; }), ); - for (const count of await Promise.all(scopeCancellations)) { - cancelled += count; + try { + for (const count of await Promise.all(scopeCancellations)) { + cancelled += count; + } + for (const count of await Promise.all(leaseCancellations)) { + cancelled += count; + } + return cancelled; + } finally { + /** Delivery may fail after the deletion committed. Receipt persistence for + * removed rows is still terminal and must not poison graceful shutdown. */ + if (removed.size > 0) { + this.dropDeletedControlReceiptWork(userId, removed, tenantId); + } } - for (const count of await Promise.all(leaseCancellations)) { - cancelled += count; + } + + /** A successful deletion makes false receipt writes permanent, not retryable. + * Remove only work whose authorized parent or child was actually deleted. */ + private dropDeletedControlReceiptWork( + userId: string, + removedConversationIds: ReadonlySet, + tenantId?: string, + ): void { + const matchesDeletedScope = (scopeId: string): boolean => { + const scope = parseScope(scopeId); + return ( + scope.userId === userId && + matchesTenant(scope.tenantId, tenantId) && + removedConversationIds.has(scope.parentConversationId) + ); + }; + for (const [key, pending] of this.pendingControlReceipts) { + const task = parseControlTaskKey(key); + if (task == null) continue; + const deleteWholeTask = matchesDeletedScope(task.scopeId); + for (const [invocationId, candidate] of pending) { + if (deleteWholeTask || removedConversationIds.has(candidate.threadId)) { + pending.delete(invocationId); + } + } + if (pending.size === 0) { + this.pendingControlReceipts.delete(key); + const retry = this.controlPersistenceRetryTimers.get(key); + if (retry != null) clearTimeout(retry); + this.controlPersistenceRetryTimers.delete(key); + } } - return cancelled; + const dropInvocation = (key: string, invocation: ControlInvocationRecord): void => { + const resultThreadId = + 'task' in invocation.result ? invocation.result.task.threadId : undefined; + if ( + !matchesDeletedScope(invocation.scopeId) && + (resultThreadId == null || !removedConversationIds.has(resultThreadId)) + ) { + return; + } + this.controlInvocations.delete(key); + this.terminalControlInvocations.delete(key); + if (invocation.result.status === 'accepted' && invocation.result.controlId != null) { + this.controlInvocationByReceipt.delete( + controlReceiptKey(invocation.scopeId, invocation.taskId, invocation.result.controlId), + ); + } + }; + for (const [key, invocation] of this.controlInvocations) dropInvocation(key, invocation); + for (const [key, invocation] of this.terminalControlInvocations) + dropInvocation(key, invocation); } /** Cancels this process's live children for one scope, optionally narrowed to threads. */ @@ -2639,6 +3345,7 @@ const REQUIRED_THREAD_METHODS = [ 'deleteConvos', 'deleteMessages', 'getConvo', + 'getSubagentTaskControlReplay', 'getMessages', 'listActiveSubagentThreadLeases', 'recordSubagentTaskControlReceipt', @@ -2666,6 +3373,7 @@ export function createSubagentThreadTaskStore( MessageMethods, | 'claimSubagentTaskResult' | 'deleteMessages' + | 'getSubagentTaskControlReplay' | 'getMessages' | 'recordSubagentTaskControlReceipt' | 'saveMessage' diff --git a/packages/api/src/agents/view.spec.ts b/packages/api/src/agents/view.spec.ts index a3527d6813..1e99c99c70 100644 --- a/packages/api/src/agents/view.spec.ts +++ b/packages/api/src/agents/view.spec.ts @@ -207,6 +207,14 @@ describe('subagent thread parent-scoped view', () => { it('returns bounded authoritative control receipts without private fingerprints', async () => { const input = message('task-1:user', 'running', true); input.subagentTask!.controlReceipts = [ + { + invocationId: 'private-reservation', + fingerprint: 'private-reservation-fingerprint', + action: 'queue' as const, + status: 'reserved' as const, + createdAt: new Date('2026-08-21T09:59:59.000Z'), + updatedAt: new Date('2026-08-21T09:59:59.000Z'), + }, ...Array.from({ length: 32 }, (_, index) => ({ invocationId: `earlier-${index}`, fingerprint: `private-${index}`, @@ -258,6 +266,7 @@ describe('subagent thread parent-scoped view', () => { expect(Buffer.byteLength(projected?.message ?? '', 'utf8')).toBeLessThanOrEqual(512); expect(view.controlReceipts).toHaveLength(32); expect(view.controlReceiptsTruncated).toBe(true); + expect(JSON.stringify(view)).not.toContain('private-reservation'); expect(JSON.stringify(view)).not.toContain('private-fingerprint'); }); diff --git a/packages/api/src/agents/view.ts b/packages/api/src/agents/view.ts index 36859a51c6..b52ab0f35c 100644 --- a/packages/api/src/agents/view.ts +++ b/packages/api/src/agents/view.ts @@ -125,8 +125,17 @@ const publicControlReceipts = ( ): { receipts: SubagentControlReceipt[]; truncated: boolean } => { const input = messages.find((message) => message.messageId === `${taskId}:user`); const stored = input?.subagentTask?.controlReceipts ?? []; - const accepted = stored.filter((receipt) => receipt.status === 'accepted'); - const terminal = stored.filter((receipt) => receipt.status !== 'accepted'); + /** A reservation only fences at-most-once application; it does not claim that + * guidance was accepted and must never appear in the public activity view. */ + const visible = stored.filter( + ( + receipt, + ): receipt is typeof receipt & { + status: 'accepted' | 'applied' | 'rejected' | 'failed'; + } => receipt.status !== 'reserved', + ); + const accepted = visible.filter((receipt) => receipt.status === 'accepted'); + const terminal = visible.filter((receipt) => receipt.status !== 'accepted'); const terminalLimit = Math.max(0, MAX_PUBLIC_CONTROL_RECEIPTS - accepted.length); const retained = [...accepted, ...(terminalLimit === 0 ? [] : terminal.slice(-terminalLimit))] .slice(0, MAX_PUBLIC_CONTROL_RECEIPTS) @@ -154,7 +163,7 @@ const publicControlReceipts = ( : {}), }; }); - return { receipts: retained, truncated: retained.length < stored.length }; + return { receipts: retained, truncated: retained.length < visible.length }; }; const publicStatus = ( diff --git a/packages/api/src/app/shutdown.spec.ts b/packages/api/src/app/shutdown.spec.ts index 9e0e76f199..3b1a4631e9 100644 --- a/packages/api/src/app/shutdown.spec.ts +++ b/packages/api/src/app/shutdown.spec.ts @@ -301,7 +301,7 @@ describe('setupGracefulShutdown', () => { expect(order).toEqual(['generation streams', 'default-first', 'default-second', 'telemetry']); }); - it('continues subsequent tasks and still exits if one task throws', async () => { + it('continues subsequent tasks and exits nonzero if one task throws', async () => { const calls: string[] = []; jest.spyOn(server, 'close').mockImplementation((cb?: (err?: Error) => void) => { if (cb) { @@ -324,7 +324,7 @@ describe('setupGracefulShutdown', () => { await flush(); await flush(); expect(calls).toEqual(['ok-before', 'throws', 'ok-after']); - expect(exitSpy).toHaveBeenCalledWith(0); + expect(exitSpy).toHaveBeenCalledWith(1); }); it('awaits async tasks before exiting', async () => { diff --git a/packages/api/src/app/shutdown.ts b/packages/api/src/app/shutdown.ts index 656936af3f..ef10189384 100644 --- a/packages/api/src/app/shutdown.ts +++ b/packages/api/src/app/shutdown.ts @@ -81,7 +81,7 @@ export function __resetShutdownStateForTests(): void { clearForceExitTimer(); } -async function runShutdownTasks(phase: ShutdownPhase): Promise { +async function runShutdownTasks(phase: ShutdownPhase): Promise { const orderedTasks = tasks .filter((task) => task.phase === phase) .sort( @@ -89,14 +89,17 @@ async function runShutdownTasks(phase: ShutdownPhase): Promise { right.priority - left.priority || left.registrationOrder - right.registrationOrder, ); + let failed = false; for (const task of orderedTasks) { try { logger.info(`Running ${phase} shutdown task: ${task.name}`); await task.fn(); } catch (err) { + failed = true; logger.error(`Shutdown task "${task.name}" failed:`, err); } } + return failed; } function clearForceExitTimer(): void { @@ -130,9 +133,9 @@ async function shutdown(signal: NodeJS.Signals): Promise { exitCode = 1; }); - await runShutdownTasks('pre-drain'); + if (await runShutdownTasks('pre-drain')) exitCode = 1; await serverClosePromise; - await runShutdownTasks('post-drain'); + if (await runShutdownTasks('post-drain')) exitCode = 1; } finally { clearTimeout(forceExit); if (forceExitTimer === forceExit) { diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index e26f9fea16..57e7940cf6 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -42,6 +42,12 @@ let claimSubagentTaskResult: ReturnType['claimSubag let recordSubagentTaskControlReceipt: ReturnType< typeof createMessageMethods >['recordSubagentTaskControlReceipt']; +let getSubagentTaskControlReceipt: ReturnType< + typeof createMessageMethods +>['getSubagentTaskControlReceipt']; +let getSubagentTaskControlReplay: ReturnType< + typeof createMessageMethods +>['getSubagentTaskControlReplay']; let releaseSubagentTaskResultClaim: ReturnType< typeof createMessageMethods >['releaseSubagentTaskResultClaim']; @@ -68,6 +74,8 @@ beforeAll(async () => { recordMessage = methods.recordMessage; claimSubagentTaskResult = methods.claimSubagentTaskResult; recordSubagentTaskControlReceipt = methods.recordSubagentTaskControlReceipt; + getSubagentTaskControlReceipt = methods.getSubagentTaskControlReceipt; + getSubagentTaskControlReplay = methods.getSubagentTaskControlReplay; releaseSubagentTaskResultClaim = methods.releaseSubagentTaskResultClaim; await mongoose.connect(mongoUri); @@ -97,6 +105,7 @@ describe('Message Operations', () => { // Clear database await Message.deleteMany({}); + await mongoose.models.Conversation.deleteMany({}); mockCtx = { userId: 'user123', @@ -2334,6 +2343,25 @@ describe('Message Operations', () => { message: 'Use the primary source.', }; + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { ...accepted, controlId: undefined, status: 'reserved' }, + }), + ).resolves.toBe(true); + /** A reservation fences competing owners but is not public evidence that + * the command was accepted or applied. */ + await expect( + getSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + invocationId: 'invocation-1', + }), + ).resolves.toBeNull(); + await expect( recordSubagentTaskControlReceipt({ userId: 'user123', @@ -2356,12 +2384,14 @@ describe('Message Operations', () => { }), ).resolves.toBe(true); /** A delayed accepted replay cannot downgrade the durable terminal receipt. */ - await recordSubagentTaskControlReceipt({ - userId: 'user123', - conversationId, - taskId: 'task-1', - receipt: accepted, - }); + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: accepted, + }), + ).resolves.toBe('unchanged'); await expect( recordSubagentTaskControlReceipt({ userId: 'another-user', @@ -2448,6 +2478,317 @@ describe('Message Operations', () => { ]); }); + it('reads an exact authorized receipt and persists a new terminal rejection', async () => { + const conversationId = uuidv4(); + await createTaskInput(conversationId); + await mongoose.models.Conversation.create({ + user: 'user123', + conversationId, + title: 'Child thread', + endpoint: 'agents', + subagentThread: { + rootConversationId: 'parent-conversation', + parentConversationId: 'parent-conversation', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + }); + await Message.create({ + user: 'user123', + conversationId, + messageId: 'task-1:assistant', + parentMessageId: 'task-1:user', + sender: 'researcher', + text: 'Done.', + endpoint: 'agents', + isCreatedByUser: false, + subagentTask: { + attemptKey: 'task-1-attempt', + parentRunId: 'parent-message', + status: 'completed', + resultClaim: { + kind: 'manual', + claimId: 'poll-1', + claimedAt: new Date('2026-08-24T12:00:02.000Z'), + }, + }, + }); + const now = new Date('2026-08-24T12:00:00.000Z'); + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId: 'terminal-invocation', + fingerprint: 'terminal-fingerprint', + action: 'cancel', + status: 'rejected', + reason: 'task_not_running', + createdAt: now, + updatedAt: now, + }, + }), + ).resolves.toBe(true); + + await expect( + getSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + invocationId: 'terminal-invocation', + }), + ).resolves.toEqual( + expect.objectContaining({ + invocationId: 'terminal-invocation', + fingerprint: 'terminal-fingerprint', + status: 'rejected', + reason: 'task_not_running', + }), + ); + await expect( + getSubagentTaskControlReceipt({ + userId: 'another-user', + conversationId, + taskId: 'task-1', + invocationId: 'terminal-invocation', + }), + ).resolves.toBeNull(); + await expect( + getSubagentTaskControlReplay({ + userId: 'user123', + parentConversationId: 'parent-conversation', + taskId: 'task-1', + invocationId: 'terminal-invocation', + }), + ).resolves.toEqual({ + receipt: expect.objectContaining({ invocationId: 'terminal-invocation' }), + task: expect.objectContaining({ + taskId: 'task-1', + threadId: conversationId, + subagentType: 'researcher', + status: 'completed', + resultAvailable: true, + resultClaimed: true, + }), + }); + await expect( + getSubagentTaskControlReplay({ + userId: 'user123', + parentConversationId: 'different-parent', + taskId: 'task-1', + invocationId: 'terminal-invocation', + }), + ).resolves.toBeNull(); + + const ordinaryConversationId = uuidv4(); + await Message.create({ + user: 'user123', + conversationId: ordinaryConversationId, + messageId: 'ordinary-task:user', + parentMessageId: Constants.NO_PARENT, + sender: 'User', + text: 'An ordinary message with a colliding id.', + endpoint: 'agents', + isCreatedByUser: true, + }); + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId: ordinaryConversationId, + taskId: 'ordinary-task', + receipt: { + invocationId: 'terminal-invocation', + fingerprint: 'terminal-fingerprint', + action: 'cancel', + status: 'rejected', + reason: 'task_not_running', + createdAt: now, + updatedAt: now, + }, + }), + ).resolves.toBe(false); + }); + + it('replays an applied cancellation as cancelled before its terminal row exists', async () => { + const conversationId = uuidv4(); + await createTaskInput(conversationId); + await mongoose.models.Conversation.create({ + user: 'user123', + conversationId, + title: 'Cancelling child thread', + endpoint: 'agents', + subagentThread: { + rootConversationId: 'parent-conversation', + parentConversationId: 'parent-conversation', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + }); + const now = new Date('2026-08-24T12:00:00.000Z'); + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId: 'cancel-invocation', + fingerprint: 'cancel-fingerprint', + action: 'cancel', + status: 'applied', + createdAt: now, + updatedAt: now, + }, + }), + ).resolves.toBe(true); + + await expect( + getSubagentTaskControlReplay({ + userId: 'user123', + parentConversationId: 'parent-conversation', + taskId: 'task-1', + invocationId: 'cancel-invocation', + }), + ).resolves.toEqual({ + receipt: expect.objectContaining({ invocationId: 'cancel-invocation' }), + task: expect.objectContaining({ + taskId: 'task-1', + threadId: conversationId, + status: 'cancelled', + resultAvailable: false, + resultClaimed: false, + updatedAt: now, + }), + }); + }); + + it('reports every accepted control when replaying one durable invocation', async () => { + const conversationId = uuidv4(); + await createTaskInput(conversationId); + await mongoose.models.Conversation.create({ + user: 'user123', + conversationId, + title: 'Controlled child thread', + endpoint: 'agents', + subagentThread: { + rootConversationId: 'parent-conversation', + parentConversationId: 'parent-conversation', + parentMessageId: 'parent-message', + parentToolCallId: 'parent-tool', + subagentType: 'researcher', + subagentKind: 'agent', + depth: 1, + }, + }); + const now = new Date('2026-08-24T12:00:00.000Z'); + for (const index of [1, 2]) { + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId: `pending-invocation-${index}`, + fingerprint: `pending-fingerprint-${index}`, + controlId: `pending-control-${index}`, + action: 'queue', + status: 'accepted', + createdAt: now, + updatedAt: now, + }, + }), + ).resolves.toBe(true); + } + + await expect( + getSubagentTaskControlReplay({ + userId: 'user123', + parentConversationId: 'parent-conversation', + taskId: 'task-1', + invocationId: 'pending-invocation-1', + }), + ).resolves.toEqual({ + receipt: expect.objectContaining({ invocationId: 'pending-invocation-1' }), + task: expect.objectContaining({ pendingControls: 2 }), + }); + }); + + it('rejects an invocation fingerprint conflict without reporting persistence', async () => { + const conversationId = uuidv4(); + await createTaskInput(conversationId); + const now = new Date('2026-08-24T12:00:00.000Z'); + const receipt = { + invocationId: 'conflicting-invocation', + fingerprint: 'first-fingerprint', + controlId: 'first-control', + action: 'queue' as const, + status: 'accepted' as const, + createdAt: now, + updatedAt: now, + }; + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt, + }), + ).resolves.toBe(true); + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { ...receipt, fingerprint: 'different-fingerprint' }, + }), + ).resolves.toBe('conflict'); + }); + + it('retains concurrent receipts without requiring an aggregation-pipeline update', async () => { + const conversationId = uuidv4(); + await createTaskInput(conversationId); + const createdAt = new Date('2026-08-24T12:00:00.000Z'); + + const invocationIds = Array.from({ length: 32 }, (_, index) => `concurrent-${index}`); + await Promise.all( + invocationIds.map((invocationId, index) => + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId, + fingerprint: `${invocationId}-fingerprint`, + controlId: `${invocationId}-control`, + action: 'queue', + status: 'accepted', + createdAt: new Date(createdAt.getTime() + index), + updatedAt: new Date(createdAt.getTime() + index), + }, + }), + ), + ); + + const stored = await Message.findOne({ + user: 'user123', + conversationId, + messageId: 'task-1:user', + }) + .select('+subagentTask') + .lean(); + expect(stored?.subagentTask?.controlReceipts).toEqual( + expect.arrayContaining( + invocationIds.map((invocationId) => expect.objectContaining({ invocationId })), + ), + ); + }); + it('retains accepted commands while bounding terminal receipt history', async () => { const conversationId = uuidv4(); await createTaskInput(conversationId); @@ -2537,26 +2878,31 @@ describe('Message Operations', () => { ); }); - it('defensively caps accepted receipts outside the supported task-store path', async () => { + it('refuses to evict active receipt fences at durable capacity', async () => { const conversationId = uuidv4(); await createTaskInput(conversationId); const createdAt = new Date('2026-08-24T12:00:00.000Z'); + const results: Array = []; for (let index = 0; index < 70; index += 1) { - await recordSubagentTaskControlReceipt({ - userId: 'user123', - conversationId, - taskId: 'task-1', - receipt: { - invocationId: `accepted-${index}`, - fingerprint: `fingerprint-${index}`, - controlId: `control-${index}`, - action: 'queue', - status: 'accepted', - createdAt: new Date(createdAt.getTime() + index), - updatedAt: new Date(createdAt.getTime() + index), - }, - }); + results.push( + await recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId: `accepted-${index}`, + fingerprint: `fingerprint-${index}`, + controlId: `control-${index}`, + action: 'queue', + status: 'accepted', + createdAt: new Date(createdAt.getTime() + index), + updatedAt: new Date(createdAt.getTime() + index), + }, + }), + ); } + expect(results.slice(0, 64)).toEqual(Array.from({ length: 64 }, () => true)); + expect(results.slice(64)).toEqual(Array.from({ length: 6 }, () => false)); const stored = await Message.findOne({ user: 'user123', @@ -2566,7 +2912,73 @@ describe('Message Operations', () => { .select('+subagentTask') .lean(); expect(stored?.subagentTask?.controlReceipts).toHaveLength(64); - expect(stored?.subagentTask?.controlReceipts?.[0]?.invocationId).toBe('accepted-6'); + expect(stored?.subagentTask?.controlReceipts?.[0]?.invocationId).toBe('accepted-0'); + expect(stored?.subagentTask?.controlReceipts?.[63]?.invocationId).toBe('accepted-63'); + + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId: 'terminal-with-no-allowance', + fingerprint: 'terminal-fingerprint', + controlId: 'terminal-control', + action: 'queue', + status: 'applied', + createdAt: new Date(createdAt.getTime() + 100), + updatedAt: new Date(createdAt.getTime() + 100), + boundary: 'turn', + }, + }), + ).resolves.toBe(false); + const afterTerminal = await Message.findOne({ + user: 'user123', + conversationId, + messageId: 'task-1:user', + }) + .select('+subagentTask') + .lean(); + expect(afterTerminal?.subagentTask?.controlReceipts).toHaveLength(64); + expect(afterTerminal?.subagentTask?.controlReceipts).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ invocationId: 'terminal-with-no-allowance' }), + ]), + ); + + /** Completing an existing active fence always frees its own slot and wins + * over terminal history, even though its occurrence timestamp is oldest. */ + await expect( + recordSubagentTaskControlReceipt({ + userId: 'user123', + conversationId, + taskId: 'task-1', + receipt: { + invocationId: 'accepted-0', + fingerprint: 'fingerprint-0', + controlId: 'control-0', + action: 'queue', + status: 'applied', + createdAt, + updatedAt: new Date(createdAt.getTime() + 101), + boundary: 'turn', + }, + }), + ).resolves.toBe(true); + const afterTransition = await Message.findOne({ + user: 'user123', + conversationId, + messageId: 'task-1:user', + }) + .select('+subagentTask') + .lean(); + expect(afterTransition?.subagentTask?.controlReceipts).toHaveLength(64); + expect(afterTransition?.subagentTask?.controlReceipts).toEqual( + expect.arrayContaining([ + expect.objectContaining({ invocationId: 'accepted-0', status: 'applied' }), + expect.objectContaining({ invocationId: 'accepted-63', status: 'accepted' }), + ]), + ); }); }); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 941147f53c..b3bd6e3d05 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -1,9 +1,9 @@ import { HITL_MESSAGE_FILTER_FIELDS, RetentionMode } from 'librechat-data-provider'; import type { UserSubmittedMessageFieldPath } from 'librechat-data-provider'; import type { DeleteResult, FilterQuery, Model, Types } from 'mongoose'; -import type { AppConfig, IMessage } from '~/types'; +import type { AppConfig, IConversation, IMessage } from '~/types'; +import { activeExpirationFilter, createFallbackRetentionDate } from '~/utils/retention'; import { createTempChatExpirationDate } from '~/utils/tempChatRetention'; -import { createFallbackRetentionDate } from '~/utils/retention'; import { tenantSafeBulkWrite } from '~/utils/tenantBulkWrite'; import logger from '~/config/winston'; @@ -15,6 +15,9 @@ const MAX_STORED_USER_SUBMITTED_FIELD_PATHS = MAX_NORMALIZED_USER_SUBMITTED_PATH const MAX_USER_SUBMITTED_PATH_LENGTH = 2048; const MAX_SUBAGENT_CONTROL_RECEIPTS = 64; const MAX_SUBAGENT_CONTROL_MESSAGE_LENGTH = 4 * 1024; +/** One owner admits at most 64 terminal control invocations. The optimistic + * writer therefore has enough rounds for every admitted receipt to converge. */ +const MAX_SUBAGENT_CONTROL_RECEIPT_CAS_ATTEMPTS = 64; const PROVENANCE_PATHS_UNION_FIELD = '__lcProvenancePathsUnion'; const PROVENANCE_FIELD_PATHS_UNION_FIELD = '__lcProvenanceFieldPathsUnion'; const HITL_MESSAGE_FILTER_FIELD_SET = new Set(HITL_MESSAGE_FILTER_FIELDS); @@ -119,6 +122,92 @@ function getStrictPipelineUpdate(Message: Model, update: Record; } +type StoredSubagentControlReceipt = NonNullable< + NonNullable['controlReceipts'] +>[number]; + +const terminalControlReceipt = (receipt: StoredSubagentControlReceipt): boolean => + receipt.status === 'applied' || receipt.status === 'rejected' || receipt.status === 'failed'; + +function retainSubagentControlReceipts( + current: StoredSubagentControlReceipt[], + receipt: StoredSubagentControlReceipt, +): { + status: 'updated' | 'unchanged' | 'conflict' | 'capacity'; + receipts: StoredSubagentControlReceipt[]; +} { + const existingIndex = current.findIndex( + (candidate) => candidate.invocationId === receipt.invocationId, + ); + let merged: StoredSubagentControlReceipt[]; + if (existingIndex < 0) { + merged = [...current, receipt]; + } else { + const existing = current[existingIndex]; + if (existing.fingerprint !== receipt.fingerprint) { + return { status: 'conflict', receipts: current }; + } + if ( + terminalControlReceipt(existing) || + existing.status === receipt.status || + (existing.status === 'accepted' && receipt.status === 'reserved') + ) { + return { status: 'unchanged', receipts: current }; + } + merged = current.map((candidate, index) => (index === existingIndex ? receipt : candidate)); + } + const accepted = merged.filter( + (candidate) => candidate.status === 'reserved' || candidate.status === 'accepted', + ); + /** Reserved and accepted receipts are idempotency fences for commands that can + * still take effect. Never evict one to admit another receipt: report capacity + * so the caller refuses the command before mutating the live task. */ + if (accepted.length > MAX_SUBAGENT_CONTROL_RECEIPTS) { + return { status: 'capacity', receipts: current }; + } + const terminalAllowance = Math.max(0, MAX_SUBAGENT_CONTROL_RECEIPTS - accepted.length); + let terminal = + terminalAllowance === 0 + ? [] + : merged + .filter((candidate) => candidate.status !== 'reserved' && candidate.status !== 'accepted') + .sort( + (left, right) => + left.createdAt.getTime() - right.createdAt.getTime() || + left.invocationId.localeCompare(right.invocationId), + ) + .slice(-terminalAllowance); + const advancesActiveFence = + existingIndex >= 0 && + !terminalControlReceipt(current[existingIndex]) && + terminalControlReceipt(receipt); + if ( + advancesActiveFence && + !terminal.some((candidate) => candidate.invocationId === receipt.invocationId) + ) { + /** A terminal transition for an active fence must outrank unrelated terminal + * history even though it retains the command's older occurrence timestamp. */ + const otherAllowance = Math.max(0, terminalAllowance - 1); + terminal = [ + ...(otherAllowance === 0 + ? [] + : terminal + .filter((candidate) => candidate.invocationId !== receipt.invocationId) + .slice(-otherAllowance)), + receipt, + ].sort( + (left, right) => + left.createdAt.getTime() - right.createdAt.getTime() || + left.invocationId.localeCompare(right.invocationId), + ); + } + const receipts = [...accepted, ...terminal]; + if (!receipts.some((candidate) => candidate.invocationId === receipt.invocationId)) { + return { status: 'capacity', receipts: current }; + } + return { status: 'updated', receipts }; +} + /** * Builds one Mongo aggregation update that merges and caps both provenance * sets. Generic path overflow promotes the message to whole-message user @@ -291,7 +380,34 @@ export interface MessageMethods { taskId: string; tenantId?: string; receipt: NonNullable['controlReceipts']>[number]; - }): Promise; + }): Promise; + getSubagentTaskControlReceipt(input: { + userId: string; + conversationId: string; + taskId: string; + invocationId: string; + tenantId?: string; + }): Promise['controlReceipts']>[number] | null>; + getSubagentTaskControlReplay(input: { + userId: string; + parentConversationId: string; + taskId: string; + invocationId: string; + tenantId?: string; + }): Promise<{ + receipt: NonNullable['controlReceipts']>[number]; + task: { + taskId: string; + threadId: string; + subagentType: string; + status: NonNullable['status']; + resultAvailable: boolean; + resultClaimed: boolean; + pendingControls: number; + createdAt: Date; + updatedAt: Date; + }; + } | null>; bulkSaveMessages( messages: Array>, overrideTimestamp?: boolean, @@ -892,9 +1008,9 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa taskId: string; tenantId?: string; receipt: NonNullable['controlReceipts']>[number]; - }): Promise { + }): Promise { const validActions = new Set(['steer', 'queue', 'interrupt', 'cancel', 'cancel_message']); - const validStatuses = new Set(['accepted', 'applied', 'rejected', 'failed']); + const validStatuses = new Set(['reserved', 'accepted', 'applied', 'rejected', 'failed']); if ( userId.length === 0 || conversationId.length === 0 || @@ -915,233 +1031,211 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa throw new TypeError('Invalid subagent task control receipt'); } const Message = mongoose.models.Message as Model; - const terminalStatuses = ['applied', 'rejected', 'failed']; - const updated = await Message.findOneAndUpdate( - { - user: userId, - conversationId, - ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), - messageId: `${taskId}:user`, - 'subagentTask.status': 'running', + const recordsTerminalRejection = + receipt.status === 'rejected' && receipt.reason === 'task_not_running'; + const identity = { + user: userId, + conversationId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + messageId: `${taskId}:user`, + /** A genuinely new command can arrive after its task settles or its final + * lease expires. Persist that authoritative rejection for retries; every + * command that could still be applied remains fenced to a running task. */ + ...(recordsTerminalRejection + ? { 'subagentTask.status': { $in: ['running', 'completed', 'error', 'cancelled'] } } + : { 'subagentTask.status': 'running' }), + }; + /** Amazon DocumentDB does not support aggregation-pipeline updates. Use a + * bounded optimistic compare-and-swap: the read is small, the write uses + * only plain operators, and concurrent writers retry rather than overwrite. */ + for (let attempt = 0; attempt < MAX_SUBAGENT_CONTROL_RECEIPT_CAS_ATTEMPTS; attempt += 1) { + const currentMessage = await Message.findOne(identity) + .select({ 'subagentTask.controlReceipts': 1, _id: 0 }) + .lean | null>(); + if (currentMessage == null) return false; + const current = currentMessage.subagentTask?.controlReceipts ?? []; + const retained = retainSubagentControlReceipts(current, receipt); + if (retained.status === 'conflict') return 'conflict'; + if (retained.status === 'unchanged') return 'unchanged'; + if (retained.status === 'capacity') return false; + const next = retained.receipts; + const currentFilter = + currentMessage.subagentTask?.controlReceipts == null + ? { 'subagentTask.controlReceipts': { $exists: false } } + : { 'subagentTask.controlReceipts': current }; + const updated = await Message.findOneAndUpdate( + { ...identity, ...currentFilter }, + { $set: { 'subagentTask.controlReceipts': next } }, + { new: false, projection: { messageId: 1 } }, + ).lean<{ messageId: string } | null>(); + if (updated != null) return true; + } + throw new Error('Subagent control receipt write contention exceeded its retry bound.'); + } + + /** Reads one bounded authoritative receipt by its exact durable task identity. + * The stored projection is already capped, and no task/runtime metadata leaves + * this method. Authorization remains part of the Mongo identity. */ + async function getSubagentTaskControlReceipt({ + userId, + conversationId, + taskId, + invocationId, + tenantId, + }: { + userId: string; + conversationId: string; + taskId: string; + invocationId: string; + tenantId?: string; + }): Promise { + if ( + userId.length === 0 || + conversationId.length === 0 || + conversationId.length > 256 || + taskId.length === 0 || + taskId.length > 256 || + invocationId.length === 0 || + invocationId.length > 128 + ) { + throw new TypeError('Invalid subagent task control receipt identity'); + } + const Message = mongoose.models.Message as Model; + const input = await Message.findOne({ + user: userId, + conversationId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + messageId: `${taskId}:user`, + 'subagentTask.controlReceipts.invocationId': invocationId, + }) + .select({ 'subagentTask.controlReceipts': 1, _id: 0 }) + .lean | null>(); + const receipt = input?.subagentTask?.controlReceipts?.find( + (candidate) => candidate.invocationId === invocationId, + ); + /** Reservations are a server-private at-most-once fence, not proof that a + * control was applied. Public HTTP callers retry through the owning store. */ + return receipt?.status === 'reserved' ? null : (receipt ?? null); + } + + /** Resolves one authoritative receipt after its live owner disappears. The + * child conversation must still belong to the caller's parent thread, so a + * task id learned in another chat cannot cross orchestration scopes. */ + async function getSubagentTaskControlReplay({ + userId, + parentConversationId, + taskId, + invocationId, + tenantId, + }: { + userId: string; + parentConversationId: string; + taskId: string; + invocationId: string; + tenantId?: string; + }): Promise<{ + receipt: StoredSubagentControlReceipt; + task: { + taskId: string; + threadId: string; + subagentType: string; + status: 'running' | 'completed' | 'error' | 'cancelled'; + resultAvailable: boolean; + resultClaimed: boolean; + pendingControls: number; + createdAt: Date; + updatedAt: Date; + }; + } | null> { + if ( + userId.length === 0 || + parentConversationId.length === 0 || + parentConversationId.length > 256 || + taskId.length === 0 || + taskId.length > 256 || + invocationId.length === 0 || + invocationId.length > 128 + ) { + throw new TypeError('Invalid subagent task control replay identity'); + } + const Message = mongoose.models.Message as Model; + const input = await Message.findOne({ + user: userId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + messageId: `${taskId}:user`, + 'subagentTask.controlReceipts.invocationId': invocationId, + }) + .select({ + conversationId: 1, + createdAt: 1, + updatedAt: 1, + 'subagentTask.status': 1, + 'subagentTask.controlReceipts': 1, + _id: 0, + }) + .lean | null>(); + const receipt = input?.subagentTask?.controlReceipts?.find( + (candidate) => candidate.invocationId === invocationId, + ); + const status = input?.subagentTask?.status; + if ( + input == null || + receipt == null || + status == null || + input.createdAt == null || + input.updatedAt == null + ) { + return null; + } + const Conversation = mongoose.models.Conversation as Model; + const conversationQuery = Conversation.findOne({ + user: userId, + conversationId: input.conversationId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + 'subagentThread.parentConversationId': parentConversationId, + ...activeExpirationFilter(), + }) + .select({ 'subagentThread.subagentType': 1, _id: 0 }) + .lean | null>(); + const terminalQuery = Message.findOne({ + user: userId, + conversationId: input.conversationId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + messageId: `${taskId}:assistant`, + 'subagentTask.status': { $in: ['completed', 'error', 'cancelled'] }, + }) + .select({ updatedAt: 1, 'subagentTask.status': 1, 'subagentTask.resultClaim': 1, _id: 0 }) + .lean | null>(); + const [conversation, terminal] = await Promise.all([conversationQuery, terminalQuery]); + const subagentType = conversation?.subagentThread?.subagentType; + if (subagentType == null || subagentType === '') return null; + /** A committed cancel receipt is itself the authoritative cancellation + * boundary. The terminal row is written asynchronously and may not exist if + * the owner exits between those two durable commits. */ + const replayStatus = + terminal?.subagentTask?.status ?? + (receipt.action === 'cancel' && receipt.status === 'applied' ? 'cancelled' : status); + return { + receipt, + task: { + taskId, + threadId: input.conversationId, + subagentType, + status: replayStatus, + resultAvailable: terminal != null, + resultClaimed: terminal?.subagentTask?.resultClaim != null, + pendingControls: + input.subagentTask?.controlReceipts?.filter( + (candidate) => candidate.status === 'accepted', + ).length ?? 0, + createdAt: input.createdAt, + updatedAt: + terminal?.updatedAt ?? + (receipt.action === 'cancel' && receipt.status === 'applied' + ? receipt.updatedAt + : input.updatedAt), }, - [ - { - $set: { - 'subagentTask.controlReceipts': { - $let: { - vars: { - current: { - $cond: [ - { $isArray: '$subagentTask.controlReceipts' }, - '$subagentTask.controlReceipts', - [], - ], - }, - }, - in: { - $let: { - vars: { - existing: { - $arrayElemAt: [ - { - $filter: { - input: '$$current', - as: 'candidate', - cond: { $eq: ['$$candidate.invocationId', receipt.invocationId] }, - }, - }, - 0, - ], - }, - }, - in: { - $let: { - vars: { - next: { - $cond: [ - { - $or: [ - { - $in: [{ $ifNull: ['$$existing.status', ''] }, terminalStatuses], - }, - { - $and: [ - { $ne: [{ $ifNull: ['$$existing', null] }, null] }, - { $ne: ['$$existing.fingerprint', receipt.fingerprint] }, - ], - }, - ], - }, - '$$existing', - { $literal: receipt }, - ], - }, - }, - in: { - $let: { - vars: { - merged: { - $concatArrays: [ - { - $filter: { - input: '$$current', - as: 'candidate', - cond: { - $ne: ['$$candidate.invocationId', receipt.invocationId], - }, - }, - }, - ['$$next'], - ], - }, - }, - in: { - $let: { - vars: { - accepted: { - /** The supported task store admits at most 32 live - * controls. Keep a defensive storage bound here so - * custom callers cannot grow the private projection. */ - $slice: [ - { - $filter: { - input: '$$merged', - as: 'candidate', - cond: { $eq: ['$$candidate.status', 'accepted'] }, - }, - }, - -MAX_SUBAGENT_CONTROL_RECEIPTS, - ], - }, - }, - in: { - $concatArrays: [ - '$$accepted', - { - $slice: [ - { - /** DocumentDB 5 does not support $sortArray. - * Insert each bounded receipt into a stable - * createdAt/invocationId order using baseline - * aggregation expressions instead. */ - $reduce: { - input: { - $filter: { - input: '$$merged', - as: 'candidate', - cond: { - $ne: ['$$candidate.status', 'accepted'], - }, - }, - }, - initialValue: [], - in: { - $concatArrays: [ - { - $filter: { - input: '$$value', - as: 'ordered', - cond: { - $or: [ - { - $lt: [ - '$$ordered.createdAt', - '$$this.createdAt', - ], - }, - { - $and: [ - { - $eq: [ - '$$ordered.createdAt', - '$$this.createdAt', - ], - }, - { - $lte: [ - '$$ordered.invocationId', - '$$this.invocationId', - ], - }, - ], - }, - ], - }, - }, - }, - ['$$this'], - { - $filter: { - input: '$$value', - as: 'ordered', - cond: { - $or: [ - { - $gt: [ - '$$ordered.createdAt', - '$$this.createdAt', - ], - }, - { - $and: [ - { - $eq: [ - '$$ordered.createdAt', - '$$this.createdAt', - ], - }, - { - $gt: [ - '$$ordered.invocationId', - '$$this.invocationId', - ], - }, - ], - }, - ], - }, - }, - }, - ], - }, - }, - }, - { - $multiply: [ - -1, - { - $max: [ - 0, - { - $subtract: [ - MAX_SUBAGENT_CONTROL_RECEIPTS, - { $size: '$$accepted' }, - ], - }, - ], - }, - ], - }, - ], - }, - ], - }, - }, - }, - }, - }, - }, - }, - }, - }, - }, - }, - }, - }, - ], - { new: true, projection: { messageId: 1 } }, - ).lean<{ messageId: string } | null>(); - return updated != null; + }; } /** Atomically assigns one durable terminal child result to either its @@ -1749,6 +1843,8 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateToolCallResult, updateMessage, recordSubagentTaskControlReceipt, + getSubagentTaskControlReceipt, + getSubagentTaskControlReplay, claimSubagentTaskResult, releaseSubagentTaskResultClaim, deleteMessagesSince, diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index 56f594d1d6..b058fcb8e3 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -188,7 +188,7 @@ const messageSchema: Schema = new Schema( }, status: { type: String, - enum: ['accepted', 'applied', 'rejected', 'failed'], + enum: ['reserved', 'accepted', 'applied', 'rejected', 'failed'], required: true, }, createdAt: { type: Date, required: true }, diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index fa36036cff..7e3bacac61 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -12,7 +12,12 @@ export type SubagentTaskControlAction = | 'cancel' | 'cancel_message'; -export type SubagentTaskControlReceiptStatus = 'accepted' | 'applied' | 'rejected' | 'failed'; +export type SubagentTaskControlReceiptStatus = + | 'reserved' + | 'accepted' + | 'applied' + | 'rejected' + | 'failed'; /** Server-private durable receipt for one parent-to-child control invocation. */ export interface ISubagentTaskControlReceipt {