@librechat/agents 3.4.2 → 3.4.4

package/src/tools/__tests__/SubagentExecutor.test.ts

@@ -10,6 +10,7 @@ import type {
   StandardGraphInput,
   SubagentUpdateEvent,
   SubagentUsageEvent,
+  ToolSessionMap,
   ToolExecuteBatchRequest,
   ToolExecuteResult,
 } from '@/types';
@@ -1536,6 +1537,253 @@ describe('SubagentExecutor', () => {
     expect(observedChildInputs!.maxSubagentDepth).toBeUndefined();
   });
 
+  it('seeds only the selected single-agent child before its first invocation', async () => {
+    const selectedSession = {
+      session_id: 'selected-storage',
+      files: [
+        {
+          id: 'selected-file',
+          name: 'selected.txt',
+          storage_session_id: 'selected-storage',
+        },
+      ],
+      lastUpdated: 1,
+    };
+    const siblingSession = {
+      session_id: 'sibling-storage',
+      files: [
+        {
+          id: 'sibling-file',
+          name: 'sibling.txt',
+          storage_session_id: 'sibling-storage',
+        },
+      ],
+      lastUpdated: 1,
+    };
+    const selected = makeConfig('selected', {
+      agentInputs: {
+        ...makeChildInputs('selected-agent'),
+        initialSessions: new Map([
+          [Constants.EXECUTE_CODE, selectedSession],
+        ]),
+      },
+    });
+    const sibling = makeConfig('sibling', {
+      agentInputs: {
+        ...makeChildInputs('sibling-agent'),
+        initialSessions: new Map([
+          [Constants.EXECUTE_CODE, siblingSession],
+        ]),
+      },
+    });
+    const childSessions: ToolSessionMap = new Map();
+    const invoke = jest.fn(async () => {
+      expect(childSessions.get(Constants.EXECUTE_CODE)).toEqual(
+        selectedSession
+      );
+      expect(childSessions.get(Constants.EXECUTE_CODE)).not.toBe(
+        selectedSession
+      );
+      return { messages: [new AIMessage('done')] };
+    });
+    const executor = createExecutor({
+      configs: new Map([
+        [selected.type, selected],
+        [sibling.type, sibling],
+      ]),
+      createChildGraph: (): StandardGraph =>
+        ({
+          sessions: childSessions,
+          createWorkflow: () => ({ invoke }),
+          clearHeavyState: jest.fn(),
+        }) as unknown as StandardGraph,
+    });
+
+    await executor.execute({
+      description: 'Use the selected file.',
+      subagentType: selected.type,
+    });
+
+    expect(invoke).toHaveBeenCalledTimes(1);
+    expect(
+      childSessions
+        .get(Constants.EXECUTE_CODE)
+        ?.files?.some((file) => file.id === 'sibling-file')
+    ).toBe(false);
+  });
+
+  it('combines member session seeds inside a selected graph child', async () => {
+    const makeSession = (storageSessionId: string, includeFileSession = true) => ({
+      session_id: storageSessionId,
+      files: [
+        {
+          id: 'shared-file-id',
+          name: `${storageSessionId}.txt`,
+          ...(includeFileSession
+            ? { storage_session_id: storageSessionId }
+            : {}),
+        },
+      ],
+      lastUpdated: 1,
+    });
+    const first = {
+      ...makeChildInputs('first'),
+      initialSessions: new Map([
+        [Constants.EXECUTE_CODE, makeSession('first-storage')],
+      ]),
+    };
+    const second = {
+      ...makeChildInputs('second'),
+      initialSessions: new Map([
+        [Constants.EXECUTE_CODE, makeSession('second-storage', false)],
+      ]),
+    };
+    const graphConfig: GraphSubagentConfig = {
+      kind: 'graph',
+      type: 'graph-team',
+      name: 'Graph Team',
+      description: 'Runs a graph.',
+      agents: [first, second],
+      edges: [{ from: 'first', to: 'second', edgeType: 'direct' }],
+      entryAgentId: 'first',
+      resultAgentId: 'second',
+    };
+    const childSessions: ToolSessionMap = new Map();
+    const invoke = jest.fn(async () => {
+      expect(
+        childSessions
+          .get(Constants.EXECUTE_CODE)
+          ?.files?.map((file) => [file.id, file.storage_session_id])
+      ).toEqual([
+        ['shared-file-id', 'first-storage'],
+        ['shared-file-id', 'second-storage'],
+      ]);
+      return {
+        messages: [new AIMessage('done')],
+        subagentResult: { agentId: 'second' },
+      };
+    });
+    const childGraph = {
+      sessions: childSessions,
+      createWorkflow: () => ({ invoke }),
+      clearHeavyState: jest.fn(),
+    } as unknown as StandardGraph;
+    const executor = createExecutor({
+      configs: new Map([[graphConfig.type, graphConfig]]),
+      createChildGraphByKind: () => childGraph,
+    });
+
+    await executor.execute({
+      description: 'Use both files.',
+      subagentType: graphConfig.type,
+    });
+
+    expect(invoke).toHaveBeenCalledTimes(1);
+  });
+
+  it('preserves the child session across an in-memory HITL resume', async () => {
+    const seededSession = {
+      session_id: 'seed-storage',
+      files: [
+        {
+          id: 'seed-file',
+          name: 'seed.txt',
+          storage_session_id: 'seed-storage',
+        },
+      ],
+      lastUpdated: 1,
+    };
+    const hitlConfig = makeConfig('hitl-child', {
+      agentInputs: {
+        ...makeChildInputs('hitl-child'),
+        initialSessions: new Map([
+          [Constants.EXECUTE_CODE, seededSession],
+        ]),
+      },
+    });
+    const childSessions: ToolSessionMap = new Map();
+    const childInterrupt = new GraphInterrupt([
+      {
+        id: 'approval-1',
+        value: {
+          type: 'tool_approval',
+          action_requests: [],
+          review_configs: [],
+        },
+      },
+    ]);
+    const invoke = jest
+      .fn()
+      .mockImplementationOnce(async () => {
+        const session = childSessions.get(Constants.EXECUTE_CODE)!;
+        childSessions.set(Constants.EXECUTE_CODE, {
+          ...session,
+          session_id: 'live-execution',
+          files: [
+            ...(session.files ?? []),
+            {
+              id: 'generated-file',
+              name: 'generated.txt',
+              storage_session_id: 'live-execution',
+            },
+          ],
+        });
+        throw childInterrupt;
+      })
+      .mockImplementationOnce(async () => {
+        expect(childSessions.get(Constants.EXECUTE_CODE)).toMatchObject({
+          session_id: 'live-execution',
+          files: [{ id: 'seed-file' }, { id: 'generated-file' }],
+        });
+        return { messages: [new AIMessage('resumed')] };
+      });
+    const executor = createExecutor({
+      configs: new Map([[hitlConfig.type, hitlConfig]]),
+      humanInTheLoop: { enabled: true },
+      checkpointer: new MemorySaver(),
+      createChildGraph: (): StandardGraph =>
+        ({
+          sessions: childSessions,
+          createWorkflow: () => ({
+            getState: jest.fn().mockResolvedValue({
+              values: {},
+              next: [],
+              tasks: [],
+            }),
+            invoke,
+          }),
+          getChildCheckpointThreadIds: jest.fn(() => []),
+          createSubagentResumeState: jest.fn(() => ({
+            toolCallSteps: [],
+            toolSessions: [],
+            toolNodes: [],
+            eagerToolUsage: [],
+            eagerToolSuppressions: [],
+          })),
+          clearHeavyState: jest.fn(),
+        }) as unknown as StandardGraph,
+    });
+    const params = {
+      description: 'Pause and resume.',
+      subagentType: hitlConfig.type,
+      threadId: 'durable-parent',
+      parentToolCallId: 'call-hitl',
+    };
+
+    await expect(executor.execute(params)).rejects.toBeInstanceOf(
+      GraphInterrupt
+    );
+    const result = await executor.execute({
+      ...params,
+      parentConfigurable: {
+        __pregel_resume_map: { 'approval-1': [] },
+      },
+    });
+
+    expect(result.content).toBe('resumed');
+    expect(invoke).toHaveBeenCalledTimes(2);
+  });
+
   describe('parentConfigurable inheritance', () => {
     type CapturingGraphFactory = {
       factory: () => StandardGraph;
@@ -2154,6 +2402,146 @@ describe('SubagentExecutor', () => {
       expect(phases[phases.length - 1]).toBe('stop');
     });
 
+    it('attributes graph updates to validated payload and step-local member identities', async () => {
+      const updates: SubagentUpdateEvent[] = [];
+      const registry = new HandlerRegistry();
+      registry.register(GraphEvents.ON_SUBAGENT_UPDATE, {
+        handle: (_event, data): void => {
+          updates.push(data as SubagentUpdateEvent);
+        },
+      });
+      const graphConfig: GraphSubagentConfig = {
+        kind: 'graph',
+        type: 'graph-team',
+        name: 'Graph Team',
+        description: 'Runs a graph.',
+        agents: [
+          makeChildInputs('entry'),
+          makeChildInputs('worker'),
+          makeChildInputs('result'),
+        ],
+        edges: [
+          { from: 'entry', to: 'worker', edgeType: 'direct' },
+          { from: 'worker', to: 'result', edgeType: 'direct' },
+        ],
+        entryAgentId: 'entry',
+        resultAgentId: 'result',
+      };
+      const graphFactory = (): StandardGraph =>
+        ({
+          createWorkflow: (): { invoke: jest.Mock } => ({
+            invoke: jest.fn().mockImplementation(async (_state, options) => {
+              const opts = options as { callbacks?: unknown[] };
+              const forwarder = (opts.callbacks ?? [])[0] as {
+                handleCustomEvent?: (
+                  eventName: string,
+                  data: unknown,
+                  runId: string,
+                  tags?: string[],
+                  metadata?: Record<string, unknown>
+                ) => Promise<void> | void;
+              };
+              const emit = async (
+                eventName: string,
+                data: unknown,
+                agentId: string
+              ): Promise<void> => {
+                await forwarder.handleCustomEvent?.(
+                  eventName,
+                  data,
+                  'child-run',
+                  [],
+                  { agentId }
+                );
+              };
+
+              await emit(
+                GraphEvents.ON_RUN_STEP,
+                {
+                  id: 'step_entry',
+                  type: StepTypes.MESSAGE_CREATION,
+                  agentId: 'entry',
+                },
+                'parent-agent'
+              );
+              await emit(
+                GraphEvents.ON_MESSAGE_DELTA,
+                { id: 'step_entry', delta: { content: 'entry' } },
+                'parent-agent'
+              );
+              await emit(
+                GraphEvents.ON_REASONING_DELTA,
+                { id: 'step_entry', delta: { content: [] } },
+                'parent-agent'
+              );
+              await emit(
+                GraphEvents.ON_RUN_STEP,
+                {
+                  id: 'step_worker',
+                  type: StepTypes.MESSAGE_CREATION,
+                  agentId: 'worker',
+                },
+                'entry'
+              );
+              await emit(
+                GraphEvents.ON_RUN_STEP,
+                {
+                  id: 'step_spoofed',
+                  type: StepTypes.MESSAGE_CREATION,
+                  agentId: 'unknown-member',
+                },
+                'worker'
+              );
+              await emit(
+                GraphEvents.ON_RUN_STEP_COMPLETED,
+                { result: { id: 'step_result', agentId: 'result' } },
+                'parent-agent'
+              );
+              await emit(
+                GraphEvents.ON_MESSAGE_DELTA,
+                { id: 'step_metadata', delta: { content: 'worker' } },
+                'worker'
+              );
+              return {
+                messages: [new AIMessage('ok')],
+                subagentResult: { agentId: 'result' },
+              };
+            }),
+          }),
+          clearHeavyState: jest.fn(),
+        }) as unknown as StandardGraph;
+      const executor = createExecutor({
+        configs: new Map([[graphConfig.type, graphConfig]]),
+        createChildGraphByKind: graphFactory,
+        parentHandlerRegistry: registry,
+      });
+
+      await executor.execute({
+        description: 'Run the graph.',
+        subagentType: graphConfig.type,
+      });
+
+      const activity = updates.filter(
+        (update) => update.phase !== 'start' && update.phase !== 'stop'
+      );
+      expect(
+        activity.map(({ phase, memberAgentId }) => ({
+          phase,
+          memberAgentId,
+        }))
+      ).toEqual([
+        { phase: 'run_step', memberAgentId: 'entry' },
+        { phase: 'message_delta', memberAgentId: 'entry' },
+        { phase: 'reasoning_delta', memberAgentId: 'entry' },
+        { phase: 'run_step', memberAgentId: 'worker' },
+        { phase: 'run_step', memberAgentId: undefined },
+        { phase: 'run_step_completed', memberAgentId: 'result' },
+        { phase: 'message_delta', memberAgentId: 'worker' },
+      ]);
+      expect(updates[0].memberAgentId).toBeUndefined();
+      expect(updates[updates.length - 1].memberAgentId).toBeUndefined();
+    });
+
     it('keeps toolDefinitions on child when registry has ON_TOOL_EXECUTE handler', async () => {
       const registry = new HandlerRegistry();
       registry.register(GraphEvents.ON_TOOL_EXECUTE, {
@@ -3138,6 +3526,54 @@ describe('summarizeEvent', () => {
 });
 
 describe('sanitizeForwardedSubagentUpdateData', () => {
+  it('preserves bounded assistant phase metadata on message-creation steps', () => {
+    const sanitized = sanitizeForwardedSubagentUpdateData(
+      GraphEvents.ON_RUN_STEP,
+      {
+        id: 'step_1',
+        stepDetails: {
+          type: StepTypes.MESSAGE_CREATION,
+          message_creation: {
+            message_id: 'message_1',
+            content_type: 'text',
+            phase: 'commentary',
+            futureSecret: 'nested-secret',
+          },
+        },
+      }
+    );
+
+    expect(sanitized).toEqual({
+      id: 'step_1',
+      stepDetails: {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: {
+          message_id: 'message_1',
+          content_type: 'text',
+          phase: 'commentary',
+        },
+      },
+    });
+
+    expect(
+      sanitizeForwardedSubagentUpdateData(GraphEvents.ON_RUN_STEP, {
+        stepDetails: {
+          type: StepTypes.MESSAGE_CREATION,
+          message_creation: {
+            message_id: 'message_2',
+            content_type: 'image',
+            phase: 'internal',
+          },
+        },
+      })
+    ).toEqual({
+      stepDetails: {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: { message_id: 'message_2' },
+      },
+    });
+  });
+
   it('uses an allowlist for run step payloads', () => {
     const sanitized = sanitizeForwardedSubagentUpdateData(
       GraphEvents.ON_RUN_STEP,

package/src/tools/subagent/SubagentExecutor.ts

@@ -40,6 +40,7 @@ import type {
   HumanInTheLoopConfig,
   InjectedMessage,
   MessageDeltaEvent,
+  MessageCreationDetails,
   ProcessedToolCall,
   ReasoningDeltaEvent,
   RunStep,
@@ -59,6 +60,7 @@ import type {
   ToolApprovalInterruptPayload,
   ToolExecuteBatchRequest,
   ToolCallDelta,
+  ToolSessionContext,
   TokenCounter,
   ToolApprovalDecision,
   ToolApprovalDecisionMap,
@@ -144,6 +146,75 @@ const SUBAGENT_CONFIG_CHANGED_MESSAGE =
 const SUBAGENT_INVOCATION_CHANGED_MESSAGE =
   'Subagent error: Subagent invocation changed for this execution.';
 
+function cloneToolSessionContext(
+  context: ToolSessionContext
+): ToolSessionContext {
+  return {
+    ...context,
+    ...(context.files == null
+      ? {}
+      : {
+        files: context.files.map((file) => ({
+          ...file,
+          storage_session_id:
+              file.storage_session_id ?? context.session_id,
+        })),
+      }),
+  };
+}
+
+function seedChildGraphSessions(
+  childGraph: StandardGraph,
+  agents: AgentInputs[]
+): void {
+  const seenFilesByTool = new Map<string, Set<string>>();
+  for (const agent of agents) {
+    if (agent.initialSessions == null) {
+      continue;
+    }
+    for (const [toolName, context] of agent.initialSessions) {
+      const existing = childGraph.sessions.get(toolName);
+      if (existing == null) {
+        const cloned = cloneToolSessionContext(context);
+        childGraph.sessions.set(toolName, cloned);
+        seenFilesByTool.set(
+          toolName,
+          new Set(
+            cloned.files?.map(
+              (file) => `${file.storage_session_id ?? ''}\0${file.id}`
+            ) ?? []
+          )
+        );
+        continue;
+      }
+      if (context.files == null || context.files.length === 0) {
+        continue;
+      }
+      const seenFiles =
+        seenFilesByTool.get(toolName) ??
+        new Set(
+          existing.files?.map(
+            (file) =>
+              `${file.storage_session_id ?? existing.session_id}\0${file.id}`
+          ) ?? []
+        );
+      const files = existing.files == null ? [] : [...existing.files];
+      for (const file of context.files) {
+        const storageSessionId =
+          file.storage_session_id ?? context.session_id;
+        const key = `${storageSessionId}\0${file.id}`;
+        if (seenFiles.has(key)) {
+          continue;
+        }
+        seenFiles.add(key);
+        files.push({ ...file, storage_session_id: storageSessionId });
+      }
+      seenFilesByTool.set(toolName, seenFiles);
+      childGraph.sessions.set(toolName, { ...existing, files });
+    }
+  }
+}
+
 async function dispatchObservationalSubagentUpdate(
   handler: EventHandler,
   event: SubagentUpdateEvent
@@ -208,9 +279,7 @@ type SanitizedRunStep = Partial<
 type SanitizedStepDetails =
   | {
       type: StepTypes.MESSAGE_CREATION;
-      message_creation?: {
-        message_id?: string;
-      };
+      message_creation?: Partial<MessageCreationDetails['message_creation']>;
     }
   | {
       type: StepTypes.TOOL_CALLS;
@@ -2132,6 +2201,9 @@ export class SubagentExecutor {
       });
     }
     childGraph ??= this.createChildGraph(childGraphInput);
+    if (cachedChildRun == null) {
+      seedChildGraphSessions(childGraph, childPlan.agents);
+    }
     let forwarding: ForwarderCallback | undefined;
     if (forwardingEnabled) {
       forwarding = this.createForwarderCallback({
@@ -2139,6 +2211,8 @@ export class SubagentExecutor {
         subagentType,
         subagentKind: childPlan.kind,
         subagentAgentId: childAgentId,
+        graphMemberInputs:
+          childPlan.kind === 'graph' ? childPlan.memberInputs : undefined,
         childRunId,
         parentToolCallId,
         executionContext: childExecutionContext,
@@ -2589,6 +2663,7 @@ export class SubagentExecutor {
     subagentType: string;
     subagentKind: 'agent' | 'graph';
     subagentAgentId: string;
+    graphMemberInputs?: ReadonlyMap<string, AgentInputs>;
     childRunId: string;
     parentToolCallId?: string;
     executionContext: SubagentExecutionContext;
@@ -2598,12 +2673,14 @@ export class SubagentExecutor {
       subagentType,
       subagentKind,
       subagentAgentId,
+      graphMemberInputs,
       childRunId,
       parentToolCallId,
       executionContext,
     } = args;
     const immediateParentRunId = this.parentRunId;
     const parentAgentId = this.parentAgentId;
+    const memberAgentIdByStepId = new Map<string, string>();
 
     const wrap = async (
       eventName: string,
@@ -2710,7 +2787,14 @@ export class SubagentExecutor {
         metadata?: Record<string, unknown>
       ): Promise<void> => {
         const memberAgentId =
-          asNonEmptyString(metadata?.agentId) ?? getEventAgentId(data);
+          graphMemberInputs == null
+            ? (asNonEmptyString(metadata?.agentId) ?? getEventAgentId(data))
+            : getGraphEventMemberAgentId({
+              data,
+              metadata,
+              memberInputs: graphMemberInputs,
+              memberAgentIdByStepId,
+            });
         if (eventName === GraphEvents.ON_TOOL_EXECUTE) {
           const toolHandler = parentRegistry.getHandler(
             GraphEvents.ON_TOOL_EXECUTE
@@ -2928,6 +3012,52 @@ function getEventAgentId(data: unknown): string | undefined {
   );
 }
 
+function getEventStepId(data: unknown): string | undefined {
+  if (data == null || typeof data !== 'object') {
+    return undefined;
+  }
+  const event = data as {
+    id?: unknown;
+    result?: { id?: unknown };
+  };
+  return asNonEmptyString(event.id) ?? asNonEmptyString(event.result?.id);
+}
+
+function getGraphEventMemberAgentId({
+  data,
+  metadata,
+  memberInputs,
+  memberAgentIdByStepId,
+}: {
+  data: unknown;
+  metadata?: Record<string, unknown>;
+  memberInputs: ReadonlyMap<string, AgentInputs>;
+  memberAgentIdByStepId: Map<string, string>;
+}): string | undefined {
+  const eventAgentId = getEventAgentId(data);
+  const stepId = getEventStepId(data);
+  if (eventAgentId != null) {
+    if (!memberInputs.has(eventAgentId)) {
+      return undefined;
+    }
+    if (stepId != null) {
+      memberAgentIdByStepId.set(stepId, eventAgentId);
+    }
+    return eventAgentId;
+  }
+
+  const stepMemberAgentId =
+    stepId == null ? undefined : memberAgentIdByStepId.get(stepId);
+  if (stepMemberAgentId != null) {
+    return stepMemberAgentId;
+  }
+
+  const metadataAgentId = asNonEmptyString(metadata?.agentId);
+  return metadataAgentId != null && memberInputs.has(metadataAgentId)
+    ? metadataAgentId
+    : undefined;
+}
+
 /**
  * Best-effort read of the configured model from a subagent's client
  * options. Providers disagree on the key (`model` vs `modelName`), and the
@@ -3182,7 +3312,11 @@ function sanitizeStepDetails(
     return undefined;
   }
   const rawDetails = stepDetails as {
-    message_creation?: { message_id?: unknown };
+    message_creation?: {
+      message_id?: unknown;
+      content_type?: unknown;
+      phase?: unknown;
+    };
     tool_calls?: unknown[];
     type?: unknown;
   };
@@ -3190,9 +3324,27 @@ function sanitizeStepDetails(
     const sanitized: SanitizedStepDetails = {
       type: StepTypes.MESSAGE_CREATION,
     };
-    const messageId = rawDetails.message_creation?.message_id;
-    if (typeof messageId === 'string') {
-      sanitized.message_creation = { message_id: messageId };
+    const rawMessageCreation = rawDetails.message_creation;
+    const messageId = rawMessageCreation?.message_id;
+    const contentType = rawMessageCreation?.content_type;
+    const phase = rawMessageCreation?.phase;
+    if (
+      typeof messageId === 'string' ||
+      contentType === ContentTypes.TEXT ||
+      contentType === ContentTypes.THINK ||
+      phase === 'commentary' ||
+      phase === 'final_answer'
+    ) {
+      sanitized.message_creation = {
+        ...(typeof messageId === 'string' ? { message_id: messageId } : {}),
+        ...(contentType === ContentTypes.TEXT ||
+        contentType === ContentTypes.THINK
+          ? { content_type: contentType }
+          : {}),
+        ...(phase === 'commentary' || phase === 'final_answer'
+          ? { phase }
+          : {}),
+      };
     }
     return sanitized;
   }