@librechat/agents 3.4.5 → 3.4.6

package/src/graphs/__tests__/Graph.closeRunStep.test.ts

@@ -0,0 +1,385 @@
+// src/graphs/__tests__/Graph.closeRunStep.test.ts
+import type * as t from '@/types';
+import { GraphEvents, StepTypes, Providers } from '@/common';
+import { HandlerRegistry } from '@/events';
+import { StandardGraph } from '../Graph';
+
+const makeAgent = (agentId: string): t.AgentInputs => ({
+  agentId,
+  provider: Providers.OPENAI,
+  instructions: `You are ${agentId}.`,
+});
+
+function createGraph(): {
+  graph: StandardGraph;
+  closed: t.RunStepClosedEvent[];
+  } {
+  const graph = new StandardGraph({
+    runId: 'run_1',
+    agents: [makeAgent('agent')],
+  });
+  const closed: t.RunStepClosedEvent[] = [];
+  const registry = new HandlerRegistry();
+  registry.register(GraphEvents.ON_RUN_STEP_CLOSED, {
+    handle: (_event, data): void => {
+      closed.push(data as t.RunStepClosedEvent);
+    },
+  });
+  graph.handlerRegistry = registry;
+  return { graph, closed };
+}
+
+function seedStep(
+  graph: StandardGraph,
+  id: string,
+  type: StepTypes = StepTypes.MESSAGE_CREATION
+): t.RunStep {
+  const index = graph.contentData.length;
+  const stepDetails: t.StepDetails =
+    type === StepTypes.TOOL_CALLS
+      ? { type: StepTypes.TOOL_CALLS, tool_calls: [] }
+      : {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: { message_id: `msg_${id}` },
+      };
+  const step: t.RunStep = {
+    id,
+    type,
+    index,
+    stepIndex: index,
+    stepDetails,
+    usage: null,
+    created_at: 1_000,
+    status: 'in_progress',
+  };
+  graph.contentData.push(step);
+  graph.contentIndexMap.set(id, index);
+  return step;
+}
+
+describe('StandardGraph.closeRunStep', () => {
+  it('stamps the terminal status + timestamp and emits ON_RUN_STEP_CLOSED once', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a');
+
+    const first = await graph.closeRunStep('step_a', 'completed', {
+      at: 2_000,
+    });
+    expect(first).toBe(true);
+    expect(step.status).toBe('completed');
+    expect(step.completed_at).toBe(2_000);
+    expect(closed).toHaveLength(1);
+    expect(closed[0]).toMatchObject({
+      id: 'step_a',
+      index: 0,
+      type: StepTypes.MESSAGE_CREATION,
+      status: 'completed',
+      created_at: 1_000,
+      closed_at: 2_000,
+    });
+
+    const second = await graph.closeRunStep('step_a', 'completed', {
+      at: 3_000,
+    });
+    expect(second).toBe(false);
+    expect(step.completed_at).toBe(2_000);
+    expect(closed).toHaveLength(1);
+  });
+
+  it('keeps a terminal status immutable across later close attempts', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+
+    await graph.closeRunStep('step_a', 'cancelled', { at: 2_000 });
+    const reclosed = await graph.closeRunStep('step_a', 'completed', {
+      at: 3_000,
+    });
+    expect(reclosed).toBe(false);
+    expect(step.status).toBe('cancelled');
+    expect(step.cancelled_at).toBe(2_000);
+    expect(step.completed_at).toBeUndefined();
+    expect(closed).toHaveLength(1);
+  });
+
+  it('keeps the stored stamp and the emitted event in agreement', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+
+    await graph.closeRunStep('step_a', 'completed', { at: 2_000 });
+    const reclosed = await graph.closeRunStep('step_a', 'completed', {
+      at: 3_000,
+    });
+
+    /** No silent correction: RunStep.completed_at can never disagree with the
+     *  closed_at already delivered to subscribers. */
+    expect(reclosed).toBe(false);
+    expect(step.completed_at).toBe(2_000);
+    expect(closed).toHaveLength(1);
+    expect(closed[0].closed_at).toBe(2_000);
+  });
+
+  it('tolerates empty and unknown step ids', async () => {
+    const { graph, closed } = createGraph();
+    expect(await graph.closeRunStep('', 'completed')).toBe(false);
+    expect(await graph.closeRunStep('step_missing', 'completed')).toBe(false);
+    expect(closed).toHaveLength(0);
+  });
+});
+
+describe('StandardGraph.recordStepCompletion', () => {
+  it('closes a TOOL_CALLS step only after every registered call completes', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+    graph.registerPendingToolCall('call_1', 'step_a');
+    graph.registerPendingToolCall('call_2', 'step_a');
+
+    await graph.recordStepCompletion('step_a', { toolCallId: 'call_1' });
+    expect(step.status).toBe('in_progress');
+    expect(closed).toHaveLength(0);
+
+    await graph.recordStepCompletion('step_a', { toolCallId: 'call_2' });
+    expect(step.status).toBe('completed');
+    expect(typeof step.completed_at).toBe('number');
+    expect(closed).toHaveLength(1);
+    expect(closed[0].status).toBe('completed');
+  });
+
+  it('closes steps without pending tracking on the first completion', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a');
+
+    await graph.recordStepCompletion('step_a');
+    expect(step.status).toBe('completed');
+    expect(closed).toHaveLength(1);
+  });
+
+  it('absorbs duplicate completions without re-emitting', async () => {
+    const { graph, closed } = createGraph();
+    seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+    graph.registerPendingToolCall('call_1', 'step_a');
+
+    await graph.recordStepCompletion('step_a', { toolCallId: 'call_1' });
+    await graph.recordStepCompletion('step_a', { toolCallId: 'call_1' });
+    expect(closed).toHaveLength(1);
+  });
+
+  it('leaves a closed step untouched if a late call registers and completes', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+    graph.registerPendingToolCall('call_1', 'step_a');
+
+    await graph.recordStepCompletion('step_a', {
+      toolCallId: 'call_1',
+      at: 2_000,
+    });
+    expect(closed).toHaveLength(1);
+
+    graph.registerPendingToolCall('call_2', 'step_a');
+    await graph.recordStepCompletion('step_a', {
+      toolCallId: 'call_2',
+      at: 3_000,
+    });
+    /** First close wins outright — no second event and no divergence between
+     *  the stored stamp and what subscribers were told. */
+    expect(closed).toHaveLength(1);
+    expect(step.completed_at).toBe(2_000);
+    expect(closed[0].closed_at).toBe(2_000);
+  });
+  it('closes a multi-call step at the latest producer timestamp', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+    graph.registerPendingToolCall('call_1', 'step_a');
+    graph.registerPendingToolCall('call_2', 'step_a');
+
+    /** call_2 finished later, but its host handler was faster, so the
+     *  earlier-finishing call_1 is what drains the pending set. */
+    await graph.recordStepCompletion('step_a', {
+      toolCallId: 'call_2',
+      at: 5_000,
+    });
+    await graph.recordStepCompletion('step_a', {
+      toolCallId: 'call_1',
+      at: 4_000,
+    });
+
+    expect(step.completed_at).toBe(5_000);
+    expect(closed).toHaveLength(1);
+    expect(closed[0].closed_at).toBe(5_000);
+  });
+
+  it('uses the producer\'s completion timestamp when one is supplied', async () => {
+    const { graph, closed } = createGraph();
+    const step = seedStep(graph, 'step_a', StepTypes.TOOL_CALLS);
+    graph.registerPendingToolCall('call_1', 'step_a');
+
+    await graph.recordStepCompletion('step_a', {
+      toolCallId: 'call_1',
+      at: 4_000,
+    });
+
+    /** A slow host handler must not inflate the recorded duration. */
+    expect(step.completed_at).toBe(4_000);
+    expect(closed).toHaveLength(1);
+    expect(closed[0].closed_at).toBe(4_000);
+  });
+});
+
+describe('StandardGraph.closeUnfinishedRunSteps', () => {
+  it('closes only non-terminal steps with the sweep status', async () => {
+    const { graph, closed } = createGraph();
+    const done = seedStep(graph, 'step_done', StepTypes.TOOL_CALLS);
+    const openMessage = seedStep(graph, 'step_msg');
+    const openTool = seedStep(graph, 'step_tool', StepTypes.TOOL_CALLS);
+    await graph.closeRunStep('step_done', 'completed', { at: 2_000 });
+    closed.length = 0;
+
+    await graph.closeUnfinishedRunSteps('cancelled', 5_000);
+
+    expect(done.completed_at).toBe(2_000);
+    expect(openMessage.status).toBe('cancelled');
+    expect(openMessage.cancelled_at).toBe(5_000);
+    expect(openTool.status).toBe('cancelled');
+    expect(openTool.cancelled_at).toBe(5_000);
+    expect(closed.map((event) => event.id)).toEqual(['step_msg', 'step_tool']);
+    expect(closed.every((event) => event.status === 'cancelled')).toBe(true);
+    expect(graph.pendingToolCallsByStep.size).toBe(0);
+    expect(graph.openMessageStepByAgent.size).toBe(0);
+  });
+
+  it('stamps failed_at for failed sweeps', async () => {
+    const { graph } = createGraph();
+    const step = seedStep(graph, 'step_a');
+
+    await graph.closeUnfinishedRunSteps('failed', 5_000);
+    expect(step.status).toBe('failed');
+    expect(step.failed_at).toBe(5_000);
+    expect(step.cancelled_at).toBeUndefined();
+    expect(step.completed_at).toBeUndefined();
+  });
+  it('keeps sweeping when a closure handler throws', async () => {
+    const graph = new StandardGraph({
+      runId: 'run_1',
+      agents: [makeAgent('agent')],
+    });
+    const seen: string[] = [];
+    const registry = new HandlerRegistry();
+    registry.register(GraphEvents.ON_RUN_STEP_CLOSED, {
+      handle: (_event, data): void => {
+        const event = data as t.RunStepClosedEvent;
+        seen.push(event.id);
+        if (event.id === 'step_1') {
+          throw new Error('host handler blew up');
+        }
+      },
+    });
+    graph.handlerRegistry = registry;
+    const first = seedStep(graph, 'step_1');
+    const second = seedStep(graph, 'step_2');
+    const third = seedStep(graph, 'step_3');
+
+    await graph.closeUnfinishedRunSteps('cancelled', 5_000);
+
+    expect(seen).toEqual(['step_1', 'step_2', 'step_3']);
+    for (const step of [first, second, third]) {
+      expect(step.status).toBe('cancelled');
+      expect(step.cancelled_at).toBe(5_000);
+    }
+  });
+});
+
+describe('StandardGraph.trackDispatchedRunStep', () => {
+  it('reserves distinct content indexes when parallel lanes dispatch concurrently', async () => {
+    const graph = new StandardGraph({
+      runId: 'run_1',
+      agents: [makeAgent('agent')],
+    });
+    /** A host handler that yields — the window in which two lanes could
+     *  otherwise both read the same contentData.length. */
+    const registry = new HandlerRegistry();
+    registry.register(GraphEvents.ON_RUN_STEP_CLOSED, {
+      handle: async (): Promise<void> => {
+        await new Promise((resolve) => setTimeout(resolve, 5));
+      },
+    });
+    graph.handlerRegistry = registry;
+
+    const openA = seedStep(graph, 'open_a');
+    openA.agentId = 'agent_a';
+    const openB = seedStep(graph, 'open_b');
+    openB.agentId = 'agent_b';
+    graph.openMessageStepByAgent.set('agent_a', 'open_a');
+    graph.openMessageStepByAgent.set('agent_b', 'open_b');
+
+    const makeSuccessor = (id: string, agentId: string): t.RunStep => ({
+      id,
+      agentId,
+      type: StepTypes.MESSAGE_CREATION,
+      /** Both lanes computed the same stale index before dispatching. */
+      index: graph.contentData.length,
+      stepDetails: {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: { message_id: `msg_${id}` },
+      },
+      usage: null,
+      created_at: 1_000,
+      status: 'in_progress',
+    });
+    const successorA = makeSuccessor('next_a', 'agent_a');
+    const successorB = makeSuccessor('next_b', 'agent_b');
+
+    const track = (
+      graph as unknown as {
+        trackDispatchedRunStep: (step: t.RunStep) => Promise<void>;
+      }
+    ).trackDispatchedRunStep.bind(graph);
+    await Promise.all([track(successorA), track(successorB)]);
+
+    expect(successorA.index).not.toBe(successorB.index);
+    expect(graph.getRunStep('next_a')).toBe(successorA);
+    expect(graph.getRunStep('next_b')).toBe(successorB);
+    expect(graph.contentData.filter((s) => s.id === 'next_a')).toHaveLength(1);
+    const indexes = graph.contentData.map((step) => step.index);
+    expect(new Set(indexes).size).toBe(indexes.length);
+  });
+  it('stamps created_at after the predecessor closure is delivered', async () => {
+    const graph = new StandardGraph({
+      runId: 'run_1',
+      agents: [makeAgent('agent')],
+    });
+    let predecessorClosedAt = 0;
+    const registry = new HandlerRegistry();
+    registry.register(GraphEvents.ON_RUN_STEP_CLOSED, {
+      handle: async (): Promise<void> => {
+        /** A slow predecessor handler must not be charged to the successor. */
+        await new Promise((resolve) => setTimeout(resolve, 20));
+        predecessorClosedAt = Date.now();
+      },
+    });
+    graph.handlerRegistry = registry;
+
+    seedStep(graph, 'open_a');
+    graph.openMessageStepByAgent.set('', 'open_a');
+
+    const successor: t.RunStep = {
+      id: 'next_a',
+      type: StepTypes.MESSAGE_CREATION,
+      index: 0,
+      stepDetails: {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: { message_id: 'msg_next_a' },
+      },
+      usage: null,
+      status: 'in_progress',
+    };
+    await (
+      graph as unknown as {
+        trackDispatchedRunStep: (step: t.RunStep) => Promise<void>;
+      }
+    ).trackDispatchedRunStep.bind(graph)(successor);
+
+    expect(predecessorClosedAt).toBeGreaterThan(0);
+    expect(successor.created_at as number).toBeGreaterThanOrEqual(
+      predecessorClosedAt
+    );
+  });
+});

package/src/langfuseTraceShaping.ts

@@ -14,6 +14,8 @@ const CHAIN_OBSERVATION_TYPE = 'chain';
 const TOOL_OBSERVATION_TYPE = 'tool';
 const AGENT_TRACE_TAG = 'agent';
 const TITLE_TRACE_TAG = 'title';
+const ACTIVITY_PHASE_TRACE_TAG = 'activity-phase';
+const ACTIVITY_PHASE_ROOT_NAME = 'summarize-activity-phase';
 const EPHEMERAL_AGENT_SENDER_SEPARATOR = '___';
 const EPHEMERAL_AGENT_INDEX_SEPARATOR = '____';
 const OBSERVATION_METADATA_LANGGRAPH_NODE = `${LangfuseOtelSpanAttributes.OBSERVATION_METADATA}.langgraph_node`;
@@ -473,14 +475,18 @@ function shapeRootObservationType(span: MutableSpan): void {
   if (isGenerationSpan(span)) {
     return;
   }
-  if (hasTraceTag(span, AGENT_TRACE_TAG)) {
+  if (
+    hasTraceTag(span, TITLE_TRACE_TAG) ||
+    (span.name === ACTIVITY_PHASE_ROOT_NAME &&
+      hasTraceTag(span, ACTIVITY_PHASE_TRACE_TAG))
+  ) {
     span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_TYPE] =
-      ROOT_OBSERVATION_TYPE;
+      CHAIN_OBSERVATION_TYPE;
     return;
   }
-  if (hasTraceTag(span, TITLE_TRACE_TAG)) {
+  if (hasTraceTag(span, AGENT_TRACE_TAG)) {
     span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_TYPE] =
-      CHAIN_OBSERVATION_TYPE;
+      ROOT_OBSERVATION_TYPE;
   }
 }
 
@@ -497,9 +503,10 @@ function shapeRootObservationType(span: MutableSpan): void {
  * - Agent nodes become `agent` observations, while tool-dispatch nodes become
  *   stable `chain` observations whose input is scoped to the pending calls.
  *   Individual child calls remain `tool` observations (items 3 & 4).
- * - Agent trace roots become `agent` observations and title trace roots become
- *   `chain` observations. Root and trace input/output are reduced to the user
- *   question and assistant response when chat messages are available (item 2).
+ * - Agent trace roots become `agent` observations, while title and activity
+ *   summary roots become `chain` observations. Root and trace input/output are
+ *   reduced to the user question and assistant response when chat messages are
+ *   available (item 2).
  */
 export function shapeLangfuseSpan(span: ReadableSpan): void {
   const mutable = span as MutableSpan;