Skip to content

Commit 666d2bd

Browse files
committed
fix: preserve child continuation context
1 parent 4822845 commit 666d2bd

4 files changed

Lines changed: 307 additions & 38 deletions

File tree

‎api/server/services/Endpoints/agents/subagentThreadStore.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ const subagentThreadTaskStore = createSubagentThreadTaskStore({
77
deleteConvos: db.deleteConvos,
88
deleteMessages: db.deleteMessages,
99
getConvo: db.getConvo,
10+
getFiles: db.getFiles,
1011
getMessages: db.getMessages,
1112
saveConvo: db.saveConvo,
1213
saveMessage: db.saveMessage,

‎packages/api/src/agents/guard.spec.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ function makeStore(): SubagentThreadTaskStore {
4040
deleteConvos: unused as AllMethods['deleteConvos'],
4141
deleteMessages: unused as AllMethods['deleteMessages'],
4242
getConvo: unused as AllMethods['getConvo'],
43+
getFiles: unused as AllMethods['getFiles'],
4344
getMessages: unused as AllMethods['getMessages'],
4445
saveConvo: unused as AllMethods['saveConvo'],
4546
saveMessage: unused as AllMethods['saveMessage'],

‎packages/api/src/agents/subagentThreads.spec.ts‎

Lines changed: 157 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import mongoose from 'mongoose';
22
import { randomUUID } from 'node:crypto';
3-
import { EModelEndpoint } from 'librechat-data-provider';
43
import { MongoMemoryServer } from 'mongodb-memory-server';
4+
import { EModelEndpoint, FileSources } from 'librechat-data-provider';
55
import { AIMessage, HumanMessage } from '@librechat/agents/langchain/messages';
66
import { createMethods, createModels, logger, tenantStorage } from '@librechat/data-schemas';
77
import type {
@@ -112,6 +112,7 @@ beforeEach(async () => {
112112
await Promise.all([
113113
(mongoose.models.Message as mongoose.Model<IMessage>).deleteMany({}),
114114
(mongoose.models.Conversation as mongoose.Model<IConversation>).deleteMany({}),
115+
mongoose.models.File.deleteMany({}),
115116
]);
116117
});
117118

@@ -201,6 +202,51 @@ describe('SubagentThreadTaskStore', () => {
201202
});
202203
});
203204

205+
it('does not overwrite a child title changed while detached execution is running', async () => {
206+
const userId = 'user-title-race';
207+
const parentConversationId = randomUUID();
208+
await saveParent(userId, parentConversationId);
209+
const store = new SubagentThreadTaskStore(methods);
210+
const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId });
211+
let release = (_value: { content: string; messages: BaseMessage[] }): void => undefined;
212+
let markEntered = (): void => undefined;
213+
const entered = new Promise<void>((resolve) => {
214+
markEntered = resolve;
215+
});
216+
const result = new Promise<{ content: string; messages: BaseMessage[] }>((resolve) => {
217+
release = resolve;
218+
});
219+
const started = store.start(
220+
taskRequest(config.scopeId, {
221+
run: async () => {
222+
markEntered();
223+
return result;
224+
},
225+
}),
226+
);
227+
if (!started.accepted) return;
228+
await entered;
229+
230+
await methods.saveConvo(
231+
{ userId },
232+
{ conversationId: requireThreadId(started), title: 'Renamed while running' },
233+
{ noUpsert: true },
234+
);
235+
release({
236+
content: 'Renamed child completed.',
237+
messages: [
238+
new HumanMessage('Investigate the issue.'),
239+
new AIMessage('Renamed child completed.'),
240+
],
241+
});
242+
await waitForSettled(store, config.scopeId, started);
243+
244+
expect(await methods.getConvo(userId, requireThreadId(started))).toMatchObject({
245+
title: 'Renamed while running',
246+
subagentThread: { userRunnable: true },
247+
});
248+
});
249+
204250
it('atomically excludes detached and ordinary turns with the same child lease', async () => {
205251
const userId = 'user-shared-lease';
206252
const parentConversationId = randomUUID();
@@ -294,13 +340,7 @@ describe('SubagentThreadTaskStore', () => {
294340
const flakyMethods = {
295341
...methods,
296342
saveConvo: jest.fn(async (...args: Parameters<AllMethods['saveConvo']>) => {
297-
const lineage = args[1].subagentThread;
298-
if (
299-
lineage != null &&
300-
typeof lineage === 'object' &&
301-
'userRunnable' in lineage &&
302-
lineage.userRunnable === true
303-
) {
343+
if (args[1]['subagentThread.userRunnable'] === true) {
304344
runnableRefreshes += 1;
305345
if (runnableRefreshes === 2) {
306346
throw new Error('sidebar refresh failed');
@@ -660,7 +700,7 @@ describe('SubagentThreadTaskStore', () => {
660700
});
661701

662702
it('folds ordinary child-chat turns into the next parent-driven continuation', async () => {
663-
const userId = 'user-manual';
703+
const userId = new mongoose.Types.ObjectId().toString();
664704
const parentConversationId = randomUUID();
665705
await saveParent(userId, parentConversationId);
666706
const store = new SubagentThreadTaskStore(methods);
@@ -672,6 +712,26 @@ describe('SubagentThreadTaskStore', () => {
672712
const manualUserId = randomUUID();
673713
const manualAssistantId = randomUUID();
674714
const manualAttachmentId = randomUUID();
715+
await methods.createFile({
716+
file_id: 'manual-file',
717+
user: new mongoose.Types.ObjectId(userId),
718+
filename: 'brief.txt',
719+
filepath: '/uploads/brief.txt',
720+
type: 'text/plain',
721+
bytes: 32,
722+
source: FileSources.text,
723+
text: 'The launch code is ORBIT-7.',
724+
});
725+
await methods.createFile({
726+
file_id: 'foreign-file',
727+
user: new mongoose.Types.ObjectId(),
728+
filename: 'foreign.txt',
729+
filepath: '/uploads/foreign.txt',
730+
type: 'text/plain',
731+
bytes: 40,
732+
source: FileSources.text,
733+
text: 'This cross-owner secret must not be restored.',
734+
});
675735
await methods.saveMessage(
676736
{ userId },
677737
{
@@ -709,8 +769,11 @@ describe('SubagentThreadTaskStore', () => {
709769
parentMessageId: manualAssistantId,
710770
sender: 'User',
711771
text: '',
712-
files: [{ file_id: 'manual-file', filename: 'brief.pdf' }],
713-
attachments: [{ conversationId: first.task.threadId, filename: 'brief.pdf' }],
772+
files: [
773+
{ file_id: 'manual-file', filename: 'brief.txt' },
774+
{ file_id: 'foreign-file', filename: 'foreign.txt' },
775+
],
776+
attachments: [{ conversationId: first.task.threadId, filename: 'brief.txt' }],
714777
endpoint: EModelEndpoint.agents,
715778
isCreatedByUser: true,
716779
},
@@ -750,7 +813,10 @@ describe('SubagentThreadTaskStore', () => {
750813
{ type: 'image_url', image_url: { url: 'data:image/png;base64,AAAA' } },
751814
]);
752815
expect(restored[3].content).toBe('The child answered the human.');
753-
expect(restored[4].content).toBe(ATTACHMENT_ONLY_TEXT);
816+
expect(restored[4].content).toContain('Attached document(s):');
817+
expect(restored[4].content).toContain('The launch code is ORBIT-7.');
818+
expect(restored[4].content).not.toContain('cross-owner secret');
819+
expect(restored[4].content).not.toBe(ATTACHMENT_ONLY_TEXT);
754820

755821
const afterFirstContinuation = await methods.getMessages(
756822
{ user: userId, conversationId: first.task.threadId },
@@ -784,12 +850,89 @@ describe('SubagentThreadTaskStore', () => {
784850
expect(restoredAgain[2].content).toEqual(restored[2].content);
785851
expect(restoredAgain.map((message) => message.content).slice(3)).toEqual([
786852
'The child answered the human.',
787-
ATTACHMENT_ONLY_TEXT,
853+
restored[4].content,
788854
'Return to the parent task.',
789855
'Returned.',
790856
]);
791857
});
792858

859+
it('restores persisted Gemini thought signatures from ordinary child turns', async () => {
860+
const userId = 'user-thought-signatures';
861+
const parentConversationId = randomUUID();
862+
await saveParent(userId, parentConversationId);
863+
const store = new SubagentThreadTaskStore(methods);
864+
const config = buildSubagentThreadTaskConfig(store, { userId, parentConversationId });
865+
const first = store.start(taskRequest(config.scopeId));
866+
await waitForSettled(store, config.scopeId, first);
867+
if (!first.accepted) return;
868+
869+
const manualUserId = randomUUID();
870+
await methods.saveMessage(
871+
{ userId },
872+
{
873+
messageId: manualUserId,
874+
conversationId: first.task.threadId,
875+
parentMessageId: `${first.task.taskId}:assistant`,
876+
sender: 'User',
877+
text: 'Run the signed tool step.',
878+
endpoint: EModelEndpoint.agents,
879+
isCreatedByUser: true,
880+
},
881+
);
882+
await methods.saveMessage(
883+
{ userId },
884+
{
885+
messageId: randomUUID(),
886+
conversationId: first.task.threadId,
887+
parentMessageId: manualUserId,
888+
sender: 'researcher-agent',
889+
text: '',
890+
content: [
891+
{ type: 'text', text: '', tool_call_ids: ['signed-call'] },
892+
{
893+
type: 'tool_call',
894+
tool_call: {
895+
id: 'signed-call',
896+
name: 'lookup',
897+
args: '{}',
898+
output: 'signed result',
899+
},
900+
},
901+
],
902+
metadata: { thoughtSignatures: { 'signed-call': 'gemini-signature' } },
903+
endpoint: EModelEndpoint.agents,
904+
isCreatedByUser: false,
905+
},
906+
);
907+
908+
let restored: BaseMessage[] = [];
909+
const continued = store.start(
910+
taskRequest(config.scopeId, {
911+
threadId: first.task.threadId,
912+
input: 'Continue after the signed step.',
913+
run: async (_runtime, initialMessages = []) => {
914+
restored = initialMessages;
915+
return {
916+
content: 'Continued.',
917+
messages: [
918+
...initialMessages,
919+
new HumanMessage('Continue after the signed step.'),
920+
new AIMessage('Continued.'),
921+
],
922+
};
923+
},
924+
}),
925+
);
926+
await waitForSettled(store, config.scopeId, continued);
927+
928+
const signedMessage = restored.find(
929+
(message) =>
930+
message instanceof AIMessage &&
931+
message.tool_calls?.some((call) => call.id === 'signed-call'),
932+
);
933+
expect(signedMessage?.additional_kwargs.signatures).toEqual(['gemini-signature']);
934+
});
935+
793936
it('persists a compacted canonical replacement without replaying superseded history', async () => {
794937
const userId = 'user-compaction';
795938
const parentConversationId = randomUUID();
@@ -1129,13 +1272,7 @@ describe('SubagentThreadTaskStore', () => {
11291272
const retryingMethods = {
11301273
...methods,
11311274
saveConvo: jest.fn(async (...args: Parameters<AllMethods['saveConvo']>) => {
1132-
const lineage = args[1].subagentThread;
1133-
if (
1134-
lineage != null &&
1135-
typeof lineage === 'object' &&
1136-
'userRunnable' in lineage &&
1137-
lineage.userRunnable === true
1138-
) {
1275+
if (args[1]['subagentThread.userRunnable'] === true) {
11391276
runnableAttempts += 1;
11401277
if (runnableAttempts === 1) {
11411278
throw new Error('transient writability update failure');

0 commit comments

Comments
 (0)