mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-27 12:13:30 +00:00
🧾 feat: Persist Authoritative Subagent Control Receipts (#15168)
* feat: persist subagent control receipts * fix: require control receipt persistence * fix: preserve authoritative control history
This commit is contained in:
parent
e0d5e11cdf
commit
a997275902
13 changed files with 1608 additions and 28 deletions
|
|
@ -64,6 +64,7 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore(
|
|||
getConvo: db.getConvo,
|
||||
getMessages: db.getMessages,
|
||||
listActiveSubagentThreadLeases: db.listActiveSubagentThreadLeases,
|
||||
recordSubagentTaskControlReceipt: db.recordSubagentTaskControlReceipt,
|
||||
releaseSubagentThreadLease: db.releaseSubagentThreadLease,
|
||||
reserveSubagentThread: db.reserveSubagentThread,
|
||||
renewSubagentThreadLease: db.renewSubagentThreadLease,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ jest.mock('~/models', () => ({
|
|||
getConvo: jest.fn(),
|
||||
getMessages: jest.fn(),
|
||||
listActiveSubagentThreadLeases: jest.fn(),
|
||||
recordSubagentTaskControlReceipt: jest.fn(),
|
||||
releaseSubagentThreadLease: jest.fn(),
|
||||
reserveSubagentThread: jest.fn(),
|
||||
renewSubagentThreadLease: jest.fn(),
|
||||
|
|
@ -55,11 +56,19 @@ const {
|
|||
const subagentThreadTaskStore = require('./subagentThreadStore');
|
||||
const { configureSubagentTaskRouting } = subagentThreadTaskStore;
|
||||
const taskStoreOptions = createSubagentThreadTaskStore.mock.calls[0][1];
|
||||
const taskStoreMethods = createSubagentThreadTaskStore.mock.calls[0][0];
|
||||
const db = require('~/models');
|
||||
const activityPrepareRegistration = registerShutdownTask.mock.calls.find(
|
||||
([name]) => name === 'subagent activity streams prepare',
|
||||
);
|
||||
|
||||
describe('subagent thread Redis lifecycle', () => {
|
||||
it('wires durable control receipt persistence into the host store', () => {
|
||||
expect(taskStoreMethods.recordSubagentTaskControlReceipt).toBe(
|
||||
db.recordSubagentTaskControlReceipt,
|
||||
);
|
||||
});
|
||||
|
||||
it('reads completion wakeup rollout state at task preparation time', async () => {
|
||||
isEnabled.mockReturnValueOnce(false);
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ function makeStore(): SubagentThreadTaskStore {
|
|||
getConvo: unused as AllMethods['getConvo'],
|
||||
getMessages: unused as AllMethods['getMessages'],
|
||||
listActiveSubagentThreadLeases: unused as AllMethods['listActiveSubagentThreadLeases'],
|
||||
recordSubagentTaskControlReceipt: unused as AllMethods['recordSubagentTaskControlReceipt'],
|
||||
releaseSubagentThreadLease: unused as AllMethods['releaseSubagentThreadLease'],
|
||||
reserveSubagentThread: unused as AllMethods['reserveSubagentThread'],
|
||||
renewSubagentThreadLease: unused as AllMethods['renewSubagentThreadLease'],
|
||||
|
|
|
|||
|
|
@ -172,7 +172,7 @@ export interface SubagentTaskControlHandler {
|
|||
taskId: string,
|
||||
command: SubagentTaskControlCommand,
|
||||
invocationId: string,
|
||||
): SubagentTaskControlResult;
|
||||
): Promise<SubagentTaskControlResult> | SubagentTaskControlResult;
|
||||
list(scopeId: string): SubagentTaskSnapshot[];
|
||||
cancelScope(scopeId: string, threadIds: string[] | null): number;
|
||||
}
|
||||
|
|
@ -1021,7 +1021,12 @@ export class RedisSubagentTaskControlTransport implements SubagentTaskControlTra
|
|||
result = claim;
|
||||
} else {
|
||||
result = boundedControlResult(
|
||||
handler.control(request.scopeId, request.taskId, request.command, request.invocationId),
|
||||
await handler.control(
|
||||
request.scopeId,
|
||||
request.taskId,
|
||||
request.command,
|
||||
request.invocationId,
|
||||
),
|
||||
);
|
||||
}
|
||||
const serializedResult = JSON.stringify(result);
|
||||
|
|
|
|||
|
|
@ -53,6 +53,24 @@ class TestTaskRoutingHub {
|
|||
}
|
||||
}
|
||||
|
||||
class ReceiptTestSubagentThreadTaskStore extends SubagentThreadTaskStore {
|
||||
emitControlReceiptForTest(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
receipt: {
|
||||
controlId: string;
|
||||
action: 'steer' | 'queue' | 'interrupt';
|
||||
status: 'accepted' | 'applied' | 'rejected' | 'failed';
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
boundary?: 'preempt' | 'tool' | 'turn';
|
||||
reason?: 'withdrawn' | 'task_completed' | 'task_cancelled' | 'task_failed';
|
||||
},
|
||||
): void {
|
||||
this.onControlReceipt(scopeId, taskId, receipt);
|
||||
}
|
||||
}
|
||||
|
||||
class TestTaskControlTransport implements SubagentTaskControlTransport {
|
||||
private handler?: SubagentTaskControlHandler;
|
||||
readonly registrations: Array<{ scopeId: string; taskId: string; ttlMs: number }> = [];
|
||||
|
|
@ -187,9 +205,12 @@ async function waitForSettled(
|
|||
throw new Error('Timed out waiting for the subagent task.');
|
||||
}
|
||||
|
||||
async function waitUntil(condition: () => boolean, description: string): Promise<void> {
|
||||
async function waitUntil(
|
||||
condition: () => boolean | Promise<boolean>,
|
||||
description: string,
|
||||
): Promise<void> {
|
||||
for (let attempt = 0; attempt < 400; attempt += 1) {
|
||||
if (condition()) {
|
||||
if (await condition()) {
|
||||
return;
|
||||
}
|
||||
await new Promise<void>((resolve) => setTimeout(resolve, 10));
|
||||
|
|
@ -1964,6 +1985,20 @@ describe('SubagentThreadTaskStore', () => {
|
|||
await expect(requesterStore.listTasks(config.scopeId)).resolves.toEqual([
|
||||
expect.objectContaining({ taskId, status: 'running' }),
|
||||
]);
|
||||
await waitUntil(
|
||||
async () =>
|
||||
(
|
||||
await methods.getMessages(
|
||||
{
|
||||
user: userId,
|
||||
conversationId: requireThreadId(started),
|
||||
messageId: `${taskId}:user`,
|
||||
},
|
||||
'+subagentTask',
|
||||
)
|
||||
).length === 1,
|
||||
'the durable control receipt target',
|
||||
);
|
||||
await expect(
|
||||
requesterStore.controlTask(config.scopeId, taskId, {
|
||||
action: 'queue',
|
||||
|
|
@ -1991,7 +2026,7 @@ describe('SubagentThreadTaskStore', () => {
|
|||
const parentConversationId = randomUUID();
|
||||
await saveParent(userId, parentConversationId);
|
||||
const hub = new TestTaskRoutingHub();
|
||||
const ownerStore = new SubagentThreadTaskStore(methods);
|
||||
const ownerStore = new ReceiptTestSubagentThreadTaskStore(methods);
|
||||
const requesterStore = new SubagentThreadTaskStore(methods);
|
||||
await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub));
|
||||
await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub));
|
||||
|
|
@ -2006,12 +2041,35 @@ describe('SubagentThreadTaskStore', () => {
|
|||
}),
|
||||
);
|
||||
const taskId = requireAccepted(started).task.taskId;
|
||||
await Promise.resolve();
|
||||
const threadId = requireThreadId(started);
|
||||
let durableInput: IMessage | undefined;
|
||||
for (let attempt = 0; attempt < 200; attempt += 1) {
|
||||
[durableInput] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
if (durableInput != null) break;
|
||||
await new Promise<void>((resolve) => setTimeout(resolve, 10));
|
||||
}
|
||||
expect(durableInput).toBeDefined();
|
||||
|
||||
const steer = { action: 'queue' as const, message: 'Verify the primary source too.' };
|
||||
const routed = await requesterStore.controlTask(config.scopeId, taskId, steer, 'invocation-1');
|
||||
expect(routed).toMatchObject({ status: 'accepted' });
|
||||
|
||||
[durableInput] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
expect(durableInput?.subagentTask?.controlReceipts).toEqual([
|
||||
expect.objectContaining({
|
||||
invocationId: 'invocation-1',
|
||||
action: 'queue',
|
||||
status: 'accepted',
|
||||
message: steer.message,
|
||||
}),
|
||||
]);
|
||||
|
||||
/** The same invocation reaching the owner directly replays that result rather than
|
||||
* queueing a second steer, so local and routed callers agree. */
|
||||
await expect(
|
||||
|
|
@ -2019,6 +2077,49 @@ describe('SubagentThreadTaskStore', () => {
|
|||
).resolves.toEqual(routed);
|
||||
expect(ownerStore.get(config.scopeId, taskId)?.pendingControls).toBe(1);
|
||||
|
||||
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',
|
||||
});
|
||||
await waitUntil(
|
||||
() => ownerStore.get(config.scopeId, taskId) != null,
|
||||
'the owner task to remain available',
|
||||
);
|
||||
for (let attempt = 0; attempt < 200; attempt += 1) {
|
||||
[durableInput] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
if (durableInput?.subagentTask?.controlReceipts?.[0]?.status === 'applied') break;
|
||||
await new Promise<void>((resolve) => setTimeout(resolve, 10));
|
||||
}
|
||||
expect(durableInput?.subagentTask?.controlReceipts).toEqual([
|
||||
expect.objectContaining({
|
||||
invocationId: 'invocation-1',
|
||||
status: 'applied',
|
||||
boundary: 'turn',
|
||||
}),
|
||||
]);
|
||||
|
||||
/** A delayed retry can replay accepted in memory but cannot downgrade the
|
||||
* already-applied durable receipt. */
|
||||
await expect(
|
||||
ownerStore.controlTask(config.scopeId, taskId, steer, 'invocation-1'),
|
||||
).resolves.toEqual(routed);
|
||||
[durableInput] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
expect(durableInput?.subagentTask?.controlReceipts?.[0]?.status).toBe('applied');
|
||||
|
||||
/** Reusing one invocation id for different content is a caller error, not a retry. */
|
||||
await expect(
|
||||
requesterStore.controlTask(
|
||||
|
|
@ -2038,6 +2139,406 @@ describe('SubagentThreadTaskStore', () => {
|
|||
]);
|
||||
});
|
||||
|
||||
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 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 durable task input',
|
||||
);
|
||||
|
||||
const persistReceipt = methods.recordSubagentTaskControlReceipt;
|
||||
const persistence = jest
|
||||
.spyOn(methods, 'recordSubagentTaskControlReceipt')
|
||||
.mockResolvedValueOnce(false)
|
||||
.mockImplementation(persistReceipt);
|
||||
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);
|
||||
|
||||
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 {
|
||||
persistence.mockRestore();
|
||||
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();
|
||||
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;
|
||||
const threadId = requireThreadId(started);
|
||||
await waitUntil(
|
||||
async () =>
|
||||
(
|
||||
await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
)
|
||||
).length === 1,
|
||||
'the durable task input',
|
||||
);
|
||||
|
||||
await expect(
|
||||
store.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'queue', message: 'x'.repeat(64 * 1024) },
|
||||
'bounded-payload-invocation',
|
||||
),
|
||||
).resolves.toMatchObject({ status: 'accepted' });
|
||||
const invocations = (
|
||||
store as unknown as {
|
||||
controlInvocations: Map<string, { command: SubagentTaskControlCommand }>;
|
||||
}
|
||||
).controlInvocations;
|
||||
expect(invocations.values().next().value?.command).toEqual({
|
||||
action: 'queue',
|
||||
message: 'x'.repeat(4 * 1024),
|
||||
});
|
||||
await expect(
|
||||
store.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'queue', message: `${'x'.repeat(64 * 1024 - 1)}y` },
|
||||
'bounded-payload-invocation',
|
||||
),
|
||||
).resolves.toMatchObject({ status: 'invalid' });
|
||||
|
||||
finish({ content: 'Done.' });
|
||||
await waitForSettled(store, config.scopeId, started);
|
||||
await store.destroyTaskControlTransport();
|
||||
});
|
||||
|
||||
it('restores tenant context for a routed control receipt write', async () => {
|
||||
const userId = 'routed-receipt-tenant-user';
|
||||
const tenantId = 'routed-receipt-tenant';
|
||||
const parentConversationId = randomUUID();
|
||||
await tenantStorage.run({ tenantId, userId }, async () =>
|
||||
saveParent(userId, parentConversationId, { tenantId }),
|
||||
);
|
||||
const hub = new TestTaskRoutingHub();
|
||||
const ownerStore = new SubagentThreadTaskStore(methods);
|
||||
const requesterStore = new SubagentThreadTaskStore(methods);
|
||||
await ownerStore.configureTaskControlTransport(new TestTaskControlTransport(hub));
|
||||
await requesterStore.configureTaskControlTransport(new TestTaskControlTransport(hub));
|
||||
const config = buildSubagentThreadTaskConfig(ownerStore, {
|
||||
userId,
|
||||
tenantId,
|
||||
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(
|
||||
() =>
|
||||
tenantStorage.run({ tenantId, userId }, async () =>
|
||||
Boolean(
|
||||
(
|
||||
await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
)
|
||||
)[0],
|
||||
),
|
||||
),
|
||||
'the tenant-scoped task input',
|
||||
);
|
||||
|
||||
const contexts: Array<{ tenantId?: string; userId?: string }> = [];
|
||||
const persistReceipt = methods.recordSubagentTaskControlReceipt;
|
||||
const persistence = jest
|
||||
.spyOn(methods, 'recordSubagentTaskControlReceipt')
|
||||
.mockImplementation((input) => {
|
||||
contexts.push({ tenantId: getTenantId(), userId: getUserId() });
|
||||
return persistReceipt(input);
|
||||
});
|
||||
try {
|
||||
await expect(
|
||||
requesterStore.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'queue', message: 'Check the tenant source.' },
|
||||
'tenant-invocation',
|
||||
),
|
||||
).resolves.toMatchObject({ status: 'accepted' });
|
||||
expect(contexts).toEqual([{ tenantId, userId }]);
|
||||
} finally {
|
||||
persistence.mockRestore();
|
||||
finish({ content: 'Done.' });
|
||||
await waitForSettled(ownerStore, config.scopeId, started);
|
||||
await Promise.all([
|
||||
ownerStore.destroyTaskControlTransport(),
|
||||
requesterStore.destroyTaskControlTransport(),
|
||||
]);
|
||||
}
|
||||
});
|
||||
|
||||
it('persists the target control id for cancel_message receipts', async () => {
|
||||
const userId = 'cancel-message-receipt-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;
|
||||
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 queued = await store.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'queue', message: 'Withdraw me.' },
|
||||
'queued-invocation',
|
||||
);
|
||||
expect(queued).toMatchObject({ status: 'accepted' });
|
||||
const targetControlId = queued.status === 'accepted' ? queued.controlId : undefined;
|
||||
expect(targetControlId).toBeDefined();
|
||||
await expect(
|
||||
store.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'cancel_message', controlId: targetControlId as string },
|
||||
'cancel-message-invocation',
|
||||
),
|
||||
).resolves.toMatchObject({ status: 'accepted' });
|
||||
await expect(
|
||||
store.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'cancel_message', controlId: 'missing-control' },
|
||||
'missing-cancel-message-invocation',
|
||||
),
|
||||
).resolves.toMatchObject({ status: 'control_not_found' });
|
||||
|
||||
await waitUntil(async () => {
|
||||
const [input] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
return (
|
||||
input?.subagentTask?.controlReceipts?.some(
|
||||
(receipt) => receipt.invocationId === 'cancel-message-invocation',
|
||||
) === true
|
||||
);
|
||||
}, 'the cancel_message receipt');
|
||||
const [input] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
expect(input?.subagentTask?.controlReceipts).toContainEqual(
|
||||
expect.objectContaining({
|
||||
invocationId: 'cancel-message-invocation',
|
||||
controlId: targetControlId,
|
||||
action: 'cancel_message',
|
||||
status: 'applied',
|
||||
}),
|
||||
);
|
||||
expect(input?.subagentTask?.controlReceipts).toContainEqual(
|
||||
expect.objectContaining({
|
||||
invocationId: 'missing-cancel-message-invocation',
|
||||
controlId: 'missing-control',
|
||||
action: 'cancel_message',
|
||||
status: 'rejected',
|
||||
reason: 'control_not_found',
|
||||
}),
|
||||
);
|
||||
|
||||
finish({ content: 'Done.' });
|
||||
await waitForSettled(store, config.scopeId, started);
|
||||
await store.destroyTaskControlTransport();
|
||||
});
|
||||
|
||||
it('retries a terminal control receipt after settlement when storage recovers', async () => {
|
||||
const userId = 'receipt-retry-user';
|
||||
const parentConversationId = randomUUID();
|
||||
await saveParent(userId, parentConversationId);
|
||||
const store = new ReceiptTestSubagentThreadTaskStore(methods, {
|
||||
controlReceiptRetryMs: 10,
|
||||
});
|
||||
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 durable task input',
|
||||
);
|
||||
const command = { action: 'queue' as const, message: 'Verify the source.' };
|
||||
const accepted = await store.controlTask(config.scopeId, taskId, command, 'retry-invocation');
|
||||
expect(accepted).toMatchObject({ status: 'accepted' });
|
||||
const controlId = accepted.status === 'accepted' ? accepted.controlId : undefined;
|
||||
expect(controlId).toBeDefined();
|
||||
|
||||
const persistReceipt = methods.recordSubagentTaskControlReceipt;
|
||||
const persistence = jest
|
||||
.spyOn(methods, 'recordSubagentTaskControlReceipt')
|
||||
.mockRejectedValueOnce(new Error('database temporarily unavailable'))
|
||||
.mockResolvedValueOnce(false)
|
||||
.mockImplementation(persistReceipt);
|
||||
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 applied transition persistence attempt',
|
||||
);
|
||||
finish({ content: 'Done.' });
|
||||
await waitForSettled(store, config.scopeId, started);
|
||||
|
||||
await waitUntil(
|
||||
() => persistence.mock.calls.length >= 3,
|
||||
'the post-settlement receipt retry',
|
||||
);
|
||||
await waitUntil(async () => {
|
||||
const [input] = await methods.getMessages(
|
||||
{ user: userId, conversationId: threadId, messageId: `${taskId}:user` },
|
||||
'+subagentTask',
|
||||
);
|
||||
return input?.subagentTask?.controlReceipts?.[0]?.status === 'applied';
|
||||
}, 'the receipt retry to converge after terminal settlement');
|
||||
expect(persistence.mock.calls.length).toBeGreaterThanOrEqual(3);
|
||||
} finally {
|
||||
persistence.mockRestore();
|
||||
finish({ content: 'Done.' });
|
||||
await store.destroyTaskControlTransport();
|
||||
}
|
||||
});
|
||||
|
||||
it('flushes a pending control receipt once during graceful shutdown', async () => {
|
||||
const userId = 'receipt-shutdown-user';
|
||||
const parentConversationId = randomUUID();
|
||||
await saveParent(userId, parentConversationId);
|
||||
const store = new ReceiptTestSubagentThreadTaskStore(methods, {
|
||||
controlReceiptRetryMs: 60_000,
|
||||
});
|
||||
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');
|
||||
const accepted = await store.controlTask(
|
||||
config.scopeId,
|
||||
taskId,
|
||||
{ action: 'queue', message: 'Persist before shutdown.' },
|
||||
'shutdown-invocation',
|
||||
);
|
||||
expect(accepted).toMatchObject({ status: 'accepted' });
|
||||
|
||||
const persistReceipt = methods.recordSubagentTaskControlReceipt;
|
||||
const persistence = jest
|
||||
.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',
|
||||
});
|
||||
await waitUntil(() => persistence.mock.calls.length === 1, 'the failed receipt write');
|
||||
|
||||
await store.destroyTaskControlTransport();
|
||||
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(persistence).toHaveBeenCalledTimes(2);
|
||||
|
||||
persistence.mockRestore();
|
||||
finish({ content: 'Done.' });
|
||||
await waitForSettled(store, config.scopeId, started);
|
||||
});
|
||||
|
||||
it('fails a child closed when its owner address cannot be published', async () => {
|
||||
const userId = 'unregistered-user';
|
||||
const parentConversationId = randomUUID();
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ import type {
|
|||
AllMethods,
|
||||
IActiveSubagentThreadLease,
|
||||
IConversation,
|
||||
ISubagentTaskControlReceipt,
|
||||
IMessage,
|
||||
MessageMethods,
|
||||
ConversationMethods,
|
||||
|
|
@ -62,6 +63,8 @@ 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;
|
||||
const MAX_DURABLE_CONTROL_MESSAGE_CHARS = 4 * 1024;
|
||||
const DEFAULT_CONTROL_RECEIPT_RETRY_MS = 5_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. */
|
||||
|
|
@ -99,6 +102,7 @@ type SubagentThreadMethods = Pick<
|
|||
| 'listActiveSubagentThreadLeases'
|
||||
| 'reserveSubagentThread'
|
||||
| 'releaseSubagentThreadLease'
|
||||
| 'recordSubagentTaskControlReceipt'
|
||||
| 'renewSubagentThreadLease'
|
||||
| 'saveConvo'
|
||||
| 'saveMessage'
|
||||
|
|
@ -132,6 +136,33 @@ type ThreadMessage = Pick<
|
|||
'messageId' | 'parentMessageId' | 'text' | 'createdAt' | 'subagentTranscript' | 'subagentTask'
|
||||
>;
|
||||
|
||||
type SdkControlReceipt = {
|
||||
controlId: string;
|
||||
action: 'steer' | 'queue' | 'interrupt';
|
||||
status: 'accepted' | 'applied' | 'rejected' | 'failed';
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
boundary?: 'preempt' | 'tool' | 'turn';
|
||||
reason?: 'withdrawn' | 'task_completed' | 'task_cancelled' | 'task_failed';
|
||||
};
|
||||
|
||||
type SnapshotWithControlReceipts = SubagentTaskSnapshot & {
|
||||
controlReceipts?: SdkControlReceipt[];
|
||||
};
|
||||
|
||||
type ControlInvocationRecord = {
|
||||
scopeId: string;
|
||||
taskId: string;
|
||||
invocationId: string;
|
||||
fingerprint: string;
|
||||
command: SubagentTaskControlCommand;
|
||||
result: SubagentTaskControlResult;
|
||||
createdAt: number;
|
||||
/** Last authoritative SDK transition, retained for idempotent retries even
|
||||
* after the bounded SDK snapshot evicts older receipt history. */
|
||||
receipt?: ISubagentTaskControlReceipt;
|
||||
};
|
||||
|
||||
interface TaskThreadLease {
|
||||
idempotencyKey: string;
|
||||
taskId: string;
|
||||
|
|
@ -188,6 +219,7 @@ export interface SubagentThreadTaskStoreOptions extends InMemorySubagentTaskStor
|
|||
taskRoutingTtlMs?: number;
|
||||
isOwnerActive?: (userId: string) => Promise<boolean>;
|
||||
maxControlInvocations?: number;
|
||||
controlReceiptRetryMs?: number;
|
||||
ownerFenceGraceMs?: number;
|
||||
fenceOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise<void>;
|
||||
renewOwnerAdmission?: (userId: string, token: string, fencedUntil: Date) => Promise<boolean>;
|
||||
|
|
@ -431,6 +463,42 @@ function drainKey(parentConversationId: string, taskId: string): string {
|
|||
return `${parentConversationId}\u0000${taskId}`;
|
||||
}
|
||||
|
||||
function controlTaskKey(scopeId: string, taskId: string): string {
|
||||
return `${scopeId}\u0000${taskId}`;
|
||||
}
|
||||
|
||||
function parseControlTaskKey(key: string): { scopeId: string; taskId: string } | undefined {
|
||||
const separator = key.lastIndexOf('\u0000');
|
||||
if (separator < 0 || separator === key.length - 1) return undefined;
|
||||
return { scopeId: key.slice(0, separator), taskId: key.slice(separator + 1) };
|
||||
}
|
||||
|
||||
function controlReceiptKey(scopeId: string, taskId: string, controlId: string): string {
|
||||
return `${scopeId}\u0000${taskId}\u0000${controlId}`;
|
||||
}
|
||||
|
||||
function boundedControlMessage(command: SubagentTaskControlCommand): {
|
||||
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.slice(0, MAX_DURABLE_CONTROL_MESSAGE_CHARS),
|
||||
messageTruncated: true,
|
||||
};
|
||||
}
|
||||
|
||||
function boundedControlCommand(command: SubagentTaskControlCommand): SubagentTaskControlCommand {
|
||||
if (!('message' in command)) return command;
|
||||
return {
|
||||
action: command.action,
|
||||
message: command.message.slice(0, MAX_DURABLE_CONTROL_MESSAGE_CHARS),
|
||||
};
|
||||
}
|
||||
|
||||
function safeErrorMessage(error: unknown): string {
|
||||
return `Subagent task failed: ${publicFailureDetail(error).slice(0, 2_000)}`;
|
||||
}
|
||||
|
|
@ -453,11 +521,18 @@ async function observeSlowPreparation<T>(
|
|||
export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
||||
readonly supportsThreadContinuation = true;
|
||||
private readonly activeThreads = new Map<string, TaskThreadLease>();
|
||||
private readonly controlInvocations = new Map<
|
||||
private readonly controlInvocations = new Map<string, ControlInvocationRecord>();
|
||||
|
||||
private readonly controlInvocationByReceipt = new Map<string, ControlInvocationRecord>();
|
||||
private readonly pendingControlReceipts = new Map<
|
||||
string,
|
||||
{ scopeId: string; taskId: string; fingerprint: string; result: SubagentTaskControlResult }
|
||||
Map<string, { threadId: string; receipt: ISubagentTaskControlReceipt }>
|
||||
>();
|
||||
|
||||
private readonly controlPersistenceTails = new Map<string, Promise<void>>();
|
||||
private readonly controlPersistenceRetryTimers = new Map<string, ReturnType<typeof setTimeout>>();
|
||||
private controlPersistenceStopping = false;
|
||||
|
||||
private readonly parentPersistence = new Map<string, Promise<unknown>>();
|
||||
private readonly maxThreadDepth: number;
|
||||
private readonly leaseTtlMs: number;
|
||||
|
|
@ -466,6 +541,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
private readonly ownerDrainPollMs: number;
|
||||
private readonly taskRoutingTtlMs: number;
|
||||
private readonly maxControlInvocations: number;
|
||||
private readonly controlReceiptRetryMs: number;
|
||||
private readonly ownerFenceGraceMs: number;
|
||||
private readonly isOwnerActive: (userId: string) => Promise<boolean>;
|
||||
private readonly fenceOwnerAdmission?: (
|
||||
|
|
@ -510,6 +586,10 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
options.maxControlInvocations,
|
||||
MAX_CONTROL_INVOCATIONS,
|
||||
);
|
||||
this.controlReceiptRetryMs = positiveInteger(
|
||||
options.controlReceiptRetryMs,
|
||||
DEFAULT_CONTROL_RECEIPT_RETRY_MS,
|
||||
);
|
||||
this.ownerFenceGraceMs = positiveInteger(options.ownerFenceGraceMs, OWNER_FENCE_GRACE_MS);
|
||||
this.isOwnerActive = options.isOwnerActive ?? (async () => true);
|
||||
this.fenceOwnerAdmission = options.fenceOwnerAdmission;
|
||||
|
|
@ -519,6 +599,191 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
this.onTaskPrepared = options.onTaskPrepared;
|
||||
}
|
||||
|
||||
/** Receives payload-free authoritative transitions from the SDK task store. */
|
||||
protected onControlReceipt(scopeId: string, taskId: string, receipt: SdkControlReceipt): void {
|
||||
const invocation = this.controlInvocationByReceipt.get(
|
||||
controlReceiptKey(scopeId, taskId, receipt.controlId),
|
||||
);
|
||||
const threadId = this.get(scopeId, taskId)?.threadId;
|
||||
if (invocation == null || threadId == null) return;
|
||||
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);
|
||||
});
|
||||
}
|
||||
|
||||
private durableReceipt(
|
||||
invocation: ControlInvocationRecord,
|
||||
receipt: SdkControlReceipt,
|
||||
): ISubagentTaskControlReceipt {
|
||||
return {
|
||||
invocationId: invocation.invocationId,
|
||||
fingerprint: invocation.fingerprint,
|
||||
controlId: receipt.controlId,
|
||||
action: receipt.action,
|
||||
status: receipt.status,
|
||||
createdAt: new Date(receipt.createdAt),
|
||||
updatedAt: new Date(receipt.updatedAt),
|
||||
...(receipt.boundary == null ? {} : { boundary: receipt.boundary }),
|
||||
...(receipt.reason == null ? {} : { reason: receipt.reason }),
|
||||
...boundedControlMessage(invocation.command),
|
||||
};
|
||||
}
|
||||
|
||||
private controlResultReceipt(
|
||||
invocation: ControlInvocationRecord,
|
||||
): ISubagentTaskControlReceipt | undefined {
|
||||
const { command, result } = invocation;
|
||||
if (result.status === 'not_found' || result.status === 'invalid') return undefined;
|
||||
if (invocation.receipt != null) return invocation.receipt;
|
||||
const snapshot = result.task as SnapshotWithControlReceipts;
|
||||
if (
|
||||
result.status === 'accepted' &&
|
||||
result.controlId != null &&
|
||||
(command.action === 'steer' || command.action === 'queue' || command.action === 'interrupt')
|
||||
) {
|
||||
const sdkReceipt = snapshot.controlReceipts?.find(
|
||||
(receipt) => receipt.controlId === result.controlId,
|
||||
);
|
||||
if (sdkReceipt != null) return this.durableReceipt(invocation, sdkReceipt);
|
||||
return {
|
||||
invocationId: invocation.invocationId,
|
||||
fingerprint: invocation.fingerprint,
|
||||
controlId: result.controlId,
|
||||
action: command.action,
|
||||
status: 'accepted',
|
||||
createdAt: new Date(invocation.createdAt),
|
||||
updatedAt: new Date(invocation.createdAt),
|
||||
...boundedControlMessage(command),
|
||||
};
|
||||
}
|
||||
const now = new Date();
|
||||
let reason: string | undefined;
|
||||
if (result.status === 'not_running') {
|
||||
reason = 'task_not_running';
|
||||
} else if (result.status === 'control_not_found') {
|
||||
reason = 'control_not_found';
|
||||
}
|
||||
let targetControlId: string | undefined;
|
||||
if (command.action === 'cancel_message') {
|
||||
targetControlId = command.controlId;
|
||||
} else if (result.status === 'accepted') {
|
||||
targetControlId = result.controlId;
|
||||
}
|
||||
return {
|
||||
invocationId: invocation.invocationId,
|
||||
fingerprint: invocation.fingerprint,
|
||||
...(targetControlId == null ? {} : { controlId: targetControlId }),
|
||||
action: command.action,
|
||||
status:
|
||||
result.status === 'accepted' || result.status === 'cancelled' ? 'applied' : 'rejected',
|
||||
createdAt: new Date(invocation.createdAt),
|
||||
updatedAt: now,
|
||||
...(reason == null ? {} : { reason }),
|
||||
...boundedControlMessage(command),
|
||||
};
|
||||
}
|
||||
|
||||
private queueControlReceipt(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
threadId: string,
|
||||
receipt: ISubagentTaskControlReceipt,
|
||||
): Promise<void> {
|
||||
const key = controlTaskKey(scopeId, taskId);
|
||||
const pending = this.pendingControlReceipts.get(key) ?? new Map();
|
||||
pending.set(receipt.invocationId, { threadId, receipt });
|
||||
this.pendingControlReceipts.set(key, pending);
|
||||
return this.flushControlReceipts(scopeId, taskId);
|
||||
}
|
||||
|
||||
private flushControlReceipts(scopeId: string, taskId: string): Promise<void> {
|
||||
const key = controlTaskKey(scopeId, taskId);
|
||||
const prior = this.controlPersistenceTails.get(key) ?? Promise.resolve();
|
||||
const operation = prior
|
||||
.catch(() => undefined)
|
||||
.then(async () => {
|
||||
const pending = this.pendingControlReceipts.get(key);
|
||||
if (pending == null) return;
|
||||
const scope = parseScope(scopeId);
|
||||
for (const [invocationId, candidate] of [...pending]) {
|
||||
const current = pending.get(invocationId);
|
||||
if (current !== candidate) continue;
|
||||
const persisted = await this.runWithOwnerContext(scope, () =>
|
||||
this.methods.recordSubagentTaskControlReceipt({
|
||||
userId: scope.userId,
|
||||
conversationId: candidate.threadId,
|
||||
taskId,
|
||||
...(scope.tenantId == null ? {} : { tenantId: scope.tenantId }),
|
||||
receipt: candidate.receipt,
|
||||
}),
|
||||
);
|
||||
if (!persisted) {
|
||||
throw new Error('The child control receipt target is not ready.');
|
||||
}
|
||||
if (pending.get(invocationId) === candidate) {
|
||||
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);
|
||||
}
|
||||
});
|
||||
this.controlPersistenceTails.set(key, operation);
|
||||
void operation.then(
|
||||
() => {
|
||||
if (this.controlPersistenceTails.get(key) === operation) {
|
||||
this.controlPersistenceTails.delete(key);
|
||||
}
|
||||
this.scheduleControlReceiptRetry(scopeId, taskId);
|
||||
},
|
||||
() => {
|
||||
if (this.controlPersistenceTails.get(key) === operation) {
|
||||
this.controlPersistenceTails.delete(key);
|
||||
}
|
||||
this.scheduleControlReceiptRetry(scopeId, taskId);
|
||||
},
|
||||
);
|
||||
return operation;
|
||||
}
|
||||
|
||||
/** 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. */
|
||||
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) ||
|
||||
!this.pendingControlReceipts.has(key)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
const timer = setTimeout(() => {
|
||||
this.controlPersistenceRetryTimers.delete(key);
|
||||
void this.flushControlReceipts(scopeId, taskId).catch((error) => {
|
||||
logger.warn('[subagentThreads] Failed to retry child control receipts', error);
|
||||
});
|
||||
}, this.controlReceiptRetryMs);
|
||||
this.controlPersistenceRetryTimers.set(key, timer);
|
||||
}
|
||||
|
||||
private async flushControlReceiptsForSettlement(scopeId: string, taskId: string): Promise<void> {
|
||||
try {
|
||||
await this.flushControlReceipts(scopeId, taskId);
|
||||
} catch (error) {
|
||||
logger.warn('[subagentThreads] Failed to flush child control receipts', error);
|
||||
}
|
||||
}
|
||||
|
||||
/** Enables optional cross-replica lookup after the host's Redis service is ready. */
|
||||
async configureTaskControlTransport(transport: SubagentTaskControlTransport): Promise<void> {
|
||||
if (this.taskControlTransport != null) {
|
||||
|
|
@ -527,7 +792,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
await transport.bind({
|
||||
claim: (scopeId, taskId) => super.claim(scopeId, taskId),
|
||||
control: (scopeId, taskId, command, invocationId) =>
|
||||
this.controlInvocation(scopeId, taskId, command, invocationId),
|
||||
this.controlInvocationAndPersist(scopeId, taskId, command, invocationId),
|
||||
list: (scopeId) => super.list(scopeId),
|
||||
cancelScope: (scopeId, threadIds) => this.cancelForScope(scopeId, threadIds),
|
||||
});
|
||||
|
|
@ -535,6 +800,16 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
}
|
||||
|
||||
async destroyTaskControlTransport(): Promise<void> {
|
||||
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());
|
||||
const transport = this.taskControlTransport;
|
||||
this.taskControlTransport = undefined;
|
||||
await transport?.destroy();
|
||||
|
|
@ -1039,7 +1314,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
command: SubagentTaskControlCommand,
|
||||
invocationId: string = randomUUID(),
|
||||
): Promise<SubagentTaskControlResult> {
|
||||
const local = this.controlInvocation(scopeId, taskId, command, invocationId);
|
||||
const local = await this.controlInvocationAndPersist(scopeId, taskId, command, invocationId);
|
||||
if (local.status !== 'not_found') {
|
||||
return local;
|
||||
}
|
||||
|
|
@ -1093,7 +1368,50 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
if (result.status === 'not_found') {
|
||||
return result;
|
||||
}
|
||||
this.controlInvocations.set(key, { scopeId, taskId, fingerprint, result });
|
||||
const invocation: ControlInvocationRecord = {
|
||||
scopeId,
|
||||
taskId,
|
||||
invocationId,
|
||||
fingerprint,
|
||||
command: boundedControlCommand(command),
|
||||
result,
|
||||
createdAt: Date.now(),
|
||||
};
|
||||
this.controlInvocations.set(key, invocation);
|
||||
if (
|
||||
result.status === 'accepted' &&
|
||||
result.controlId != null &&
|
||||
(command.action === 'steer' || command.action === 'queue' || command.action === 'interrupt')
|
||||
) {
|
||||
this.controlInvocationByReceipt.set(
|
||||
controlReceiptKey(scopeId, taskId, result.controlId),
|
||||
invocation,
|
||||
);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private async controlInvocationAndPersist(
|
||||
scopeId: string,
|
||||
taskId: string,
|
||||
command: SubagentTaskControlCommand,
|
||||
invocationId: string,
|
||||
): Promise<SubagentTaskControlResult> {
|
||||
const result = this.controlInvocation(scopeId, taskId, command, invocationId);
|
||||
const invocation = this.controlInvocations.get(
|
||||
`${scopeId}\u0000${taskId}\u0000${invocationId}`,
|
||||
);
|
||||
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();
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
|
|
@ -1111,6 +1429,13 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
for (const [key, invocation] of this.controlInvocations) {
|
||||
if (this.get(invocation.scopeId, invocation.taskId) == null) {
|
||||
this.controlInvocations.delete(key);
|
||||
const result = invocation.result;
|
||||
if (result.status === 'accepted' && result.controlId != null) {
|
||||
this.controlInvocationByReceipt.delete(
|
||||
controlReceiptKey(invocation.scopeId, invocation.taskId, result.controlId),
|
||||
);
|
||||
}
|
||||
this.pendingControlReceipts.delete(controlTaskKey(invocation.scopeId, invocation.taskId));
|
||||
}
|
||||
}
|
||||
return this.controlInvocations.size < this.maxControlInvocations;
|
||||
|
|
@ -1298,25 +1623,29 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
(lease) => removed.has(lease.parentConversationId) || removed.has(lease.conversationId),
|
||||
)
|
||||
.map((lease) =>
|
||||
cancelSlot(() =>
|
||||
this.controlTask(
|
||||
serializeScope({
|
||||
userId,
|
||||
parentConversationId: lease.parentConversationId,
|
||||
...(tenantId ? { tenantId } : {}),
|
||||
}),
|
||||
lease.taskId,
|
||||
{ action: 'cancel' },
|
||||
),
|
||||
),
|
||||
cancelSlot(async () => {
|
||||
const scopeId = serializeScope({
|
||||
userId,
|
||||
parentConversationId: lease.parentConversationId,
|
||||
...(tenantId ? { tenantId } : {}),
|
||||
});
|
||||
const stopped = await transport.cancelScope(scopeId, [lease.conversationId]);
|
||||
if (stopped > 0 || this.cancelUnroutedTask == null) return stopped;
|
||||
return (await this.cancelUnroutedTask({
|
||||
userId,
|
||||
parentConversationId: lease.parentConversationId,
|
||||
taskId: lease.taskId,
|
||||
...(tenantId ? { tenantId } : {}),
|
||||
}))
|
||||
? 1
|
||||
: 0;
|
||||
}),
|
||||
);
|
||||
for (const count of await Promise.all(scopeCancellations)) {
|
||||
cancelled += count;
|
||||
}
|
||||
for (const result of await Promise.all(leaseCancellations)) {
|
||||
if (result.status === 'cancelled') {
|
||||
cancelled += 1;
|
||||
}
|
||||
for (const count of await Promise.all(leaseCancellations)) {
|
||||
cancelled += count;
|
||||
}
|
||||
return cancelled;
|
||||
}
|
||||
|
|
@ -1943,6 +2272,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
if (savedUserMessage == null) {
|
||||
throw new Error('Unable to persist the child-thread input.');
|
||||
}
|
||||
await this.flushControlReceiptsForSettlement(scopeId, taskId);
|
||||
const currentParent = await this.methods.getConvo(scope.userId, scope.parentConversationId);
|
||||
if (currentParent == null || !matchesTenant(currentParent.tenantId, scope.tenantId)) {
|
||||
throw new SubagentThreadPublicError('Parent thread is unavailable.');
|
||||
|
|
@ -2013,6 +2343,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
if (prepared.userMessageId == null) {
|
||||
throw new Error('The child-thread input was not prepared.');
|
||||
}
|
||||
await this.flushControlReceiptsForSettlement(request.scopeId, taskId);
|
||||
const subagentTranscript = serializeTranscript(
|
||||
taskId,
|
||||
prepared.initialStoredMessages,
|
||||
|
|
@ -2063,6 +2394,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
error: unknown,
|
||||
detachedUsage: UsageMetadata[],
|
||||
): Promise<void> {
|
||||
await this.flushControlReceiptsForSettlement(request.scopeId, taskId);
|
||||
const conversation = await this.currentConversation(scope, request, threadId);
|
||||
if (conversation == null || !(await this.taskInputExists(scope, threadId, taskId))) {
|
||||
return;
|
||||
|
|
@ -2128,6 +2460,7 @@ export class SubagentThreadTaskStore extends InMemorySubagentTaskStore {
|
|||
taskId: string,
|
||||
detachedUsage: UsageMetadata[],
|
||||
): Promise<void> {
|
||||
await this.flushControlReceiptsForSettlement(request.scopeId, taskId);
|
||||
const conversation = await this.currentConversation(scope, request, threadId);
|
||||
if (conversation == null || !(await this.taskInputExists(scope, threadId, taskId))) {
|
||||
return;
|
||||
|
|
@ -2308,6 +2641,7 @@ const REQUIRED_THREAD_METHODS = [
|
|||
'getConvo',
|
||||
'getMessages',
|
||||
'listActiveSubagentThreadLeases',
|
||||
'recordSubagentTaskControlReceipt',
|
||||
'releaseSubagentThreadLease',
|
||||
'renewSubagentThreadLease',
|
||||
'reserveSubagentThread',
|
||||
|
|
@ -2330,7 +2664,11 @@ export function createSubagentThreadTaskStore(
|
|||
> &
|
||||
Pick<
|
||||
MessageMethods,
|
||||
'claimSubagentTaskResult' | 'deleteMessages' | 'getMessages' | 'saveMessage'
|
||||
| 'claimSubagentTaskResult'
|
||||
| 'deleteMessages'
|
||||
| 'getMessages'
|
||||
| 'recordSubagentTaskControlReceipt'
|
||||
| 'saveMessage'
|
||||
>,
|
||||
options?: SubagentThreadTaskStoreOptions,
|
||||
): SubagentThreadTaskStore {
|
||||
|
|
|
|||
|
|
@ -127,6 +127,7 @@ describe('subagent thread parent-scoped view', () => {
|
|||
status: 'completed',
|
||||
activity: [],
|
||||
activityTruncated: false,
|
||||
controlReceipts: [],
|
||||
messages: [
|
||||
expect.objectContaining({ messageId: 'task-1:user', role: 'user' }),
|
||||
expect.objectContaining({
|
||||
|
|
@ -203,6 +204,63 @@ describe('subagent thread parent-scoped view', () => {
|
|||
expect(view.messages[0]).not.toHaveProperty('subagentTranscript');
|
||||
});
|
||||
|
||||
it('returns bounded authoritative control receipts without private fingerprints', async () => {
|
||||
const input = message('task-1:user', 'running', true);
|
||||
input.subagentTask!.controlReceipts = [
|
||||
...Array.from({ length: 32 }, (_, index) => ({
|
||||
invocationId: `earlier-${index}`,
|
||||
fingerprint: `private-${index}`,
|
||||
action: 'queue' as const,
|
||||
status: 'applied' as const,
|
||||
createdAt: new Date(`2026-08-21T10:00:${String(index).padStart(2, '0')}.000Z`),
|
||||
updatedAt: new Date(`2026-08-21T10:00:${String(index).padStart(2, '0')}.000Z`),
|
||||
})),
|
||||
{
|
||||
invocationId: 'invocation-1',
|
||||
fingerprint: 'private-fingerprint',
|
||||
controlId: 'control-1',
|
||||
action: 'steer',
|
||||
status: 'applied',
|
||||
createdAt: new Date('2026-08-21T11:00:01.000Z'),
|
||||
updatedAt: new Date('2026-08-21T11:00:02.000Z'),
|
||||
boundary: 'tool',
|
||||
message: 'x'.repeat(1_000),
|
||||
},
|
||||
];
|
||||
const handler = createSubagentThreadViewHandler({
|
||||
getConvoOwnership: jest.fn().mockResolvedValue(parent),
|
||||
getSubagentThreadForParent: jest.fn().mockResolvedValue(child),
|
||||
getMessagesForSubagentThreadView: jest
|
||||
.fn()
|
||||
.mockResolvedValue([message('task-1:assistant', 'completed'), input]),
|
||||
});
|
||||
const { response, json } = createResponse();
|
||||
|
||||
await handler(createRequest({}, { taskId: 'task-1' }), response);
|
||||
|
||||
const view = json.mock.calls[0][0];
|
||||
expect(view.controlReceipts).toEqual(
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
invocationId: 'invocation-1',
|
||||
controlId: 'control-1',
|
||||
action: 'steer',
|
||||
status: 'applied',
|
||||
boundary: 'tool',
|
||||
messageTruncated: true,
|
||||
}),
|
||||
]),
|
||||
);
|
||||
const projected = view.controlReceipts.find(
|
||||
(receipt: { invocationId: string }) => receipt.invocationId === 'invocation-1',
|
||||
);
|
||||
expect(projected).toBeDefined();
|
||||
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-fingerprint');
|
||||
});
|
||||
|
||||
it('fences replacement activity to the exact selected task input', async () => {
|
||||
const selected = {
|
||||
...message('task-1:assistant', 'completed'),
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import type {
|
|||
ParentSubagentIndex,
|
||||
ParentSubagentSummary,
|
||||
ParentSubagentTaskSummary,
|
||||
SubagentControlReceipt,
|
||||
SubagentThreadMessage,
|
||||
SubagentThreadStatus,
|
||||
SubagentThreadView,
|
||||
|
|
@ -30,6 +31,8 @@ const MAX_TITLE_BYTES = 1024;
|
|||
const MAX_PARENT_CHILDREN = 64;
|
||||
const MAX_PARENT_TASKS_PER_CHILD = 20;
|
||||
const MAX_PARENT_INDEX_BYTES = 96 * 1024;
|
||||
const MAX_PUBLIC_CONTROL_RECEIPTS = 32;
|
||||
const MAX_PUBLIC_CONTROL_MESSAGE_BYTES = 512;
|
||||
type SubagentThreadViewDependencies = Pick<
|
||||
ConversationMethods,
|
||||
'getConvoOwnership' | 'getSubagentThreadForParent'
|
||||
|
|
@ -116,6 +119,44 @@ const publicMessage = (
|
|||
};
|
||||
};
|
||||
|
||||
const publicControlReceipts = (
|
||||
messages: SubagentThreadViewMessageRecord[],
|
||||
taskId: string,
|
||||
): { 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');
|
||||
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)
|
||||
.map((receipt) => {
|
||||
const message =
|
||||
receipt.message == null
|
||||
? undefined
|
||||
: truncateUtf8(receipt.message, MAX_PUBLIC_CONTROL_MESSAGE_BYTES);
|
||||
return {
|
||||
invocationId: truncateUtf8(receipt.invocationId, MAX_PUBLIC_ID_BYTES).text,
|
||||
...(receipt.controlId == null
|
||||
? {}
|
||||
: { controlId: truncateUtf8(receipt.controlId, MAX_PUBLIC_ID_BYTES).text }),
|
||||
action: receipt.action,
|
||||
status: receipt.status,
|
||||
createdAt: isoDate(receipt.createdAt) ?? new Date(0).toISOString(),
|
||||
updatedAt: isoDate(receipt.updatedAt) ?? new Date(0).toISOString(),
|
||||
...(receipt.boundary == null ? {} : { boundary: receipt.boundary }),
|
||||
...(receipt.reason == null
|
||||
? {}
|
||||
: { reason: truncateUtf8(receipt.reason, MAX_PUBLIC_ID_BYTES).text }),
|
||||
...(message == null ? {} : { message: message.text }),
|
||||
...(receipt.messageTruncated === true || message?.truncated === true
|
||||
? { messageTruncated: true }
|
||||
: {}),
|
||||
};
|
||||
});
|
||||
return { receipts: retained, truncated: retained.length < stored.length };
|
||||
};
|
||||
|
||||
const publicStatus = (
|
||||
messages: SubagentThreadViewMessageRecord[],
|
||||
activeLeaseTaskId: string | undefined,
|
||||
|
|
@ -444,6 +485,10 @@ export function createSubagentThreadViewHandler(deps: SubagentThreadViewDependen
|
|||
projectedNewestFirst.push(projected.message);
|
||||
remainingTextBytes -= projected.bytes;
|
||||
}
|
||||
const projectedControls =
|
||||
requestedTaskId == null
|
||||
? { receipts: [], truncated: false }
|
||||
: publicControlReceipts(newestFirst, requestedTaskId);
|
||||
const view: SubagentThreadView = {
|
||||
threadId,
|
||||
parentConversationId,
|
||||
|
|
@ -461,6 +506,8 @@ export function createSubagentThreadViewHandler(deps: SubagentThreadViewDependen
|
|||
status: publicStatus(newestFirst, activeLeaseTaskId, requestedTaskId),
|
||||
activity: projectedActivity.activity,
|
||||
activityTruncated: projectedActivity.truncated,
|
||||
controlReceipts: projectedControls.receipts,
|
||||
...(projectedControls.truncated ? { controlReceiptsTruncated: true } : {}),
|
||||
messages: projectedNewestFirst.reverse(),
|
||||
historyTruncated: historyTruncated || projectedNewestFirst.length < newestFirst.length,
|
||||
...(isoDate(child.updatedAt) == null ? {} : { updatedAt: isoDate(child.updatedAt) }),
|
||||
|
|
|
|||
|
|
@ -66,6 +66,19 @@ export type SubagentActivityItem =
|
|||
outputTruncated?: boolean;
|
||||
};
|
||||
|
||||
export type SubagentControlReceipt = {
|
||||
invocationId: string;
|
||||
controlId?: string;
|
||||
action: 'steer' | 'queue' | 'interrupt' | 'cancel' | 'cancel_message';
|
||||
status: 'accepted' | 'applied' | 'rejected' | 'failed';
|
||||
createdAt: string;
|
||||
updatedAt: string;
|
||||
boundary?: 'preempt' | 'tool' | 'turn';
|
||||
reason?: string;
|
||||
message?: string;
|
||||
messageTruncated?: boolean;
|
||||
};
|
||||
|
||||
export type SubagentThreadMessage = {
|
||||
messageId: string;
|
||||
parentMessageId: string | null;
|
||||
|
|
@ -89,6 +102,10 @@ export type SubagentThreadView = {
|
|||
/** Activity for the exact task requested by the parent card, when retained. */
|
||||
activity: SubagentActivityItem[];
|
||||
activityTruncated: boolean;
|
||||
/** Bounded authoritative parent-to-child command receipts for this task. */
|
||||
controlReceipts?: SubagentControlReceipt[];
|
||||
/** True when older authoritative command receipts were omitted from this view. */
|
||||
controlReceiptsTruncated?: boolean;
|
||||
messages: SubagentThreadMessage[];
|
||||
historyTruncated: boolean;
|
||||
updatedAt?: string;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import mongoose from 'mongoose';
|
||||
import { v4 as uuidv4 } from 'uuid';
|
||||
import { RetentionMode } from 'librechat-data-provider';
|
||||
import { MongoMemoryServer } from 'mongodb-memory-server';
|
||||
import { Constants, RetentionMode } from 'librechat-data-provider';
|
||||
import type { IMessage } from '..';
|
||||
import {
|
||||
createMessageMethods,
|
||||
|
|
@ -39,6 +39,9 @@ let updateMessageText: ReturnType<typeof createMessageMethods>['updateMessageTex
|
|||
let deleteMessagesSince: ReturnType<typeof createMessageMethods>['deleteMessagesSince'];
|
||||
let recordMessage: ReturnType<typeof createMessageMethods>['recordMessage'];
|
||||
let claimSubagentTaskResult: ReturnType<typeof createMessageMethods>['claimSubagentTaskResult'];
|
||||
let recordSubagentTaskControlReceipt: ReturnType<
|
||||
typeof createMessageMethods
|
||||
>['recordSubagentTaskControlReceipt'];
|
||||
let releaseSubagentTaskResultClaim: ReturnType<
|
||||
typeof createMessageMethods
|
||||
>['releaseSubagentTaskResultClaim'];
|
||||
|
|
@ -64,6 +67,7 @@ beforeAll(async () => {
|
|||
deleteMessagesSince = methods.deleteMessagesSince;
|
||||
recordMessage = methods.recordMessage;
|
||||
claimSubagentTaskResult = methods.claimSubagentTaskResult;
|
||||
recordSubagentTaskControlReceipt = methods.recordSubagentTaskControlReceipt;
|
||||
releaseSubagentTaskResultClaim = methods.releaseSubagentTaskResultClaim;
|
||||
|
||||
await mongoose.connect(mongoUri);
|
||||
|
|
@ -2298,6 +2302,274 @@ describe('Message Operations', () => {
|
|||
expect((doc as Record<string, unknown> | null)?.unknownPipelineField).toBeUndefined();
|
||||
});
|
||||
});
|
||||
describe('recordSubagentTaskControlReceipt', () => {
|
||||
const createTaskInput = async (conversationId: string, taskId = 'task-1') => {
|
||||
await Message.create({
|
||||
user: 'user123',
|
||||
conversationId,
|
||||
messageId: `${taskId}:user`,
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
sender: 'User',
|
||||
text: 'Do the work',
|
||||
endpoint: 'agents',
|
||||
isCreatedByUser: true,
|
||||
subagentTask: {
|
||||
attemptKey: `attempt-${taskId}`,
|
||||
status: 'running',
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
it('advances one invocation monotonically and enforces ownership', async () => {
|
||||
const conversationId = uuidv4();
|
||||
await createTaskInput(conversationId);
|
||||
const accepted = {
|
||||
invocationId: 'invocation-1',
|
||||
fingerprint: 'fingerprint-1',
|
||||
controlId: 'control-1',
|
||||
action: 'steer' as const,
|
||||
status: 'accepted' as const,
|
||||
createdAt: new Date('2026-08-24T12:00:00.000Z'),
|
||||
updatedAt: new Date('2026-08-24T12:00:00.000Z'),
|
||||
message: 'Use the primary source.',
|
||||
};
|
||||
|
||||
await expect(
|
||||
recordSubagentTaskControlReceipt({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: 'task-1',
|
||||
receipt: accepted,
|
||||
}),
|
||||
).resolves.toBe(true);
|
||||
await expect(
|
||||
recordSubagentTaskControlReceipt({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: 'task-1',
|
||||
receipt: {
|
||||
...accepted,
|
||||
status: 'applied',
|
||||
boundary: 'tool',
|
||||
updatedAt: new Date('2026-08-24T12:00:01.000Z'),
|
||||
},
|
||||
}),
|
||||
).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: 'another-user',
|
||||
conversationId,
|
||||
taskId: 'task-1',
|
||||
receipt: accepted,
|
||||
}),
|
||||
).resolves.toBe(false);
|
||||
|
||||
const stored = await Message.findOne({
|
||||
user: 'user123',
|
||||
conversationId,
|
||||
messageId: 'task-1:user',
|
||||
})
|
||||
.select('+subagentTask')
|
||||
.lean<IMessage>();
|
||||
expect(stored).not.toBeNull();
|
||||
if (stored == null) throw new Error('Expected the durable task input.');
|
||||
expect(stored.subagentTask?.status).toBe('running');
|
||||
expect(stored.subagentTask?.controlReceipts).toEqual([
|
||||
expect.objectContaining({
|
||||
invocationId: 'invocation-1',
|
||||
action: 'steer',
|
||||
status: 'applied',
|
||||
boundary: 'tool',
|
||||
}),
|
||||
]);
|
||||
});
|
||||
|
||||
it('updates only the authorized tenant when message identities collide', async () => {
|
||||
const conversationId = uuidv4();
|
||||
const taskId = 'tenant-task';
|
||||
await Promise.all(
|
||||
['tenant-a', 'tenant-b'].map((tenantId) =>
|
||||
Message.create({
|
||||
user: 'user123',
|
||||
tenantId,
|
||||
conversationId,
|
||||
messageId: `${taskId}:user`,
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
sender: 'User',
|
||||
text: 'Do the tenant work',
|
||||
endpoint: 'agents',
|
||||
isCreatedByUser: true,
|
||||
subagentTask: { attemptKey: `attempt-${tenantId}`, status: 'running' },
|
||||
}),
|
||||
),
|
||||
);
|
||||
const now = new Date('2026-08-24T12:00:00.000Z');
|
||||
|
||||
await expect(
|
||||
recordSubagentTaskControlReceipt({
|
||||
userId: 'user123',
|
||||
tenantId: 'tenant-b',
|
||||
conversationId,
|
||||
taskId,
|
||||
receipt: {
|
||||
invocationId: 'tenant-invocation',
|
||||
fingerprint: 'tenant-fingerprint',
|
||||
controlId: 'tenant-control',
|
||||
action: 'queue',
|
||||
status: 'accepted',
|
||||
createdAt: now,
|
||||
updatedAt: now,
|
||||
},
|
||||
}),
|
||||
).resolves.toBe(true);
|
||||
|
||||
const [tenantA, tenantB] = await Promise.all(
|
||||
['tenant-a', 'tenant-b'].map((tenantId) =>
|
||||
Message.findOne({
|
||||
user: 'user123',
|
||||
tenantId,
|
||||
conversationId,
|
||||
messageId: `${taskId}:user`,
|
||||
})
|
||||
.select('+subagentTask')
|
||||
.lean<IMessage>(),
|
||||
),
|
||||
);
|
||||
expect(tenantA?.subagentTask?.controlReceipts).toBeUndefined();
|
||||
expect(tenantB?.subagentTask?.controlReceipts).toEqual([
|
||||
expect.objectContaining({ invocationId: 'tenant-invocation' }),
|
||||
]);
|
||||
});
|
||||
|
||||
it('retains accepted commands while bounding terminal receipt history', async () => {
|
||||
const conversationId = uuidv4();
|
||||
await createTaskInput(conversationId);
|
||||
const createdAt = new Date('2026-08-24T12:00:00.000Z');
|
||||
await recordSubagentTaskControlReceipt({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: 'task-1',
|
||||
receipt: {
|
||||
invocationId: 'pending',
|
||||
fingerprint: 'pending-fingerprint',
|
||||
controlId: 'pending-control',
|
||||
action: 'queue',
|
||||
status: 'accepted',
|
||||
createdAt,
|
||||
updatedAt: createdAt,
|
||||
},
|
||||
});
|
||||
for (let index = 0; index < 70; index += 1) {
|
||||
await recordSubagentTaskControlReceipt({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: 'task-1',
|
||||
receipt: {
|
||||
invocationId: `terminal-${index}`,
|
||||
fingerprint: `fingerprint-${index}`,
|
||||
controlId: `control-${index}`,
|
||||
action: 'steer',
|
||||
status: 'applied',
|
||||
createdAt: new Date(createdAt.getTime() + index + 1),
|
||||
updatedAt: new Date(createdAt.getTime() + index + 1),
|
||||
boundary: 'tool',
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
const stored = await Message.findOne({
|
||||
user: 'user123',
|
||||
conversationId,
|
||||
messageId: 'task-1:user',
|
||||
})
|
||||
.select('+subagentTask')
|
||||
.lean<IMessage>();
|
||||
expect(stored).not.toBeNull();
|
||||
if (stored == null) throw new Error('Expected the durable task input.');
|
||||
expect(stored.subagentTask?.controlReceipts).toHaveLength(64);
|
||||
expect(stored.subagentTask?.controlReceipts).toEqual(
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ invocationId: 'pending', status: 'accepted' }),
|
||||
expect.objectContaining({ invocationId: 'terminal-69', status: 'applied' }),
|
||||
]),
|
||||
);
|
||||
expect(stored.subagentTask?.controlReceipts).not.toEqual(
|
||||
expect.arrayContaining([expect.objectContaining({ invocationId: 'terminal-0' })]),
|
||||
);
|
||||
|
||||
/** An old idempotent retry retains its original occurrence ordering and
|
||||
* cannot evict newer terminal history merely by arriving again. */
|
||||
await recordSubagentTaskControlReceipt({
|
||||
userId: 'user123',
|
||||
conversationId,
|
||||
taskId: 'task-1',
|
||||
receipt: {
|
||||
invocationId: 'terminal-0',
|
||||
fingerprint: 'fingerprint-0',
|
||||
controlId: 'control-0',
|
||||
action: 'steer',
|
||||
status: 'applied',
|
||||
createdAt: new Date(createdAt.getTime() + 1),
|
||||
updatedAt: new Date(createdAt.getTime() + 1),
|
||||
boundary: 'tool',
|
||||
},
|
||||
});
|
||||
const afterReplay = await Message.findOne({
|
||||
user: 'user123',
|
||||
conversationId,
|
||||
messageId: 'task-1:user',
|
||||
})
|
||||
.select('+subagentTask')
|
||||
.lean<IMessage>();
|
||||
expect(afterReplay?.subagentTask?.controlReceipts).toHaveLength(64);
|
||||
expect(afterReplay?.subagentTask?.controlReceipts).toEqual(
|
||||
expect.arrayContaining([expect.objectContaining({ invocationId: 'terminal-7' })]),
|
||||
);
|
||||
expect(afterReplay?.subagentTask?.controlReceipts).not.toEqual(
|
||||
expect.arrayContaining([expect.objectContaining({ invocationId: 'terminal-0' })]),
|
||||
);
|
||||
});
|
||||
|
||||
it('defensively caps accepted receipts outside the supported task-store path', async () => {
|
||||
const conversationId = uuidv4();
|
||||
await createTaskInput(conversationId);
|
||||
const createdAt = new Date('2026-08-24T12:00:00.000Z');
|
||||
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),
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
const stored = await Message.findOne({
|
||||
user: 'user123',
|
||||
conversationId,
|
||||
messageId: 'task-1:user',
|
||||
})
|
||||
.select('+subagentTask')
|
||||
.lean<IMessage>();
|
||||
expect(stored?.subagentTask?.controlReceipts).toHaveLength(64);
|
||||
expect(stored?.subagentTask?.controlReceipts?.[0]?.invocationId).toBe('accepted-6');
|
||||
});
|
||||
});
|
||||
|
||||
describe('claimSubagentTaskResult', () => {
|
||||
const terminalResult = async (taskId: string, conversationId: string, status: string) =>
|
||||
saveMessage({ userId: 'user123' }, {
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ const MAX_STORED_USER_SUBMITTED_PATHS = 256;
|
|||
const MAX_NORMALIZED_USER_SUBMITTED_PATHS = MAX_STORED_USER_SUBMITTED_PATHS + 1;
|
||||
const MAX_STORED_USER_SUBMITTED_FIELD_PATHS = MAX_NORMALIZED_USER_SUBMITTED_PATHS;
|
||||
const MAX_USER_SUBMITTED_PATH_LENGTH = 2048;
|
||||
const MAX_SUBAGENT_CONTROL_RECEIPTS = 64;
|
||||
const MAX_SUBAGENT_CONTROL_MESSAGE_LENGTH = 4 * 1024;
|
||||
const PROVENANCE_PATHS_UNION_FIELD = '__lcProvenancePathsUnion';
|
||||
const PROVENANCE_FIELD_PATHS_UNION_FIELD = '__lcProvenanceFieldPathsUnion';
|
||||
const HITL_MESSAGE_FILTER_FIELD_SET = new Set<string>(HITL_MESSAGE_FILTER_FIELDS);
|
||||
|
|
@ -283,6 +285,13 @@ export interface MessageMethods {
|
|||
params: Partial<IMessage> & { newMessageId?: string },
|
||||
metadata?: { context?: string },
|
||||
): Promise<IMessage | null | undefined>;
|
||||
recordSubagentTaskControlReceipt(input: {
|
||||
userId: string;
|
||||
conversationId: string;
|
||||
taskId: string;
|
||||
tenantId?: string;
|
||||
receipt: NonNullable<NonNullable<IMessage['subagentTask']>['controlReceipts']>[number];
|
||||
}): Promise<boolean>;
|
||||
bulkSaveMessages(
|
||||
messages: Array<Partial<IMessage>>,
|
||||
overrideTimestamp?: boolean,
|
||||
|
|
@ -866,6 +875,275 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Atomically records one bounded parent-to-child control receipt on the
|
||||
* durable task input. Terminal receipt states are monotonic, and accepted
|
||||
* receipts are retained ahead of older terminal history when the bound fills.
|
||||
*/
|
||||
async function recordSubagentTaskControlReceipt({
|
||||
userId,
|
||||
conversationId,
|
||||
taskId,
|
||||
tenantId,
|
||||
receipt,
|
||||
}: {
|
||||
userId: string;
|
||||
conversationId: string;
|
||||
taskId: string;
|
||||
tenantId?: string;
|
||||
receipt: NonNullable<NonNullable<IMessage['subagentTask']>['controlReceipts']>[number];
|
||||
}): Promise<boolean> {
|
||||
const validActions = new Set(['steer', 'queue', 'interrupt', 'cancel', 'cancel_message']);
|
||||
const validStatuses = new Set(['accepted', 'applied', 'rejected', 'failed']);
|
||||
if (
|
||||
userId.length === 0 ||
|
||||
conversationId.length === 0 ||
|
||||
conversationId.length > 256 ||
|
||||
taskId.length === 0 ||
|
||||
taskId.length > 256 ||
|
||||
receipt.invocationId.length === 0 ||
|
||||
receipt.invocationId.length > 128 ||
|
||||
receipt.fingerprint.length === 0 ||
|
||||
receipt.fingerprint.length > 128 ||
|
||||
!validActions.has(receipt.action) ||
|
||||
!validStatuses.has(receipt.status) ||
|
||||
(receipt.controlId != null && receipt.controlId.length > 256) ||
|
||||
(receipt.message != null && receipt.message.length > MAX_SUBAGENT_CONTROL_MESSAGE_LENGTH) ||
|
||||
!Number.isFinite(receipt.createdAt.getTime()) ||
|
||||
!Number.isFinite(receipt.updatedAt.getTime())
|
||||
) {
|
||||
throw new TypeError('Invalid subagent task control receipt');
|
||||
}
|
||||
const Message = mongoose.models.Message as Model<IMessage>;
|
||||
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',
|
||||
},
|
||||
[
|
||||
{
|
||||
$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
|
||||
* explicit poller or one idempotent automatic wakeup delivery. */
|
||||
async function claimSubagentTaskResult({
|
||||
|
|
@ -1470,6 +1748,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa
|
|||
updateMessageText,
|
||||
updateToolCallResult,
|
||||
updateMessage,
|
||||
recordSubagentTaskControlReceipt,
|
||||
claimSubagentTaskResult,
|
||||
releaseSubagentTaskResultClaim,
|
||||
deleteMessagesSince,
|
||||
|
|
|
|||
|
|
@ -175,6 +175,33 @@ const messageSchema: Schema<IMessage> = new Schema(
|
|||
_id: false,
|
||||
default: undefined,
|
||||
},
|
||||
controlReceipts: {
|
||||
type: [
|
||||
{
|
||||
invocationId: { type: String, required: true },
|
||||
fingerprint: { type: String, required: true },
|
||||
controlId: { type: String },
|
||||
action: {
|
||||
type: String,
|
||||
enum: ['steer', 'queue', 'interrupt', 'cancel', 'cancel_message'],
|
||||
required: true,
|
||||
},
|
||||
status: {
|
||||
type: String,
|
||||
enum: ['accepted', 'applied', 'rejected', 'failed'],
|
||||
required: true,
|
||||
},
|
||||
createdAt: { type: Date, required: true },
|
||||
updatedAt: { type: Date, required: true },
|
||||
boundary: { type: String, enum: ['preempt', 'tool', 'turn'] },
|
||||
reason: { type: String },
|
||||
message: { type: String },
|
||||
messageTruncated: { type: Boolean },
|
||||
_id: false,
|
||||
},
|
||||
],
|
||||
default: undefined,
|
||||
},
|
||||
},
|
||||
_id: false,
|
||||
select: false,
|
||||
|
|
|
|||
|
|
@ -5,6 +5,30 @@ import type {
|
|||
} from 'librechat-data-provider';
|
||||
import type { Document } from 'mongoose';
|
||||
|
||||
export type SubagentTaskControlAction =
|
||||
| 'steer'
|
||||
| 'queue'
|
||||
| 'interrupt'
|
||||
| 'cancel'
|
||||
| 'cancel_message';
|
||||
|
||||
export type SubagentTaskControlReceiptStatus = 'accepted' | 'applied' | 'rejected' | 'failed';
|
||||
|
||||
/** Server-private durable receipt for one parent-to-child control invocation. */
|
||||
export interface ISubagentTaskControlReceipt {
|
||||
invocationId: string;
|
||||
fingerprint: string;
|
||||
controlId?: string;
|
||||
action: SubagentTaskControlAction;
|
||||
status: SubagentTaskControlReceiptStatus;
|
||||
createdAt: Date;
|
||||
updatedAt: Date;
|
||||
boundary?: 'preempt' | 'tool' | 'turn';
|
||||
reason?: string;
|
||||
message?: string;
|
||||
messageTruncated?: boolean;
|
||||
}
|
||||
|
||||
// @ts-ignore
|
||||
export interface IMessage extends Document {
|
||||
messageId: string;
|
||||
|
|
@ -70,6 +94,7 @@ export interface IMessage extends Document {
|
|||
claimId: string;
|
||||
claimedAt: Date;
|
||||
};
|
||||
controlReceipts?: ISubagentTaskControlReceipt[];
|
||||
};
|
||||
contextMeta?: {
|
||||
calibrationRatio?: number;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue